updates
This commit is contained in:
parent
7608047531
commit
bc741f30e8
13
detect.py
13
detect.py
|
@ -64,10 +64,10 @@ def detect(save_img=False):
|
||||||
if webcam:
|
if webcam:
|
||||||
view_img = True
|
view_img = True
|
||||||
torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference
|
torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference
|
||||||
dataset = LoadStreams(source, img_size=img_size, half=half)
|
dataset = LoadStreams(source, img_size=img_size)
|
||||||
else:
|
else:
|
||||||
save_img = True
|
save_img = True
|
||||||
dataset = LoadImages(source, img_size=img_size, half=half)
|
dataset = LoadImages(source, img_size=img_size)
|
||||||
|
|
||||||
# Get names and colors
|
# Get names and colors
|
||||||
names = load_classes(opt.names)
|
names = load_classes(opt.names)
|
||||||
|
@ -77,15 +77,14 @@ def detect(save_img=False):
|
||||||
t0 = time.time()
|
t0 = time.time()
|
||||||
for path, img, im0s, vid_cap in dataset:
|
for path, img, im0s, vid_cap in dataset:
|
||||||
t = time.time()
|
t = time.time()
|
||||||
|
|
||||||
# Get detections
|
|
||||||
img = torch.from_numpy(img).to(device)
|
img = torch.from_numpy(img).to(device)
|
||||||
|
img = img.half() if half else img.float() # uint8 to fp16/32
|
||||||
|
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||||||
if img.ndimension() == 3:
|
if img.ndimension() == 3:
|
||||||
img = img.unsqueeze(0)
|
img = img.unsqueeze(0)
|
||||||
pred = model(img)[0]
|
|
||||||
|
|
||||||
if opt.half:
|
# Inference
|
||||||
pred = pred.float()
|
pred = model(img)[0].float() if half else model(img)[0]
|
||||||
|
|
||||||
# Apply NMS
|
# Apply NMS
|
||||||
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
|
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
|
||||||
|
|
|
@ -42,7 +42,7 @@ def exif_size(img):
|
||||||
|
|
||||||
|
|
||||||
class LoadImages: # for inference
|
class LoadImages: # for inference
|
||||||
def __init__(self, path, img_size=416, half=False):
|
def __init__(self, path, img_size=416):
|
||||||
path = str(Path(path)) # os-agnostic
|
path = str(Path(path)) # os-agnostic
|
||||||
files = []
|
files = []
|
||||||
if os.path.isdir(path):
|
if os.path.isdir(path):
|
||||||
|
@ -59,7 +59,6 @@ class LoadImages: # for inference
|
||||||
self.nF = nI + nV # number of files
|
self.nF = nI + nV # number of files
|
||||||
self.video_flag = [False] * nI + [True] * nV
|
self.video_flag = [False] * nI + [True] * nV
|
||||||
self.mode = 'images'
|
self.mode = 'images'
|
||||||
self.half = half # half precision fp16 images
|
|
||||||
if any(videos):
|
if any(videos):
|
||||||
self.new_video(videos[0]) # new video
|
self.new_video(videos[0]) # new video
|
||||||
else:
|
else:
|
||||||
|
@ -104,8 +103,7 @@ class LoadImages: # for inference
|
||||||
|
|
||||||
# Convert
|
# Convert
|
||||||
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
||||||
img = np.ascontiguousarray(img, dtype=np.float16 if self.half else np.float32) # uint8 to fp16/fp32
|
img = np.ascontiguousarray(img)
|
||||||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
|
||||||
|
|
||||||
# cv2.imwrite(path + '.letterbox.jpg', 255 * img.transpose((1, 2, 0))[:, :, ::-1]) # save letterbox image
|
# cv2.imwrite(path + '.letterbox.jpg', 255 * img.transpose((1, 2, 0))[:, :, ::-1]) # save letterbox image
|
||||||
return path, img, img0, self.cap
|
return path, img, img0, self.cap
|
||||||
|
@ -120,9 +118,8 @@ class LoadImages: # for inference
|
||||||
|
|
||||||
|
|
||||||
class LoadWebcam: # for inference
|
class LoadWebcam: # for inference
|
||||||
def __init__(self, pipe=0, img_size=416, half=False):
|
def __init__(self, pipe=0, img_size=416):
|
||||||
self.img_size = img_size
|
self.img_size = img_size
|
||||||
self.half = half # half precision fp16 images
|
|
||||||
|
|
||||||
if pipe == '0':
|
if pipe == '0':
|
||||||
pipe = 0 # local camera
|
pipe = 0 # local camera
|
||||||
|
@ -177,8 +174,7 @@ class LoadWebcam: # for inference
|
||||||
|
|
||||||
# Convert
|
# Convert
|
||||||
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
||||||
img = np.ascontiguousarray(img, dtype=np.float16 if self.half else np.float32) # uint8 to fp16/fp32
|
img = np.ascontiguousarray(img)
|
||||||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
|
||||||
|
|
||||||
return img_path, img, img0, None
|
return img_path, img, img0, None
|
||||||
|
|
||||||
|
@ -187,10 +183,9 @@ class LoadWebcam: # for inference
|
||||||
|
|
||||||
|
|
||||||
class LoadStreams: # multiple IP or RTSP cameras
|
class LoadStreams: # multiple IP or RTSP cameras
|
||||||
def __init__(self, sources='streams.txt', img_size=416, half=False):
|
def __init__(self, sources='streams.txt', img_size=416):
|
||||||
self.mode = 'images'
|
self.mode = 'images'
|
||||||
self.img_size = img_size
|
self.img_size = img_size
|
||||||
self.half = half # half precision fp16 images
|
|
||||||
|
|
||||||
if os.path.isfile(sources):
|
if os.path.isfile(sources):
|
||||||
with open(sources, 'r') as f:
|
with open(sources, 'r') as f:
|
||||||
|
@ -251,9 +246,8 @@ class LoadStreams: # multiple IP or RTSP cameras
|
||||||
img = np.stack(img, 0)
|
img = np.stack(img, 0)
|
||||||
|
|
||||||
# Convert
|
# Convert
|
||||||
img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to 3x416x416, uint8 to float32
|
img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416
|
||||||
img = np.ascontiguousarray(img, dtype=np.float16 if self.half else np.float32)
|
img = np.ascontiguousarray(img)
|
||||||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
|
||||||
|
|
||||||
return self.sources, img, img0, None
|
return self.sources, img, img0, None
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue