286 lines
10 KiB
Python
286 lines
10 KiB
Python
"""
|
|
Empty Shelf / Out-of-Stock (OOS) Slot Detector
|
|
Algorithm: Planogram Geometric Row-Clustering & Gap Analysis
|
|
Methodology:
|
|
1. Clusters YOLO product bounding boxes into horizontal shelf tiers/rows.
|
|
2. Filters out vertical architectural dividers / central pillars.
|
|
3. Identifies horizontal gaps exceeding expected product pitch.
|
|
4. Computes precise red bounding boxes for empty shelf slots.
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from typing import List, Dict, Any, Tuple, Optional
|
|
|
|
|
|
class EmptyShelfDetector:
|
|
def __init__(
|
|
self,
|
|
min_gap_ratio: float = 0.75,
|
|
min_row_products: int = 4,
|
|
row_height_tol: float = 0.35,
|
|
pillar_margin: int = 25
|
|
):
|
|
"""
|
|
Args:
|
|
min_gap_ratio: Minimum ratio of gap width to median item width in row (default 0.75).
|
|
min_row_products: Minimum products required to establish a valid shelf row.
|
|
row_height_tol: Vertical clustering tolerance relative to item height.
|
|
pillar_margin: Buffer zone around detected vertical pillars/dividers.
|
|
"""
|
|
self.min_gap_ratio = min_gap_ratio
|
|
self.min_row_products = min_row_products
|
|
self.row_height_tol = row_height_tol
|
|
self.pillar_margin = pillar_margin
|
|
|
|
def cluster_rows(self, boxes: List[List[int]]) -> List[List[List[int]]]:
|
|
"""Clusters bounding boxes [x1, y1, x2, y2] into horizontal shelf rows."""
|
|
if not boxes:
|
|
return []
|
|
|
|
# Sort by vertical center
|
|
sorted_boxes = sorted(boxes, key=lambda b: (b[1] + b[3]) / 2)
|
|
rows: List[List[List[int]]] = []
|
|
|
|
for b in sorted_boxes:
|
|
yc = (b[1] + b[3]) / 2
|
|
placed = False
|
|
for r in rows:
|
|
ryc = float(np.mean([(x[1] + x[3]) / 2 for x in r]))
|
|
rh = float(np.mean([x[3] - x[1] for x in r]))
|
|
if abs(yc - ryc) < self.row_height_tol * rh:
|
|
r.append(b)
|
|
placed = True
|
|
break
|
|
if not placed:
|
|
rows.append([b])
|
|
|
|
# Filter rows with sufficient products and sort each row left-to-right
|
|
valid_rows = []
|
|
for r in rows:
|
|
if len(r) >= self.min_row_products:
|
|
r.sort(key=lambda b: b[0])
|
|
valid_rows.append(r)
|
|
|
|
return valid_rows
|
|
|
|
def detect_vertical_pillars(
|
|
self,
|
|
rows: List[List[List[int]]],
|
|
image_width: int
|
|
) -> List[Tuple[int, int]]:
|
|
"""
|
|
Identifies structural vertical shelf dividers (pillars) across rows.
|
|
A pillar exhibits consistent gaps across multiple vertical shelf rows.
|
|
"""
|
|
if len(rows) < 2:
|
|
return []
|
|
|
|
# Count gap intersections across rows in 10px bins
|
|
bin_size = 10
|
|
num_bins = int(np.ceil(image_width / bin_size))
|
|
gap_votes = np.zeros(num_bins, dtype=int)
|
|
|
|
for row in rows:
|
|
for j in range(len(row) - 1):
|
|
x1 = row[j][2]
|
|
x2 = row[j + 1][0]
|
|
gap_w = x2 - x1
|
|
# Architectural dividers typically span 40px - 180px
|
|
if 40 <= gap_w <= 180:
|
|
start_bin = int(x1 // bin_size)
|
|
end_bin = min(num_bins, int(np.ceil(x2 / bin_size)))
|
|
gap_votes[start_bin:end_bin] += 1
|
|
|
|
# A pillar should have gap votes in at least 50% of the shelf rows
|
|
min_votes = max(2, int(len(rows) * 0.5))
|
|
pillars = []
|
|
in_pillar = False
|
|
p_start = 0
|
|
|
|
for b in range(num_bins):
|
|
if gap_votes[b] >= min_votes:
|
|
if not in_pillar:
|
|
in_pillar = True
|
|
p_start = b * bin_size
|
|
else:
|
|
if in_pillar:
|
|
in_pillar = False
|
|
p_end = b * bin_size
|
|
pillars.append((max(0, p_start - self.pillar_margin), min(image_width, p_end + self.pillar_margin)))
|
|
|
|
if in_pillar:
|
|
pillars.append((max(0, p_start - self.pillar_margin), image_width))
|
|
|
|
return pillars
|
|
|
|
def detect_empty_slots(
|
|
self,
|
|
image_shape: Tuple[int, int],
|
|
product_boxes: List[Any],
|
|
custom_pillars: Optional[List[Tuple[int, int]]] = None
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Detects empty shelf slots from product boxes.
|
|
|
|
Args:
|
|
image_shape: (height, width) of the shelf image.
|
|
product_boxes: List of bounding boxes [x1, y1, x2, y2] or dicts with 'box' key.
|
|
custom_pillars: Optional pre-defined pillar intervals [(x1, x2), ...].
|
|
|
|
Returns:
|
|
List of empty slot dictionaries:
|
|
[
|
|
{
|
|
"id": "empty_1",
|
|
"box": [x1, y1, x2, y2],
|
|
"width": w,
|
|
"height": h,
|
|
"missing_items_est": int,
|
|
"center": [cx, cy],
|
|
"row_index": int
|
|
}, ...
|
|
]
|
|
"""
|
|
img_h, img_w = image_shape[:2]
|
|
if not product_boxes:
|
|
return []
|
|
|
|
# Standardize box format
|
|
boxes: List[List[int]] = []
|
|
for b in product_boxes:
|
|
if isinstance(b, dict) and "box" in b:
|
|
boxes.append([int(v) for v in b["box"]])
|
|
elif isinstance(b, (list, tuple)) and len(b) >= 4:
|
|
boxes.append([int(b[0]), int(b[1]), int(b[2]), int(b[3])])
|
|
|
|
rows = self.cluster_rows(boxes)
|
|
if not rows:
|
|
return []
|
|
|
|
# Detect structural vertical pillars/frames
|
|
pillars = custom_pillars if custom_pillars is not None else self.detect_vertical_pillars(rows, img_w)
|
|
|
|
empty_slots: List[Dict[str, Any]] = []
|
|
slot_counter = 1
|
|
|
|
for row_idx, row in enumerate(rows):
|
|
med_y1 = int(np.median([b[1] for b in row]))
|
|
med_y2 = int(np.median([b[3] for b in row]))
|
|
med_w = float(np.median([b[2] - b[0] for b in row]))
|
|
med_h = med_y2 - med_y1
|
|
|
|
for j in range(len(row) - 1):
|
|
x1 = row[j][2]
|
|
x2 = row[j + 1][0]
|
|
gap_w = x2 - x1
|
|
|
|
if gap_w < self.min_gap_ratio * med_w:
|
|
continue
|
|
|
|
# Check if this gap overlaps significantly with a known structural pillar
|
|
is_pillar = False
|
|
for px1, px2 in pillars:
|
|
overlap_start = max(x1, px1)
|
|
overlap_end = min(x2, px2)
|
|
if overlap_end > overlap_start:
|
|
overlap_len = overlap_end - overlap_start
|
|
if overlap_len >= 0.6 * gap_w or (x1 >= px1 and x2 <= px2) or (x1 <= px1 and x2 >= px2):
|
|
is_pillar = True
|
|
break
|
|
|
|
if is_pillar:
|
|
continue
|
|
|
|
missing_est = max(1, int(round(gap_w / med_w)))
|
|
cx = int((x1 + x2) / 2)
|
|
cy = int((med_y1 + med_y2) / 2)
|
|
|
|
empty_slots.append({
|
|
"id": f"empty_{slot_counter}",
|
|
"box": [int(x1), int(med_y1), int(x2), int(med_y2)],
|
|
"width": int(gap_w),
|
|
"height": int(med_h),
|
|
"missing_items_est": missing_est,
|
|
"center": [cx, cy],
|
|
"row_index": row_idx + 1
|
|
})
|
|
slot_counter += 1
|
|
|
|
return empty_slots
|
|
|
|
def draw_empty_slots(
|
|
self,
|
|
image: np.ndarray,
|
|
empty_slots: List[Dict[str, Any]],
|
|
box_color: Tuple[int, int, int] = (0, 0, 255),
|
|
thickness: int = 3,
|
|
draw_fill: bool = True,
|
|
draw_label: bool = False
|
|
) -> np.ndarray:
|
|
"""
|
|
Draws red bounding boxes around empty shelf slots.
|
|
|
|
Args:
|
|
image: BGR numpy image.
|
|
empty_slots: List of empty slot dictionaries.
|
|
box_color: BGR color (default Red: (0, 0, 255)).
|
|
thickness: Bounding box line thickness.
|
|
draw_fill: Whether to render a semi-transparent fill tint.
|
|
draw_label: Whether to render label text.
|
|
"""
|
|
annotated = image.copy()
|
|
|
|
if draw_fill and empty_slots:
|
|
overlay = annotated.copy()
|
|
for slot in empty_slots:
|
|
x1, y1, x2, y2 = slot["box"]
|
|
cv2.rectangle(overlay, (x1, y1), (x2, y2), box_color, -1)
|
|
cv2.addWeighted(overlay, 0.22, annotated, 0.78, 0, annotated)
|
|
|
|
for slot in empty_slots:
|
|
x1, y1, x2, y2 = slot["box"]
|
|
cv2.rectangle(annotated, (x1, y1), (x2, y2), box_color, thickness)
|
|
if draw_label:
|
|
est = slot.get("missing_items_est", 1)
|
|
text = f"Empty (~{est})"
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
font_scale = 0.5
|
|
tw, th = cv2.getTextSize(text, font, font_scale, 1)[0]
|
|
cv2.rectangle(annotated, (x1, max(0, y1 - th - 6)), (x1 + tw + 6, y1), box_color, -1)
|
|
cv2.putText(annotated, text, (x1 + 3, y1 - 4), font, font_scale, (255, 255, 255), 1, cv2.LINE_AA)
|
|
|
|
return annotated
|
|
|
|
|
|
def detect_and_draw_empty_slots(
|
|
image_path: str,
|
|
detector_weights: str = "runs/detect/train/weights/best.pt",
|
|
output_path: Optional[str] = "detected_empty_slots.jpg"
|
|
) -> Tuple[np.ndarray, List[Dict[str, Any]]]:
|
|
"""
|
|
Convenience function: runs YOLO detection, extracts empty shelf slots,
|
|
and draws red bounding boxes.
|
|
"""
|
|
from detector import ShelfProductDetector
|
|
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
raise ValueError(f"Failed to read image at: {image_path}")
|
|
|
|
detector = ShelfProductDetector(detector_weights)
|
|
det_res = detector.detect(image)
|
|
boxes = det_res.get("boxes", [])
|
|
|
|
empty_detector = EmptyShelfDetector()
|
|
empty_slots = empty_detector.detect_empty_slots(image.shape, boxes)
|
|
|
|
# Draw products in existing clean bounding boxes (green)
|
|
annotated = detector.draw_detections(image, boxes, draw_labels=False)
|
|
# Draw empty slots in RED
|
|
annotated = empty_detector.draw_empty_slots(annotated, empty_slots, box_color=(0, 0, 255), thickness=3)
|
|
|
|
if output_path:
|
|
cv2.imwrite(output_path, annotated)
|
|
|
|
return annotated, empty_slots
|