Merge remote-tracking branch 'origin/master'
Browse files- detect.py +3 -3
- models/yolo.py +2 -1
- train.py +1 -2
detect.py
CHANGED
@@ -102,7 +102,7 @@ def detect(save_img=False):
|
|
102 |
|
103 |
if save_img or view_img: # Add bbox to image
|
104 |
label = '%s %.2f' % (names[int(cls)], conf)
|
105 |
-
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=
|
106 |
|
107 |
# Print time (inference + NMS)
|
108 |
print('%sDone. (%.3fs)' % (s, t2 - t1))
|
@@ -139,10 +139,10 @@ def detect(save_img=False):
|
|
139 |
|
140 |
if __name__ == '__main__':
|
141 |
parser = argparse.ArgumentParser()
|
142 |
-
parser.add_argument('--weights', type=str, default='weights/
|
143 |
parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam
|
144 |
parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder
|
145 |
-
parser.add_argument('--img-size', type=int, default=
|
146 |
parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
|
147 |
parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
|
148 |
parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
|
|
|
102 |
|
103 |
if save_img or view_img: # Add bbox to image
|
104 |
label = '%s %.2f' % (names[int(cls)], conf)
|
105 |
+
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
|
106 |
|
107 |
# Print time (inference + NMS)
|
108 |
print('%sDone. (%.3fs)' % (s, t2 - t1))
|
|
|
139 |
|
140 |
if __name__ == '__main__':
|
141 |
parser = argparse.ArgumentParser()
|
142 |
+
parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='model.pt path')
|
143 |
parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam
|
144 |
parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder
|
145 |
+
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
146 |
parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
|
147 |
parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
|
148 |
parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
|
models/yolo.py
CHANGED
@@ -52,7 +52,8 @@ class Model(nn.Module):
|
|
52 |
self.md = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
53 |
|
54 |
# Define model
|
55 |
-
if nc:
|
|
|
56 |
self.md['nc'] = nc # override yaml value
|
57 |
self.model, self.save = parse_model(self.md, ch=[ch]) # model, savelist, ch_out
|
58 |
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
|
|
52 |
self.md = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
53 |
|
54 |
# Define model
|
55 |
+
if nc and nc != self.md['nc']:
|
56 |
+
print('Overriding %s nc=%g with nc=%g' % (model_cfg, self.md['nc'], nc))
|
57 |
self.md['nc'] = nc # override yaml value
|
58 |
self.model, self.save = parse_model(self.md, ch=[ch]) # model, savelist, ch_out
|
59 |
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
train.py
CHANGED
@@ -77,8 +77,7 @@ def train(hyp):
|
|
77 |
os.remove(f)
|
78 |
|
79 |
# Create model
|
80 |
-
model = Model(opt.cfg).to(device)
|
81 |
-
assert model.md['nc'] == nc, '%s nc=%g classes but %s nc=%g classes' % (opt.data, nc, opt.cfg, model.md['nc'])
|
82 |
|
83 |
# Image sizes
|
84 |
gs = int(max(model.stride)) # grid size (max stride)
|
|
|
77 |
os.remove(f)
|
78 |
|
79 |
# Create model
|
80 |
+
model = Model(opt.cfg, nc=data_dict['nc']).to(device)
|
|
|
81 |
|
82 |
# Image sizes
|
83 |
gs = int(max(model.stride)) # grid size (max stride)
|