Initial commit for push folder
This commit is contained in:
Nguyen Hong Van
commit
c4af9fcbb5
82 changed files with 3294 additions and 0 deletions
286
recognization/code/empty_shelf_detector.py
Normal file
286
recognization/code/empty_shelf_detector.py
Normal file
|
|
@ -0,0 +1,286 @@
|
|||
"""
|
||||
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
|
||||
Loading…
Add table
Add a link
Reference in a new issue