From 1fdaa4987ba1e4b35bf38bfbb4dbc0b521a6585b Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 4 Jul 2020 17:44:53 -0700 Subject: [PATCH 1/2] update detect.py --- detect.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/detect.py b/detect.py index f9b12eb..2650c20 100644 --- a/detect.py +++ b/detect.py @@ -102,7 +102,7 @@ def detect(save_img=False): if save_img or view_img: # Add bbox to image label = '%s %.2f' % (names[int(cls)], conf) - plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1) + plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) # Print time (inference + NMS) print('%sDone. (%.3fs)' % (s, t2 - t1)) @@ -139,10 +139,10 @@ def detect(save_img=False): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--weights', type=str, default='weights/yolov5m.pt', help='model.pt path') + parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='model.pt path') parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder - parser.add_argument('--img-size', type=int, default=1024, help='inference size (pixels)') + parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)') From bb3c34691639b2aafb1a1f567e1dfb551a9a2563 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 4 Jul 2020 17:51:54 -0700 Subject: [PATCH 2/2] model.yaml nc inherited from dataset.yaml --- models/yolo.py | 3 ++- train.py | 3 +-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/models/yolo.py b/models/yolo.py index 7cc86b2..5d96c60 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -52,7 +52,8 @@ class Model(nn.Module): self.md = yaml.load(f, Loader=yaml.FullLoader) # model dict # Define model - if nc: + if nc and nc != self.md['nc']: + print('Overriding %s nc=%g with nc=%g' % (model_cfg, self.md['nc'], nc)) self.md['nc'] = nc # override yaml value self.model, self.save = parse_model(self.md, ch=[ch]) # model, savelist, ch_out # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) diff --git a/train.py b/train.py index 1c8ff7d..e07cbb1 100644 --- a/train.py +++ b/train.py @@ -77,8 +77,7 @@ def train(hyp): os.remove(f) # Create model - model = Model(opt.cfg).to(device) - assert model.md['nc'] == nc, '%s nc=%g classes but %s nc=%g classes' % (opt.data, nc, opt.cfg, model.md['nc']) + model = Model(opt.cfg, nc=data_dict['nc']).to(device) # Image sizes gs = int(max(model.stride)) # grid size (max stride)