update `detect.py` in order to support torch script (#5109)
Browse files* update detect.py in order to support torch script
This change assumes the torchscrip file was previously saved with `export.py`
* update `detect.py` for torchscript support
Simple update for torchscript support. Assumes the torchscript file has been generated with `export.py`
* Cleanup
Co-authored-by: Glenn Jocher <[email protected]>
detect.py
CHANGED
|
@@ -79,7 +79,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
|
| 79 |
pt, onnx, tflite, pb, saved_model = (suffix == x for x in suffixes) # backend booleans
|
| 80 |
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
|
| 81 |
if pt:
|
| 82 |
-
model = attempt_load(weights, map_location=device)
|
| 83 |
stride = int(model.stride.max()) # model stride
|
| 84 |
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
| 85 |
if half:
|
|
|
|
| 79 |
pt, onnx, tflite, pb, saved_model = (suffix == x for x in suffixes) # backend booleans
|
| 80 |
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
|
| 81 |
if pt:
|
| 82 |
+
model = torch.jit.load(w) if 'torchscript' in w else attempt_load(weights, map_location=device)
|
| 83 |
stride = int(model.stride.max()) # model stride
|
| 84 |
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
| 85 |
if half:
|