89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
"""
|
|
Test Script for image.jpg
|
|
1. Runs YOLO detection on image.jpg.
|
|
2. Draws ONLY bounding box rectangles (no class names, no confidence labels).
|
|
3. Saves output to detected_image.jpg.
|
|
4. Matches the detected bounding boxes against stock templates.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import cv2
|
|
import numpy as np
|
|
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
if current_dir not in sys.path:
|
|
sys.path.insert(0, current_dir)
|
|
|
|
from detector import ShelfProductDetector
|
|
from feature_matcher import ProductFeatureMatcher
|
|
from inventory import InventoryManager
|
|
|
|
|
|
def main():
|
|
image_path = "image.jpg"
|
|
print(f"[1/4] Loading {image_path}...")
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
print(f"[Error] Failed to read {image_path}")
|
|
return
|
|
|
|
h, w = image.shape[:2]
|
|
print(f" Image resolution: {w}x{h}")
|
|
|
|
print("[2/4] Detecting products using YOLO...")
|
|
detector = ShelfProductDetector("runs/detect/train/weights/best.pt")
|
|
# You can adjust conf_threshold as needed (default 0.25)
|
|
det_res = detector.detect(image, conf_threshold=0.25)
|
|
boxes = det_res.get("boxes", [])
|
|
print(f" Total detected items: {len(boxes)}")
|
|
|
|
print("[3/4] Drawing bounding boxes ONLY (no class names, no confidence)...")
|
|
annotated = detector.draw_detections(image, boxes, draw_labels=False)
|
|
output_path = "detected_image.jpg"
|
|
cv2.imwrite(output_path, annotated)
|
|
print(f" Saved detection visualization to: {output_path}")
|
|
|
|
print("[4/4] Matching against stock catalogue (SURF, ORB, BRISK, Hybrid)...")
|
|
matcher = ProductFeatureMatcher("stock")
|
|
inv = InventoryManager("stock")
|
|
|
|
print("\n--- Stock Templates Present in image.jpg ---")
|
|
matched_skus = set()
|
|
sample_matches = []
|
|
|
|
# Test top detections to find corresponding stock SKUs
|
|
for b in boxes:
|
|
crop = detector.crop_roi(image, b["box"])
|
|
# Use hybrid method from the paper
|
|
res = matcher.classify_crop(crop, method="hybrid")
|
|
if res.get("status") == "success":
|
|
sku_id = res["best_sku"]
|
|
score = res["score"]
|
|
conf = res["confidence"]
|
|
prod = inv.get_product(sku_id)
|
|
if prod:
|
|
matched_skus.add(sku_id)
|
|
sample_matches.append({
|
|
"box_id": b["id"],
|
|
"box": b["box"],
|
|
"sku_id": sku_id,
|
|
"product_name": prod["product_name"],
|
|
"price": prod["price"],
|
|
"number_on": prod["number_on"],
|
|
"score": score,
|
|
"confidence": conf
|
|
})
|
|
|
|
# Sort matches by confidence
|
|
sample_matches.sort(key=lambda x: x["score"], reverse=True)
|
|
|
|
print(f"\nSuccessfully matched {len(matched_skus)} distinct SKUs across detected shelf boxes:")
|
|
for sm in sample_matches[:10]:
|
|
print(f"Box #{sm['box_id']:<3} -> [{sm['sku_id']}] {sm['product_name']:<30} | Price: ¥{sm['price']:<5} | On-Shelf: {sm['number_on']} | Score: {sm['score']}")
|
|
|
|
print("\n[Complete] Tested image.jpg successfully.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|