You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
51 lines
1.4 KiB
51 lines
1.4 KiB
import os
|
|
from functools import lru_cache
|
|
from typing import Iterator, Optional, Tuple
|
|
|
|
from ultralytics import YOLO
|
|
|
|
|
|
MODEL_MAP = {
|
|
"floating": "model/Floating/best.pt",
|
|
"gate": "model/gate_model/best.pt",
|
|
"shore_garbage": "model/shore_garbage_model/best.pt",
|
|
"water_gauge": "model/water_gauge_model/best.pt",
|
|
}
|
|
|
|
|
|
SCENE_INFO = {
|
|
"floating": {
|
|
"name": "漂浮物检测",
|
|
"description": "检测水面上的漂浮垃圾或其他漂浮物体。",
|
|
},
|
|
"gate": {
|
|
"name": "闸口场景检测",
|
|
"description": "检测闸口区域周边的目标和异常情况。",
|
|
},
|
|
"shore_garbage": {
|
|
"name": "岸边垃圾检测",
|
|
"description": "检测岸线附近堆积或散落的垃圾。",
|
|
},
|
|
"water_gauge": {
|
|
"name": "水尺检测",
|
|
"description": "检测水位监测场景中的水尺目标。",
|
|
},
|
|
}
|
|
|
|
|
|
def get_model_path(model_type: str) -> str:
|
|
return os.path.join(os.path.dirname(__file__), MODEL_MAP[model_type])
|
|
|
|
|
|
@lru_cache(maxsize=len(MODEL_MAP))
|
|
def load_model(model_type: str) -> YOLO:
|
|
return YOLO(get_model_path(model_type))
|
|
|
|
|
|
def iter_models(model_type: Optional[str] = None) -> Iterator[Tuple[str, YOLO]]:
|
|
if model_type is not None:
|
|
yield model_type, load_model(model_type)
|
|
return
|
|
|
|
for current_model_type in MODEL_MAP:
|
|
yield current_model_type, load_model(current_model_type)
|