orisys
231 lines
7.5 KiB
Python
231 lines
7.5 KiB
Python
"""统一检测器接口：YOLO / Faster R-CNN，供推理与二阶段脚本共用。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
|
||
ROOT = Path(__file__).resolve().parent
|
||
|
||
|
||
@dataclass
|
||
class DetBox:
|
||
xyxy: tuple[int, int, int, int]
|
||
conf: float
|
||
cls_id: int = 0
|
||
|
||
|
||
def import_yolo():
|
||
"""导入 YOLO。项目下 ultralytics/ 源码目录会挡住包，需优先指向仓库根。"""
|
||
import sys
|
||
|
||
repo = ROOT / "ultralytics"
|
||
if (repo / "ultralytics" / "__init__.py").exists():
|
||
repo_s = str(repo)
|
||
while repo_s in sys.path:
|
||
sys.path.remove(repo_s)
|
||
sys.path.insert(0, repo_s)
|
||
|
||
mod = sys.modules.get("ultralytics")
|
||
if mod is not None and getattr(mod, "__file__", None) is None:
|
||
for key in list(sys.modules):
|
||
if key == "ultralytics" or key.startswith("ultralytics."):
|
||
del sys.modules[key]
|
||
|
||
from ultralytics import YOLO
|
||
|
||
return YOLO
|
||
|
||
|
||
def resolve_device(device: str | int) -> torch.device:
|
||
text = str(device).strip().lower()
|
||
if text in {"cpu", "-1"}:
|
||
return torch.device("cpu")
|
||
if torch.cuda.is_available():
|
||
try:
|
||
return torch.device(f"cuda:{int(text)}")
|
||
except Exception:
|
||
return torch.device("cuda:0")
|
||
return torch.device("cpu")
|
||
|
||
|
||
def build_faster_rcnn(num_classes: int, pretrained_backbone: bool = True):
|
||
"""num_classes = 前景类数 + 1（含背景）。"""
|
||
from torchvision.models.detection import (
|
||
FasterRCNN_ResNet50_FPN_Weights,
|
||
fasterrcnn_resnet50_fpn,
|
||
)
|
||
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
|
||
|
||
weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT if pretrained_backbone else None
|
||
model = fasterrcnn_resnet50_fpn(weights=weights)
|
||
in_features = model.roi_heads.box_predictor.cls_score.in_features
|
||
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
|
||
return model
|
||
|
||
|
||
class YOLODetector:
|
||
backend = "yolo"
|
||
|
||
def __init__(self, weights: str | Path, device: str | int = "0") -> None:
|
||
YOLO = import_yolo()
|
||
self.model = YOLO(str(weights))
|
||
self.device = str(device)
|
||
|
||
def detect(self, image_bgr: np.ndarray, conf: float = 0.25, iou: float = 0.7, imgsz: int = 640) -> list[DetBox]:
|
||
results = self.model.predict(
|
||
source=image_bgr,
|
||
imgsz=imgsz,
|
||
conf=conf,
|
||
iou=iou,
|
||
device=self.device,
|
||
verbose=False,
|
||
max_det=50,
|
||
)
|
||
if not results:
|
||
return []
|
||
boxes = results[0].boxes
|
||
if boxes is None or len(boxes) == 0:
|
||
return []
|
||
out: list[DetBox] = []
|
||
xyxy = boxes.xyxy.cpu().numpy()
|
||
confs = boxes.conf.cpu().numpy()
|
||
clss = boxes.cls.cpu().numpy() if boxes.cls is not None else np.zeros(len(boxes))
|
||
for i in range(len(boxes)):
|
||
x1, y1, x2, y2 = xyxy[i]
|
||
out.append(
|
||
DetBox(
|
||
xyxy=(int(round(x1)), int(round(y1)), int(round(x2)), int(round(y2))),
|
||
conf=float(confs[i]),
|
||
cls_id=int(clss[i]),
|
||
)
|
||
)
|
||
out.sort(key=lambda b: b.conf, reverse=True)
|
||
return out
|
||
|
||
def plot(self, image_bgr: np.ndarray, conf: float = 0.25, iou: float = 0.7, imgsz: int = 640) -> np.ndarray:
|
||
results = self.model.predict(
|
||
source=image_bgr,
|
||
imgsz=imgsz,
|
||
conf=conf,
|
||
iou=iou,
|
||
device=self.device,
|
||
verbose=False,
|
||
)
|
||
if not results:
|
||
return image_bgr.copy()
|
||
return results[0].plot()
|
||
|
||
|
||
class FasterRCNNDetector:
|
||
backend = "faster_rcnn"
|
||
|
||
def __init__(
|
||
self,
|
||
weights: str | Path | None = None,
|
||
device: str | int = "0",
|
||
num_classes: int = 2,
|
||
score_thresh: float = 0.25,
|
||
) -> None:
|
||
self.device = resolve_device(device)
|
||
self.score_thresh = float(score_thresh)
|
||
self.num_classes = int(num_classes)
|
||
self.model = build_faster_rcnn(self.num_classes, pretrained_backbone=weights is None)
|
||
if weights is not None:
|
||
ckpt = torch.load(str(weights), map_location="cpu", weights_only=False)
|
||
state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
|
||
self.model.load_state_dict(state)
|
||
self.model.to(self.device)
|
||
self.model.eval()
|
||
|
||
@torch.inference_mode()
|
||
def detect(self, image_bgr: np.ndarray, conf: float | None = None, **_: Any) -> list[DetBox]:
|
||
thr = self.score_thresh if conf is None else float(conf)
|
||
rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
|
||
tensor = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0
|
||
tensor = tensor.to(self.device)
|
||
out = self.model([tensor])[0]
|
||
boxes = out["boxes"].detach().cpu().numpy()
|
||
scores = out["scores"].detach().cpu().numpy()
|
||
labels = out["labels"].detach().cpu().numpy()
|
||
dets: list[DetBox] = []
|
||
for box, score, label in zip(boxes, scores, labels):
|
||
if float(score) < thr:
|
||
continue
|
||
# Faster R-CNN: 0=背景，前景从 1 开始；对外统一成 0-based 类别
|
||
cls_id = int(label) - 1 if int(label) > 0 else 0
|
||
x1, y1, x2, y2 = box
|
||
dets.append(
|
||
DetBox(
|
||
xyxy=(int(round(x1)), int(round(y1)), int(round(x2)), int(round(y2))),
|
||
conf=float(score),
|
||
cls_id=cls_id,
|
||
)
|
||
)
|
||
dets.sort(key=lambda b: b.conf, reverse=True)
|
||
return dets
|
||
|
||
def plot(self, image_bgr: np.ndarray, conf: float | None = None, **kwargs: Any) -> np.ndarray:
|
||
dets = self.detect(image_bgr, conf=conf, **kwargs)
|
||
vis = image_bgr.copy()
|
||
for det in dets:
|
||
x1, y1, x2, y2 = det.xyxy
|
||
cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||
cv2.putText(
|
||
vis,
|
||
f"id{det.cls_id} {det.conf:.2f}",
|
||
(x1, max(20, y1 - 6)),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.6,
|
||
(0, 255, 0),
|
||
2,
|
||
)
|
||
return vis
|
||
|
||
|
||
def guess_backend(weights: str | Path) -> str:
|
||
path = Path(weights)
|
||
name = path.name.lower()
|
||
suffix = path.suffix.lower()
|
||
if "faster" in name or "rcnn" in name or "frcnn" in name:
|
||
return "faster_rcnn"
|
||
if suffix == ".pt":
|
||
# Ultralytics YOLO 常用 .pt；torchvision 训练脚本默认 .pth
|
||
return "yolo"
|
||
if suffix in {".pth", ".pkl"}:
|
||
return "faster_rcnn"
|
||
return "yolo"
|
||
|
||
|
||
def load_detector(
|
||
backend: str,
|
||
weights: str | Path,
|
||
device: str | int = "0",
|
||
num_classes: int = 2,
|
||
conf: float = 0.25,
|
||
):
|
||
"""
|
||
backend: yolo | faster_rcnn | auto
|
||
Faster R-CNN 的 num_classes = 前景类数 + 1（背景）
|
||
"""
|
||
kind = backend.strip().lower()
|
||
if kind == "auto":
|
||
kind = guess_backend(weights)
|
||
if kind in {"yolo", "ultralytics"}:
|
||
return YOLODetector(weights, device=device)
|
||
if kind in {"faster_rcnn", "rcnn", "faster-rcnn", "frcnn"}:
|
||
return FasterRCNNDetector(weights, device=device, num_classes=num_classes, score_thresh=conf)
|
||
raise ValueError(f"不支持的检测后端: {backend}")
|
||
|
||
|
||
def best_box(dets: list[DetBox], conf_thr: float = 0.25) -> DetBox | None:
|
||
for det in dets:
|
||
if det.conf >= conf_thr:
|
||
return det
|
||
return None
|