From 786e10a1972c5354d379ba19bd68bdd0d24d0f19 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 11 Feb 2019 13:45:04 +0100 Subject: [PATCH] class labeling corrections --- detect.py | 45 +++++++++++++++++++++++++++------------------ utils/datasets.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 18 deletions(-) diff --git a/detect.py b/detect.py index b71d5d9d..ab0a071e 100755 --- a/detect.py +++ b/detect.py @@ -17,7 +17,8 @@ def detect( conf_thres=0.3, nms_thres=0.45, save_txt=False, - save_images=True + save_images=True, + webcam=False ): device = torch_utils.select_device() os.system('rm -rf ' + output) @@ -37,15 +38,20 @@ def detect( model.to(device).eval() # Set Dataloader - dataloader = LoadImages(images, img_size=img_size) + if webcam: + save_images = False + dataloader = LoadWebcam(images, img_size=img_size) + else: + dataloader = LoadImages(images, img_size=img_size) # Get classes and colors classes = load_classes(parse_data_cfg('cfg/coco.data')['names']) colors = [[random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)] for _ in range(len(classes))] for i, (path, img, im0) in enumerate(dataloader): - print("%g/%g '%s': " % (i + 1, len(dataloader), path), end='') t = time.time() + print("%g/%g '%s': " % (i + 1, len(dataloader), path), end='') + save_path = os.path.join(output, path.split('/')[-1]) # Get detections img = torch.from_numpy(img).unsqueeze(0).to(device) @@ -53,45 +59,48 @@ def detect( torch.onnx._export(model, img, 'weights/model.onnx', verbose=True) return # ONNX export pred = model(img) - pred = pred[pred[:, :, 4] > conf_thres] + pred = pred[pred[:, :, 4] > conf_thres] # remove boxes < threshold if len(pred) > 0: + # Run NMS on predictions detections = non_max_suppression(pred.unsqueeze(0), conf_thres, nms_thres)[0] - # Draw bounding boxes and labels of detections - if detections is not None: - save_path = os.path.join(output, path.split('/')[-1]) - # Rescale boxes from 416 to true image size detections[:, :4] = scale_coords(img_size, detections[:, :4], im0.shape) + # Print results to screen unique_classes = detections[:, -1].cpu().unique() for i in unique_classes: n = (detections[:, -1].cpu() == i).sum() print('%g %ss' % (n, classes[int(i)]), end=', ') + # Draw bounding boxes and labels of detections for x1, y1, x2, y2, conf, cls_conf, cls in detections: if save_txt: # Write to file with open(save_path + '.txt', 'a') as file: - file.write('%g %g %g %g %g %g\n' % (x1, y1, x2, y2, cls, cls_conf * conf)) + file.write('%g %g %g %g %g %g\n' % + (x1, y1, x2, y2, cls, cls_conf * conf)) - if save_images: # Add bbox to the image - label = '%s %.2f' % (classes[int(cls)], conf) - plot_one_box([x1, y1, x2, y2], im0, label=label, color=colors[int(cls)]) - - if save_images: # Save generated image with detections - cv2.imwrite(save_path, im0) + # Add bbox to the image + label = '%s %.2f' % (classes[int(cls)], conf) + plot_one_box([x1, y1, x2, y2], im0, label=label, color=colors[int(cls)]) print('Done. (%.3fs)' % (time.time() - t)) - if platform == 'darwin': # MacOS + if save_images: # Save generated image with detections + cv2.imwrite(save_path, im0) + + if webcam: # Show live webcam + cv2.imshow(weights, im0) + + if save_images and (platform == 'darwin'): # MacOS os.system('open ' + output + '&& open ' + save_path) if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path') - parser.add_argument('--weights', type=str, default='weights/yolov3.pt', help='path to weights file') + parser.add_argument('--cfg', type=str, default='cfg/yolov3-tiny.cfg', help='cfg file path') + parser.add_argument('--weights', type=str, default='weights/yolov3-tiny.pt', help='path to weights file') parser.add_argument('--images', type=str, default='data/samples', help='path to images') parser.add_argument('--img-size', type=int, default=32 * 13, help='size of each image dimension') parser.add_argument('--conf-thres', type=float, default=0.50, help='object confidence threshold') diff --git a/utils/datasets.py b/utils/datasets.py index ef64c4e9..66428836 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -55,6 +55,42 @@ class LoadImages: # for inference return self.nF # number of files +class LoadWebcam: # for inference + def __init__(self, path, img_size=416): + self.cam = cv2.VideoCapture(0) + self.nF = 9999 # number of image files + self.height = img_size + + def __iter__(self): + self.count = -1 + return self + + def __next__(self): + self.count += 1 + if cv2.waitKey(1) == 27: # esc to quit + cv2.destroyAllWindows() + raise StopIteration + + # Read image + ret_val, img0 = self.cam.read() + assert ret_val, 'Webcam Error' + img_path = 'webcam_%g.jpg' % self.count + img0 = cv2.flip(img0, 1) + + # Padded resize + img, _, _, _ = letterbox(img0, height=self.height) + + # Normalize RGB + img = img[:, :, ::-1].transpose(2, 0, 1) + img = np.ascontiguousarray(img, dtype=np.float32) + img /= 255.0 + + return img_path, img, img0 + + def __len__(self): + return self.nF # number of files + + class LoadImagesAndLabels: # for training def __init__(self, path, batch_size=1, img_size=608, multi_scale=False, augment=False): self.path = path