Commit
·
e8cf24b
1
Parent(s):
121d90b
Initial model ensemble capability #318
Browse files- detect.py +3 -4
- models/experimental.py +17 -0
- test.py +12 -15
- utils/datasets.py +1 -1
detect.py
CHANGED
@@ -2,7 +2,7 @@ import argparse
|
|
2 |
|
3 |
import torch.backends.cudnn as cudnn
|
4 |
|
5 |
-
from
|
6 |
from utils.datasets import *
|
7 |
from utils.utils import *
|
8 |
|
@@ -20,8 +20,7 @@ def detect(save_img=False):
|
|
20 |
half = device.type != 'cpu' # half precision only supported on CUDA
|
21 |
|
22 |
# Load model
|
23 |
-
|
24 |
-
model = torch.load(weights, map_location=device)['model'].float().eval() # load FP32 model
|
25 |
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
|
26 |
if half:
|
27 |
model.half() # to FP16
|
@@ -137,7 +136,7 @@ def detect(save_img=False):
|
|
137 |
|
138 |
if __name__ == '__main__':
|
139 |
parser = argparse.ArgumentParser()
|
140 |
-
parser.add_argument('--weights', type=str, default='
|
141 |
parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam
|
142 |
parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder
|
143 |
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
|
|
2 |
|
3 |
import torch.backends.cudnn as cudnn
|
4 |
|
5 |
+
from models.experimental import *
|
6 |
from utils.datasets import *
|
7 |
from utils.utils import *
|
8 |
|
|
|
20 |
half = device.type != 'cpu' # half precision only supported on CUDA
|
21 |
|
22 |
# Load model
|
23 |
+
model = attempt_load(weights, map_location=device) # load FP32 model
|
|
|
24 |
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
|
25 |
if half:
|
26 |
model.half() # to FP16
|
|
|
136 |
|
137 |
if __name__ == '__main__':
|
138 |
parser = argparse.ArgumentParser()
|
139 |
+
parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
|
140 |
parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam
|
141 |
parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder
|
142 |
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
models/experimental.py
CHANGED
@@ -1,6 +1,7 @@
|
|
1 |
# This file contains experimental modules
|
2 |
|
3 |
from models.common import *
|
|
|
4 |
|
5 |
|
6 |
class CrossConv(nn.Module):
|
@@ -119,3 +120,19 @@ class Ensemble(nn.ModuleList):
|
|
119 |
for module in self:
|
120 |
y.append(module(x, augment)[0])
|
121 |
return torch.cat(y, 1), None # ensembled inference output, train output
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
# This file contains experimental modules
|
2 |
|
3 |
from models.common import *
|
4 |
+
from utils import google_utils
|
5 |
|
6 |
|
7 |
class CrossConv(nn.Module):
|
|
|
120 |
for module in self:
|
121 |
y.append(module(x, augment)[0])
|
122 |
return torch.cat(y, 1), None # ensembled inference output, train output
|
123 |
+
|
124 |
+
|
125 |
+
def attempt_load(weights, map_location=None):
|
126 |
+
# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
|
127 |
+
model = Ensemble()
|
128 |
+
for w in weights if isinstance(weights, list) else [weights]:
|
129 |
+
google_utils.attempt_download(w)
|
130 |
+
model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model
|
131 |
+
|
132 |
+
if len(model) == 1:
|
133 |
+
return model[-1] # return model
|
134 |
+
else:
|
135 |
+
print('Ensemble created with %s\n' % weights)
|
136 |
+
for k in ['names', 'stride']:
|
137 |
+
setattr(model, k, getattr(model[-1], k))
|
138 |
+
return model # return ensemble
|
test.py
CHANGED
@@ -1,9 +1,8 @@
|
|
1 |
import argparse
|
2 |
import json
|
3 |
|
4 |
-
from
|
5 |
from utils.datasets import *
|
6 |
-
from utils.utils import *
|
7 |
|
8 |
|
9 |
def test(data,
|
@@ -20,28 +19,26 @@ def test(data,
|
|
20 |
dataloader=None,
|
21 |
merge=False):
|
22 |
# Initialize/load model and set device
|
23 |
-
|
24 |
-
|
25 |
-
|
|
|
|
|
26 |
device = torch_utils.select_device(opt.device, batch_size=batch_size)
|
|
|
27 |
|
28 |
# Remove previous
|
29 |
for f in glob.glob('test_batch*.jpg'):
|
30 |
os.remove(f)
|
31 |
|
32 |
# Load model
|
33 |
-
|
34 |
-
model = torch.load(weights, map_location=device)['model'].float().fuse().to(device) # load to FP32
|
35 |
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
|
36 |
|
37 |
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
|
38 |
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
|
39 |
# model = nn.DataParallel(model)
|
40 |
|
41 |
-
else: # called by train.py
|
42 |
-
training = True
|
43 |
-
device = next(model.parameters()).device # get model device
|
44 |
-
|
45 |
# Half
|
46 |
half = device.type != 'cpu' and torch.cuda.device_count() == 1 # half precision only supported on single-GPU
|
47 |
if half:
|
@@ -56,11 +53,11 @@ def test(data,
|
|
56 |
niou = iouv.numel()
|
57 |
|
58 |
# Dataloader
|
59 |
-
if
|
60 |
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
|
61 |
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
|
62 |
path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
|
63 |
-
dataloader = create_dataloader(path, imgsz, batch_size,
|
64 |
hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0]
|
65 |
|
66 |
seen = 0
|
@@ -193,7 +190,7 @@ def test(data,
|
|
193 |
if save_json and map50 and len(jdict):
|
194 |
imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files]
|
195 |
f = 'detections_val2017_%s_results.json' % \
|
196 |
-
(weights.split(os.sep)[-1].replace('.pt', '') if weights else '') # filename
|
197 |
print('\nCOCO mAP with pycocotools... saving %s...' % f)
|
198 |
with open(f, 'w') as file:
|
199 |
json.dump(jdict, file)
|
@@ -226,7 +223,7 @@ def test(data,
|
|
226 |
|
227 |
if __name__ == '__main__':
|
228 |
parser = argparse.ArgumentParser(prog='test.py')
|
229 |
-
parser.add_argument('--weights', type=str, default='
|
230 |
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
|
231 |
parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
|
232 |
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
|
|
1 |
import argparse
|
2 |
import json
|
3 |
|
4 |
+
from models.experimental import *
|
5 |
from utils.datasets import *
|
|
|
6 |
|
7 |
|
8 |
def test(data,
|
|
|
19 |
dataloader=None,
|
20 |
merge=False):
|
21 |
# Initialize/load model and set device
|
22 |
+
training = model is not None
|
23 |
+
if training: # called by train.py
|
24 |
+
device = next(model.parameters()).device # get model device
|
25 |
+
|
26 |
+
else: # called directly
|
27 |
device = torch_utils.select_device(opt.device, batch_size=batch_size)
|
28 |
+
merge = opt.merge # use Merge NMS
|
29 |
|
30 |
# Remove previous
|
31 |
for f in glob.glob('test_batch*.jpg'):
|
32 |
os.remove(f)
|
33 |
|
34 |
# Load model
|
35 |
+
model = attempt_load(weights, map_location=device) # load FP32 model
|
|
|
36 |
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
|
37 |
|
38 |
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
|
39 |
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
|
40 |
# model = nn.DataParallel(model)
|
41 |
|
|
|
|
|
|
|
|
|
42 |
# Half
|
43 |
half = device.type != 'cpu' and torch.cuda.device_count() == 1 # half precision only supported on single-GPU
|
44 |
if half:
|
|
|
53 |
niou = iouv.numel()
|
54 |
|
55 |
# Dataloader
|
56 |
+
if not training:
|
57 |
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
|
58 |
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
|
59 |
path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
|
60 |
+
dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt,
|
61 |
hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0]
|
62 |
|
63 |
seen = 0
|
|
|
190 |
if save_json and map50 and len(jdict):
|
191 |
imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files]
|
192 |
f = 'detections_val2017_%s_results.json' % \
|
193 |
+
(weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename
|
194 |
print('\nCOCO mAP with pycocotools... saving %s...' % f)
|
195 |
with open(f, 'w') as file:
|
196 |
json.dump(jdict, file)
|
|
|
223 |
|
224 |
if __name__ == '__main__':
|
225 |
parser = argparse.ArgumentParser(prog='test.py')
|
226 |
+
parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
|
227 |
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
|
228 |
parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
|
229 |
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
utils/datasets.py
CHANGED
@@ -48,7 +48,7 @@ def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=Fa
|
|
48 |
rect=rect, # rectangular training
|
49 |
cache_images=cache,
|
50 |
single_cls=opt.single_cls,
|
51 |
-
stride=stride,
|
52 |
pad=pad)
|
53 |
|
54 |
batch_size = min(batch_size, len(dataset))
|
|
|
48 |
rect=rect, # rectangular training
|
49 |
cache_images=cache,
|
50 |
single_cls=opt.single_cls,
|
51 |
+
stride=int(stride),
|
52 |
pad=pad)
|
53 |
|
54 |
batch_size = min(batch_size, len(dataset))
|