scorevision: push artifact

0d266c7 verified 3 months ago

2.14 kB

from pathlib import Path

from numpy import ndarray
from pydantic import BaseModel
from ultralytics import YOLO


class BoundingBox(BaseModel):
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    conf: float


class Polygon(BaseModel):
    cls_id: int
    conf: float
    points: list[tuple[int, int]]


class TVFrameResult(BaseModel):
    frame_id: int
    boxes: list[BoundingBox] | None = None
    polygons: list[Polygon] | None = None
    keypoints: list[tuple[int, int]] | None = None


# Element: manak0/Detect-road-signs — objects: ["road sign"] (single class, id 0).
# Trained on LVIS street_sign + signboard + stop_sign, all mapped to class 0.
# CONF_THRESHOLD is set by the sweep that replicates the element's own scorer
# (0.6*mAP50 + 0.4*(1 - FP_per_image/10)).
CONF_THRESHOLD = 0.40
IMGSZ = 1280  # matches element preproc resize_long


class Miner:
    def __init__(self, path_hf_repo: Path) -> None:
        self.model = YOLO(str(path_hf_repo / "road-sign-detection.pt"))
        print("✅ Road-sign model loaded")

    def __repr__(self) -> str:
        return f"Detect-road-signs miner: {type(self.model).__name__} @ imgsz={IMGSZ}, conf={CONF_THRESHOLD}"

    def predict_batch(
        self,
        batch_images: list[ndarray],
        offset: int,
        n_keypoints: int,
    ) -> list[TVFrameResult]:
        results = self.model.predict(
            batch_images,
            imgsz=IMGSZ,
            conf=CONF_THRESHOLD,
            verbose=False,
            half=True,
        )
        out: list[TVFrameResult] = []
        for i, det in enumerate(results):
            boxes: list[BoundingBox] = []
            if det.boxes is not None:
                for box in det.boxes.data:
                    x1, y1, x2, y2, conf, cls_id = box.tolist()
                    boxes.append(
                        BoundingBox(
                            x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2),
                            cls_id=int(cls_id), conf=float(conf),
                        )
                    )
            out.append(TVFrameResult(frame_id=offset + i, boxes=boxes))
        return out

	from pathlib import Path

	from numpy import ndarray
	from pydantic import BaseModel
	from ultralytics import YOLO


	class BoundingBox(BaseModel):
	x1: int
	y1: int
	x2: int
	y2: int
	cls_id: int
	conf: float


	class Polygon(BaseModel):
	cls_id: int
	conf: float
	points: list[tuple[int, int]]


	class TVFrameResult(BaseModel):
	frame_id: int
	boxes: list[BoundingBox] \| None = None
	polygons: list[Polygon] \| None = None
	keypoints: list[tuple[int, int]] \| None = None


	# Element: manak0/Detect-road-signs — objects: ["road sign"] (single class, id 0).
	# Trained on LVIS street_sign + signboard + stop_sign, all mapped to class 0.
	# CONF_THRESHOLD is set by the sweep that replicates the element's own scorer
	# (0.6mAP50 + 0.4(1 - FP_per_image/10)).
	CONF_THRESHOLD = 0.40
	IMGSZ = 1280 # matches element preproc resize_long


	class Miner:
	def __init__(self, path_hf_repo: Path) -> None:
	self.model = YOLO(str(path_hf_repo / "road-sign-detection.pt"))
	print("✅ Road-sign model loaded")

	def __repr__(self) -> str:
	return f"Detect-road-signs miner: {type(self.model).__name__} @ imgsz={IMGSZ}, conf={CONF_THRESHOLD}"

	def predict_batch(
	self,
	batch_images: list[ndarray],
	offset: int,
	n_keypoints: int,
	) -> list[TVFrameResult]:
	results = self.model.predict(
	batch_images,
	imgsz=IMGSZ,
	conf=CONF_THRESHOLD,
	verbose=False,
	half=True,
	)
	out: list[TVFrameResult] = []
	for i, det in enumerate(results):
	boxes: list[BoundingBox] = []
	if det.boxes is not None:
	for box in det.boxes.data:
	x1, y1, x2, y2, conf, cls_id = box.tolist()
	boxes.append(
	BoundingBox(
	x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2),
	cls_id=int(cls_id), conf=float(conf),
	)
	)
	out.append(TVFrameResult(frame_id=offset + i, boxes=boxes))
	return out