284 lines
10 KiB
Python
284 lines
10 KiB
Python
"""
|
|
Object Detection Module using YOLO
|
|
Handles inference, bounding box extraction, and spatial hit-testing for shelf images.
|
|
"""
|
|
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
from typing import List, Dict, Optional, Tuple, Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
try:
|
|
import ultralytics.nn.modules.block as _B
|
|
import ultralytics.nn.modules as _M
|
|
|
|
if not hasattr(_B, "C3k"):
|
|
class C3k(_B.C3):
|
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, k=3):
|
|
super().__init__(c1, c2, n, shortcut, g, e)
|
|
c_ = int(c2 * e)
|
|
self.m = nn.Sequential(*(_B.Bottleneck(c_, c_, shortcut, g, k=((k, k), (k, k)), e=1.0) for _ in range(n)))
|
|
|
|
class C3k2(_B.C2f):
|
|
def __init__(self, c1, c2, n=1, c3k=False, e=0.5, g=1, shortcut=True):
|
|
super().__init__(c1, c2, n, shortcut, g, e)
|
|
self.m = nn.ModuleList(
|
|
C3k(self.c, self.c, 2, shortcut, g) if c3k else _B.Bottleneck(self.c, self.c, shortcut, g) for _ in range(n)
|
|
)
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(self, dim, num_heads=8, attn_ratio=0.5):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
self.key_dim = int(self.head_dim * attn_ratio)
|
|
self.scale = self.key_dim ** -0.5
|
|
nh_kd = self.key_dim * num_heads
|
|
h = dim + nh_kd * 2
|
|
self.qkv = _B.Conv(dim, h, 1, act=False)
|
|
self.proj = _B.Conv(dim, dim, 1, act=False)
|
|
self.pe = _B.Conv(dim, dim, 3, 1, g=dim, act=False)
|
|
|
|
def forward(self, x):
|
|
b, c, h, w = x.shape
|
|
n = h * w
|
|
qkv = self.qkv(x)
|
|
q, k, v = qkv.view(b, self.num_heads, self.key_dim * 2 + self.head_dim, n).split(
|
|
[self.key_dim, self.key_dim, self.head_dim], dim=2
|
|
)
|
|
attn = (q.transpose(-2, -1) @ k) * self.scale
|
|
attn = attn.softmax(dim=-1)
|
|
x = (v @ attn.transpose(-2, -1)).view(b, c, h, w) + self.pe(x)
|
|
return self.proj(x)
|
|
|
|
class PSABlock(nn.Module):
|
|
def __init__(self, c, attn_ratio=0.5, num_heads=4, shortcut=True):
|
|
super().__init__()
|
|
self.attn = Attention(c, attn_ratio=attn_ratio, num_heads=num_heads)
|
|
self.ffn = nn.Sequential(_B.Conv(c, c * 2, 1), _B.Conv(c * 2, c, 1, act=False))
|
|
self.add = shortcut
|
|
|
|
def forward(self, x):
|
|
x = x + self.attn(x) if self.add else self.attn(x)
|
|
x = x + self.ffn(x) if self.add else self.ffn(x)
|
|
return x
|
|
|
|
class C2PSA(nn.Module):
|
|
def __init__(self, c1, c2, n=1, e=0.5):
|
|
super().__init__()
|
|
assert c1 == c2
|
|
self.c = int(c1 * e)
|
|
self.cv1 = _B.Conv(c1, 2 * self.c, 1, 1)
|
|
self.cv2 = _B.Conv(2 * self.c, c1, 1)
|
|
self.m = nn.Sequential(*(PSABlock(self.c) for _ in range(n)))
|
|
|
|
def forward(self, x):
|
|
a, b = self.cv1(x).split((self.c, self.c), dim=1)
|
|
b = self.m(b)
|
|
return self.cv2(torch.cat((a, b), 1))
|
|
|
|
_B.C3k = C3k
|
|
_B.C3k2 = C3k2
|
|
_B.C2PSA = C2PSA
|
|
_B.PSABlock = PSABlock
|
|
_B.Attention = Attention
|
|
|
|
_M.C3k = C3k
|
|
_M.C3k2 = C3k2
|
|
_M.C2PSA = C2PSA
|
|
_M.PSABlock = PSABlock
|
|
_M.Attention = Attention
|
|
except Exception as _e:
|
|
print(f"[Warning] YOLO26 patch not applied: {_e}")
|
|
|
|
from ultralytics import YOLO
|
|
|
|
|
|
class ShelfProductDetector:
|
|
"""Wraps the YOLO detector for supermarket shelf products."""
|
|
|
|
def __init__(self, weights_path: str = "runs/detect/train/weights/best.pt"):
|
|
self.weights_path = weights_path
|
|
if not os.path.exists(weights_path):
|
|
raise FileNotFoundError(f"YOLO weights not found at: {weights_path}")
|
|
|
|
print(f"[Detector] Loading YOLO model from {weights_path}...")
|
|
self.model = YOLO(weights_path)
|
|
print("[Detector] Model loaded successfully.")
|
|
|
|
def detect(
|
|
self,
|
|
image_input: Any,
|
|
conf_threshold: float = 0.25,
|
|
iou_threshold: float = 0.45
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Runs object detection on a shelf image.
|
|
image_input can be a filepath (str) or a numpy BGR array.
|
|
"""
|
|
if isinstance(image_input, str):
|
|
if not os.path.exists(image_input):
|
|
raise FileNotFoundError(f"Image not found at: {image_input}")
|
|
image_bgr = cv2.imread(image_input)
|
|
else:
|
|
image_bgr = image_input
|
|
|
|
if image_bgr is None:
|
|
return {"status": "error", "message": "Failed to read image", "boxes": []}
|
|
|
|
h, w = image_bgr.shape[:2]
|
|
results = self.model(image_bgr, conf=conf_threshold, iou=iou_threshold, verbose=False)
|
|
|
|
boxes_list = []
|
|
if results and len(results) > 0:
|
|
boxes_data = results[0].boxes
|
|
for idx, box in enumerate(boxes_data):
|
|
xyxy = box.xyxy[0].cpu().numpy()
|
|
x1, y1, x2, y2 = map(int, xyxy)
|
|
|
|
# Clamp to image boundaries
|
|
x1, y1 = max(0, x1), max(0, y1)
|
|
x2, y2 = min(w, x2), min(h, y2)
|
|
|
|
conf = float(box.conf[0].cpu().numpy())
|
|
cls_id = int(box.cls[0].cpu().numpy()) if box.cls is not None else 0
|
|
|
|
boxes_list.append({
|
|
"id": idx,
|
|
"box": [x1, y1, x2, y2],
|
|
"confidence": round(conf, 4),
|
|
"class_id": cls_id,
|
|
"width": x2 - x1,
|
|
"height": y2 - y1,
|
|
"center": [(x1 + x2) // 2, (y1 + y2) // 2]
|
|
})
|
|
|
|
return {
|
|
"status": "success",
|
|
"image_shape": (h, w),
|
|
"total_detected": len(boxes_list),
|
|
"boxes": boxes_list
|
|
}
|
|
|
|
def find_box_at_point(
|
|
self,
|
|
x: int,
|
|
y: int,
|
|
boxes: List[Dict[str, Any]],
|
|
tolerance_px: int = 40
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Finds the bounding box enclosing the click (x, y).
|
|
If clicked on overlapping boxes, selects the one with minimum distance to center.
|
|
If clicked slightly outside, selects the nearest box within tolerance_px.
|
|
"""
|
|
enclosing_boxes = []
|
|
for b_info in boxes:
|
|
x1, y1, x2, y2 = b_info["box"]
|
|
if x1 <= x <= x2 and y1 <= y <= y2:
|
|
# Calculate distance from click to center
|
|
cx, cy = b_info["center"]
|
|
dist = np.sqrt((x - cx) ** 2 + (y - cy) ** 2)
|
|
enclosing_boxes.append((dist, b_info))
|
|
|
|
if enclosing_boxes:
|
|
# Return box closest to click center
|
|
enclosing_boxes.sort(key=lambda item: item[0])
|
|
return enclosing_boxes[0][1]
|
|
|
|
# Check nearest within tolerance
|
|
nearest_box = None
|
|
min_dist = float("inf")
|
|
for b_info in boxes:
|
|
x1, y1, x2, y2 = b_info["box"]
|
|
# Distance to rectangle
|
|
dx = max(x1 - x, 0, x - x2)
|
|
dy = max(y1 - y, 0, y - y2)
|
|
dist = np.sqrt(dx ** 2 + dy ** 2)
|
|
if dist < min_dist and dist <= tolerance_px:
|
|
min_dist = dist
|
|
nearest_box = b_info
|
|
|
|
return nearest_box
|
|
|
|
def crop_roi(self, image_bgr: np.ndarray, box: List[int]) -> np.ndarray:
|
|
"""Crops a product ROI bounding box from the shelf image."""
|
|
x1, y1, x2, y2 = box
|
|
h, w = image_bgr.shape[:2]
|
|
x1, y1 = max(0, x1), max(0, y1)
|
|
x2, y2 = min(w, x2), min(h, y2)
|
|
return image_bgr[y1:y2, x1:x2].copy()
|
|
|
|
def draw_detections(
|
|
self,
|
|
image_bgr: np.ndarray,
|
|
boxes: List[Dict[str, Any]],
|
|
selected_box_id: Optional[int] = None,
|
|
draw_labels: bool = False
|
|
) -> np.ndarray:
|
|
"""
|
|
Draws bounding boxes on the shelf image.
|
|
By default (draw_labels=False), ONLY bounding boxes are drawn:
|
|
no class names, no confidence numbers.
|
|
"""
|
|
annotated = image_bgr.copy()
|
|
for b in boxes:
|
|
x1, y1, x2, y2 = b["box"]
|
|
is_selected = (selected_box_id is not None and b["id"] == selected_box_id)
|
|
|
|
color = (0, 0, 255) if is_selected else (0, 255, 0)
|
|
thickness = 3 if is_selected else 2
|
|
|
|
# Draw ONLY bounding box
|
|
cv2.rectangle(annotated, (x1, y1), (x2, y2), color, thickness)
|
|
|
|
if draw_labels:
|
|
label = f"#{b['id']} ({b['confidence']:.2f})"
|
|
if is_selected:
|
|
label = f"SELECTED #{b['id']}"
|
|
cv2.putText(
|
|
annotated,
|
|
label,
|
|
(x1, max(15, y1 - 6)),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.5,
|
|
color,
|
|
2
|
|
)
|
|
|
|
return annotated
|
|
|
|
def detect_empty_slots(
|
|
self,
|
|
image_shape: Tuple[int, int],
|
|
boxes: List[Dict[str, Any]],
|
|
min_gap_ratio: float = 0.75
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Detects empty shelf slots / out-of-stock gaps along shelf rows.
|
|
Uses planogram geometric row-clustering and architectural pillar exclusion.
|
|
"""
|
|
from empty_shelf_detector import EmptyShelfDetector
|
|
empty_detector = EmptyShelfDetector(min_gap_ratio=min_gap_ratio)
|
|
return empty_detector.detect_empty_slots(image_shape, boxes)
|
|
|
|
def draw_empty_slots(
|
|
self,
|
|
image_bgr: np.ndarray,
|
|
empty_slots: List[Dict[str, Any]],
|
|
box_color: Tuple[int, int, int] = (0, 0, 255),
|
|
thickness: int = 3,
|
|
draw_fill: bool = True
|
|
) -> np.ndarray:
|
|
"""Draws RED bounding boxes around empty shelf spaces with semi-transparent fill."""
|
|
from empty_shelf_detector import EmptyShelfDetector
|
|
empty_detector = EmptyShelfDetector()
|
|
return empty_detector.draw_empty_slots(
|
|
image_bgr,
|
|
empty_slots,
|
|
box_color=box_color,
|
|
thickness=thickness,
|
|
draw_fill=draw_fill
|
|
)
|