Spaces:
Running
Running
File size: 1,895 Bytes
6c518b6 e2c48cc 6c518b6 e2c48cc 6c518b6 1c19e5c 9cde7c4 e2c48cc 6c518b6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import json
from pathlib import Path
from dataclasses import dataclass
from enum import Enum
import numpy as np
from ultralytics import YOLO
class WebcamMode(Enum):
GRADIO = 'gradio'
WEBRTC = 'webrtc'
class Config:
MODELS_DIR: Path = Path('models')
MODELS_DIR.mkdir(exist_ok=True)
YOLO_CLASS_NAMES: dict[str, str] = json.loads(Path('yolo_classes.json').read_text())
YOLO_CLASS_NAMES: dict[int, str] = {int(k): v for k, v in YOLO_CLASS_NAMES.items()}
MODEL_URLS: dict[str, str] = {
'yolov11n.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n.pt',
'yolov11s.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11s.pt',
'yolov11m.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11m.pt',
'yolov11l.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11l.pt',
'yolov11x.pt': 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11x.pt',
}
MODEL_NAMES: list[str] = list(MODEL_URLS.keys())
IMAGE_EXTENSIONS: list[str] = ['.jpg', '.jpeg', '.png']
VIDEO_EXTENSIONS: list[str] = ['.mp4', '.avi']
DETECT_MODE_NAMES: list[str] = ['Detection', 'Tracking']
TRACKERS: dict[str, str] = {'ByteTrack': 'bytetrack.yaml', 'BoT-SORT': 'botsort.yaml'}
TRACKER_NAMES: list[str] = list(TRACKERS.keys())
WEBCAM_TIME_LIMIT: int = 60
WEBCAM_MAX_WINDOW_WIDTH: int = 600
WEBCAM_MODE: WebcamMode = WebcamMode.GRADIO # WebcamMode.WEBRTC or WebcamMode.GRADIO
@dataclass
class DetectConfig:
source: str | np.ndarray
model: YOLO
conf: float
iou: float
detect_mode: str
tracker_name: str
verbose: bool = False
save_image_predicts: bool = False
save_video_predicts: bool = True
results_dir: str | Path = 'runs'
|