Fix detect.py URL inference (#5525)
Browse files* Fix detect.py URL inference
Allows detect.py to run inference on remote URL sources, i.e.:
```python
!python detect.py --weights yolov5s.pt --source https://ultralytics.com/assets/zidane.jpg # image URL
!python detect.py --weights yolov5s.pt --source https://ultralytics.com/assets/decelera_landscape.mov # video URL
```
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
detect.py
CHANGED
@@ -24,10 +24,10 @@ if str(ROOT) not in sys.path:
|
|
24 |
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
25 |
|
26 |
from models.experimental import attempt_load
|
27 |
-
from utils.datasets import LoadImages, LoadStreams
|
28 |
-
from utils.general import (LOGGER, apply_classifier, check_img_size, check_imshow, check_requirements,
|
29 |
-
colorstr, increment_path, non_max_suppression, print_args, save_one_box,
|
30 |
-
strip_optimizer, xyxy2xywh)
|
31 |
from utils.plots import Annotator, colors
|
32 |
from utils.torch_utils import load_classifier, select_device, time_sync
|
33 |
|
@@ -61,8 +61,11 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
|
61 |
):
|
62 |
source = str(source)
|
63 |
save_img = not nosave and not source.endswith('.txt') # save inference images
|
64 |
-
|
65 |
-
|
|
|
|
|
|
|
66 |
|
67 |
# Directories
|
68 |
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
|
|
|
24 |
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
25 |
|
26 |
from models.experimental import attempt_load
|
27 |
+
from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
|
28 |
+
from utils.general import (LOGGER, apply_classifier, check_file, check_img_size, check_imshow, check_requirements,
|
29 |
+
check_suffix, colorstr, increment_path, non_max_suppression, print_args, save_one_box,
|
30 |
+
scale_coords, strip_optimizer, xyxy2xywh)
|
31 |
from utils.plots import Annotator, colors
|
32 |
from utils.torch_utils import load_classifier, select_device, time_sync
|
33 |
|
|
|
61 |
):
|
62 |
source = str(source)
|
63 |
save_img = not nosave and not source.endswith('.txt') # save inference images
|
64 |
+
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
|
65 |
+
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
|
66 |
+
webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
|
67 |
+
if is_url and is_file:
|
68 |
+
source = check_file(source) # download
|
69 |
|
70 |
# Directories
|
71 |
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
|