""" 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 )