xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 21:44:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Webcam inference fix (#202)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ayush Chaurasia 2023-01-10 22:59:11 +05:30 committed by GitHub
parent dba3f17884
commit e371e81aa0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 11 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
ultralytics/yolo/engine/predictor.py
View File

@ -57,7 +57,6 @@ class BasePredictor:
dataset (Dataset): Dataset used for prediction.
vid_path (str): Path to video file.
vid_writer (cv2.VideoWriter): Video writer for saving video output.
show (bool): Whether to view image output.
annotator (Annotator): Annotator used for prediction.
data_path (str): Path to data.
"""
@ -88,7 +87,6 @@ class BasePredictor:
self.device = None
self.dataset = None
self.vid_path, self.vid_writer = None, None
self.show = None
self.annotator = None
self.data_path = None
self.callbacks = defaultdict(list, {k: [v] for k, v in callbacks.default_callbacks.items()}) # add callbacks
@ -108,7 +106,7 @@ class BasePredictor:
def setup(self, source=None, model=None):
# source
source = str(source or self.args.source)
source = str(source if source is not None else self.args.source)
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)
@ -126,8 +124,10 @@ class BasePredictor:
# Dataloader
bs = 1 # batch_size
if self.args.show:
self.args.show = check_imshow(warn=True)
if webcam:
self.show = check_imshow(warn=True)
self.args.show = check_imshow(warn=True)
self.dataset = LoadStreams(source, imgsz=imgsz, stride=stride, auto=pt, vid_stride=self.args.vid_stride)
bs = len(self.dataset)
elif screenshot:

4
ultralytics/yolo/v8/classify/predict.py
View File

@ -4,7 +4,7 @@ import hydra
import torch
from ultralytics.yolo.engine.predictor import BasePredictor
from ultralytics.yolo.utils import DEFAULT_CONFIG
from ultralytics.yolo.utils import DEFAULT_CONFIG, ROOT
from ultralytics.yolo.utils.checks import check_imgsz
from ultralytics.yolo.utils.plotting import Annotator
@ -59,6 +59,8 @@ class ClassificationPredictor(BasePredictor):
def predict(cfg):
cfg.model = cfg.model or "squeezenet1_0"
cfg.imgsz = check_imgsz(cfg.imgsz, min_dim=2) # check image size
cfg.source = cfg.source if cfg.source is not None else ROOT / "assets"
predictor = ClassificationPredictor(cfg)
predictor()

2
ultralytics/yolo/v8/detect/predict.py
View File

@ -87,7 +87,7 @@ class DetectionPredictor(BasePredictor):
def predict(cfg):
cfg.model = cfg.model or "yolov8n.pt"
cfg.imgsz = check_imgsz(cfg.imgsz, min_dim=2) # check image size
cfg.source = cfg.source or ROOT / "assets"
cfg.source = cfg.source if cfg.source is not None else ROOT / "assets"
predictor = DetectionPredictor(cfg)
predictor()

4
ultralytics/yolo/v8/segment/predict.py
View File

@ -3,7 +3,7 @@
import hydra
import torch
from ultralytics.yolo.utils import DEFAULT_CONFIG, ops
from ultralytics.yolo.utils import DEFAULT_CONFIG, ROOT, ops
from ultralytics.yolo.utils.checks import check_imgsz
from ultralytics.yolo.utils.plotting import colors, save_one_box
@ -103,6 +103,8 @@ class SegmentationPredictor(DetectionPredictor):
def predict(cfg):
cfg.model = cfg.model or "yolov8n-seg.pt"
cfg.imgsz = check_imgsz(cfg.imgsz, min_dim=2) # check image size
cfg.source = cfg.source if cfg.source is not None else ROOT / "assets"
predictor = SegmentationPredictor(cfg)
predictor()