first
This commit is contained in:
orisys
+230
@@ -0,0 +1,230 @@
|
||||
"""统一检测器接口：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
|
||||
Reference in New Issue
Block a user