""" ============================================================================= ShelfMind AI — FastAPI Backend (HuggingFace Spaces Docker Deployment) ============================================================================= Production-ready API serving YOLO26s detection, DINOv2 embeddings, FAISS product matching, rembg auto-crop, and EasyOCR. Run locally: python api_server.py HF Spaces: Auto-starts via Dockerfile CMD ============================================================================= """ import os import sys import io import gc import re import json import time import base64 import hashlib import asyncio import logging from pathlib import Path from datetime import datetime from typing import Optional, List from contextlib import asynccontextmanager from collections import Counter import numpy as np from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageEnhance, ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True import torch import torch.nn.functional as F # ── FastAPI ─────────────────────────────────────────────── from fastapi import FastAPI, UploadFile, File, Form, HTTPException, WebSocket, WebSocketDisconnect from fastapi.responses import HTMLResponse, JSONResponse, FileResponse, Response from fastapi.middleware.cors import CORSMiddleware import uvicorn # ── Paths ───────────────────────────────────────────────── ROOT = Path(__file__).resolve().parent MODEL_DIR = ROOT / "models" DATA_DIR = ROOT / "data" CATALOG_DIR = DATA_DIR / "store_catalog" REF_IMG_DIR = CATALOG_DIR / "reference_images" PLANOGRAM_DIR = DATA_DIR / "store_planograms" for d in [DATA_DIR, CATALOG_DIR, REF_IMG_DIR, PLANOGRAM_DIR]: d.mkdir(parents=True, exist_ok=True) # ── Logging ─────────────────────────────────────────────── logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") logger = logging.getLogger("shelfmind") # ============================================================================= # MODEL MANAGER (Singleton — loaded once at startup) # ============================================================================= class ModelManager: """Manages all ML models. Loads once, serves forever.""" def __init__(self): self.yolo = None self.dinov2 = None self.projector = None self.ocr_reader = None self.rembg_session = None self.device = "cuda" if torch.cuda.is_available() else "cpu" self.ready = False def load_all(self): """Load all models at startup.""" start = time.time() logger.info(f"Loading models on device={self.device}...") self._load_yolo() self._load_dinov2() self._load_ocr() self._load_rembg() self.ready = True elapsed = time.time() - start logger.info(f"[OK] All models loaded in {elapsed:.1f}s") def _load_yolo(self): yolo_path = MODEL_DIR / "yolo_shelf_best.pt" if yolo_path.exists(): from ultralytics import YOLO self.yolo = YOLO(str(yolo_path)) logger.info(f" [OK] YOLO26s: {yolo_path.stat().st_size/1e6:.1f} MB") else: logger.warning(f" [SKIP] YOLO not found at {yolo_path}") def _load_dinov2(self): finetuned_path = MODEL_DIR / "dinov2_shelf_finetuned.pth" projector_path = MODEL_DIR / "dinov2_projector.pth" try: # Load base DINOv2 from torch.hub (cached from Docker build) self.dinov2 = torch.hub.load( "facebookresearch/dinov2", "dinov2_vitb14", pretrained=True, verbose=False ) # Load our fine-tuned weights on top if finetuned_path.exists(): state = torch.load(str(finetuned_path), map_location=self.device, weights_only=True) self.dinov2.load_state_dict(state, strict=False) logger.info(f" [OK] DINOv2: fine-tuned weights loaded ({finetuned_path.stat().st_size/1e6:.1f} MB)") else: logger.info(" [OK] DINOv2: using pretrained base (no fine-tuned weights)") self.dinov2.eval().to(self.device) # Load projector head if available if projector_path.exists(): self.projector = torch.load(str(projector_path), map_location=self.device, weights_only=True) logger.info(f" [OK] DINOv2 projector loaded ({projector_path.stat().st_size/1e6:.1f} MB)") except Exception as e: logger.warning(f" [SKIP] DINOv2 load error: {e}") self.dinov2 = None def _load_ocr(self): try: import easyocr self.ocr_reader = easyocr.Reader(["en"], gpu=(self.device == "cuda")) logger.info(" [OK] EasyOCR loaded") except Exception as e: logger.warning(f" [SKIP] EasyOCR error: {e}") def _load_rembg(self): try: from rembg import new_session self.rembg_session = new_session("u2net") logger.info(" [OK] rembg (U2-Net) loaded") except Exception as e: logger.warning(f" [SKIP] rembg error: {e}") # ── Inference Methods ───────────────────────────────── def detect(self, image: Image.Image, conf: float = 0.3, max_det: int = 200) -> list: """Run YOLO26s detection. Returns list of detections.""" if not self.yolo: raise RuntimeError("YOLO model not loaded") results = self.yolo.predict(image, conf=conf, max_det=max_det, device="cpu", verbose=False) detections = [] for r in results: for box in r.boxes: x1, y1, x2, y2 = box.xyxy[0].tolist() detections.append({ "bbox": [int(x1), int(y1), int(x2), int(y2)], "confidence": round(float(box.conf[0]), 3), "class": int(box.cls[0]) if box.cls is not None else 0, }) return detections def get_embedding(self, crop: Image.Image) -> np.ndarray: """Get DINOv2 embedding for a product crop. Returns L2-normalized 768-dim vector.""" if not self.dinov2: return np.zeros(768, dtype=np.float32) from torchvision import transforms transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) tensor = transform(crop.convert("RGB")).unsqueeze(0).to(self.device) with torch.no_grad(): emb = self.dinov2(tensor).squeeze().cpu().numpy() # L2 normalize norm = np.linalg.norm(emb) if norm > 0: emb = emb / norm return emb def get_robust_embedding(self, image: Image.Image): """Generate robust embedding by averaging 10 augmented views. Returns (embedding, view_names). """ img = image.convert("RGB") w, h = img.size views = [] view_names = [] # 1. Original views.append(img) view_names.append("Original") # 2. Horizontal flip views.append(ImageOps.mirror(img)) view_names.append("H-Flip") # 3. Center crop 80% mw, mh = int(w * 0.1), int(h * 0.1) views.append(img.crop((mw, mh, w - mw, h - mh))) view_names.append("Center 80%") # 4. Left crop views.append(img.crop((0, mh, int(w * 0.85), h - mh))) view_names.append("Left Crop") # 5. Right crop views.append(img.crop((int(w * 0.15), mh, w, h - mh))) view_names.append("Right Crop") # 6. Top crop views.append(img.crop((mw, 0, w - mw, int(h * 0.85)))) view_names.append("Top Crop") # 7. Rotation +5° views.append(img.rotate(-5, expand=False, fillcolor=(128, 128, 128))) view_names.append("Rotate +5°") # 8. Rotation -5° views.append(img.rotate(5, expand=False, fillcolor=(128, 128, 128))) view_names.append("Rotate -5°") # 9. Brightness +20% views.append(ImageEnhance.Brightness(img).enhance(1.2)) view_names.append("Bright +20%") # 10. Contrast +20% views.append(ImageEnhance.Contrast(img).enhance(1.2)) view_names.append("Contrast +20%") # Average all embeddings embeddings = [self.get_embedding(v) for v in views] avg_emb = np.mean(embeddings, axis=0) norm = np.linalg.norm(avg_emb) if norm > 0: avg_emb = avg_emb / norm return avg_emb, view_names def auto_crop_rembg(self, image: Image.Image) -> Image.Image: """Remove background with rembg and crop to product bounds.""" if not self.rembg_session: return image from rembg import remove result = remove(image, session=self.rembg_session) alpha = np.array(result.split()[-1]) coords = np.where(alpha > 30) if len(coords[0]) > 0: y_min, y_max = coords[0].min(), coords[0].max() x_min, x_max = coords[1].min(), coords[1].max() pad = 5 h, w = alpha.shape x1 = max(0, x_min - pad) y1 = max(0, y_min - pad) x2 = min(w, x_max + pad) y2 = min(h, y_max + pad) return image.crop((x1, y1, x2, y2)) return image def read_ocr(self, image: Image.Image) -> str: """Extract text from product image using EasyOCR.""" if not self.ocr_reader: return "" try: img_np = np.array(image) results = self.ocr_reader.readtext(img_np, detail=0, paragraph=True) return " ".join(results).strip() except Exception: return "" def scan_barcode(self, image: Image.Image) -> str: """Detect barcode from image using pyzbar with multiple rotations.""" try: from pyzbar.pyzbar import decode import cv2 img_cv = np.array(image) # Try 4 rotations × 2 preprocessing for angle in [0, 90, 180, 270]: for preprocess in [False, True]: rotated = img_cv if angle > 0: if angle == 90: rotated = cv2.rotate(img_cv, cv2.ROTATE_90_CLOCKWISE) elif angle == 180: rotated = cv2.rotate(img_cv, cv2.ROTATE_180) elif angle == 270: rotated = cv2.rotate(img_cv, cv2.ROTATE_90_COUNTERCLOCKWISE) if preprocess: gray = cv2.cvtColor(rotated, cv2.COLOR_RGB2GRAY) rotated = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2) rotated = cv2.cvtColor(rotated, cv2.COLOR_GRAY2RGB) barcodes = decode(rotated) if barcodes: return barcodes[0].data.decode("utf-8") return "" except Exception: return "" def lookup_barcode(self, barcode: str) -> dict: """Look up product info from Open Food Facts API.""" try: import requests resp = requests.get( f"https://world.openfoodfacts.org/api/v0/product/{barcode}.json", timeout=5 ) if resp.status_code == 200: data = resp.json() if data.get("status") == 1: product = data.get("product", {}) return { "name": product.get("product_name", ""), "category": product.get("categories_tags", [""])[0].replace("en:", "").replace("-", " ").title() if product.get("categories_tags") else "", "brand": product.get("brands", ""), } except Exception: pass return {} def detect_shelf_levels(self, detections: list, img_height: int): """Auto-detect shelf levels from product Y-positions.""" if len(detections) < 3: return [0.0, float(img_height)], 1 y_centers = sorted([(d["bbox"][1] + d["bbox"][3]) / 2 for d in detections]) gaps = [(y_centers[i] - y_centers[i-1], i) for i in range(1, len(y_centers))] min_shelf_gap = img_height * 0.15 small_gaps = sorted([g for g, _ in gaps])[:max(len(gaps)//2, 1)] avg_small = sum(small_gaps) / len(small_gaps) if small_gaps else 10 stat_threshold = avg_small * 3 sig_threshold = max(min_shelf_gap, stat_threshold) gaps_sorted = sorted(gaps, key=lambda x: x[0], reverse=True) top_gaps = [(g, idx) for g, idx in gaps_sorted if g > sig_threshold][:6] top_gaps_sorted = sorted(top_gaps, key=lambda x: y_centers[x[1]]) boundaries = [0.0] for _, idx in top_gaps_sorted: boundaries.append((y_centers[idx-1] + y_centers[idx]) / 2) boundaries.append(float(img_height)) return boundaries, len(boundaries) - 1 def assign_to_shelves(self, detections: list, boundaries: list) -> dict: """Assign each detection to a shelf level.""" shelf_assignments = {} for det in detections: y_center = (det["bbox"][1] + det["bbox"][3]) / 2 for s in range(len(boundaries) - 1): if boundaries[s] <= y_center < boundaries[s + 1]: shelf_id = s + 1 if shelf_id not in shelf_assignments: shelf_assignments[shelf_id] = [] det["shelf"] = shelf_id shelf_assignments[shelf_id].append(det) break for shelf_id in shelf_assignments: shelf_assignments[shelf_id].sort(key=lambda d: d["bbox"][0]) return shelf_assignments def search_product_faiss(self, query_embedding, products, threshold=0.3, query_bbox=None, expected_products=None, size_weight=0.15): """Match product embedding against catalog using FAISS-style cosine similarity with size-ratio fusion.""" if not products: return None, 0.0 # Build embeddings matrix from catalog valid_products = [p for p in products if p.get("embedding")] if not valid_products: return None, 0.0 catalog_embeddings = np.array([p["embedding"] for p in valid_products], dtype=np.float32) query = np.array(query_embedding, dtype=np.float32).reshape(1, -1) # Cosine similarity norms_cat = np.linalg.norm(catalog_embeddings, axis=1, keepdims=True) norms_cat[norms_cat == 0] = 1 catalog_norm = catalog_embeddings / norms_cat norm_q = np.linalg.norm(query) if norm_q > 0: query_norm = query / norm_q else: query_norm = query scores = (catalog_norm @ query_norm.T).flatten() # Get top-3 candidates k = min(3, len(scores)) top_indices = np.argsort(scores)[::-1][:k] if scores[top_indices[0]] < threshold: return None, float(scores[top_indices[0]]) if expected_products is None or query_bbox is None: best_idx = top_indices[0] return valid_products[best_idx], float(scores[best_idx]) # Size-ratio fusion qx1, qy1, qx2, qy2 = query_bbox query_height = qy2 - qy1 best_match = None best_combined = 0.0 for idx in top_indices: visual_score = float(scores[idx]) if visual_score < threshold * 0.8: continue candidate = valid_products[idx] candidate_sku = candidate["sku"] size_score = 1.0 for exp_prod in expected_products: if exp_prod.get("sku") == candidate_sku: exp_bbox = exp_prod.get("bbox") if exp_bbox and len(exp_bbox) == 4: exp_height = exp_bbox[3] - exp_bbox[1] if exp_height > 0 and query_height > 0: size_score = min(query_height, exp_height) / max(query_height, exp_height) break combined = visual_score * (1 - size_weight) + size_score * size_weight if combined > best_combined: best_combined = combined best_match = candidate if best_match and best_combined >= threshold: return best_match, best_combined return valid_products[top_indices[0]], float(scores[top_indices[0]]) def draw_annotated(self, image: Image.Image, detections: list) -> Image.Image: """Draw bounding boxes with color-coded labels.""" img = image.copy() draw = ImageDraw.Draw(img) try: font = ImageFont.truetype("arial.ttf", 14) except: font = ImageFont.load_default() for det in detections: x1, y1, x2, y2 = det["bbox"] score = det.get("match_score", det.get("confidence", 0)) if det.get("status") == "unknown" or score < 0.3: color = "#ff4343" elif score >= 0.7: color = "#00d4aa" elif score >= 0.5: color = "#ffaa00" else: color = "#ff8c00" draw.rectangle([x1, y1, x2, y2], outline=color, width=2) label = det.get("product_name", f"P{det.get('confidence', 0):.0%}") text = f"{label} ({score:.0%})" bbox = draw.textbbox((x1, y1 - 18), text, font=font) draw.rectangle([bbox[0]-2, bbox[1]-2, bbox[2]+2, bbox[3]+2], fill=color) draw.text((x1, y1 - 18), text, fill="#0a0a1a", font=font) return img # ============================================================================= # DATABASE LAYER # ============================================================================= sys.path.insert(0, str(ROOT)) try: from db import ( setup_database, add_product, get_products, get_product_count, get_next_product_id, delete_product, clear_all_products, get_catalog_as_dict, save_planogram_db, get_planograms, delete_planogram, log_compliance, log_alert, get_compliance_logs_as_list, get_analytics_summary, get_alerts_history, ) setup_database() logger.info("[OK] Database initialized") except ImportError as e: logger.warning(f"[SKIP] db.py not found: {e}") # Minimal fallback _products = [] def get_products(): return _products def add_product(**kwargs): _products.append(kwargs); return kwargs.get("sku") def get_product_count(): return len(_products) def get_next_product_id(): return len(_products) + 1 def delete_product(sku): pass def clear_all_products(): _products.clear() def get_catalog_as_dict(): return {"products": _products, "next_id": len(_products)+1} def get_planograms(): return {} def save_planogram_db(name, data): pass def delete_planogram(name): pass def get_compliance_logs_as_list(): return [] def get_analytics_summary(): return {} def get_alerts_history(): return [] def log_compliance(**kwargs): return 0 def log_alert(**kwargs): pass # ============================================================================= # FASTAPI APP # ============================================================================= models = ModelManager() @asynccontextmanager async def lifespan(app: FastAPI): models.load_all() yield del models.yolo, models.dinov2 torch.cuda.empty_cache() gc.collect() app = FastAPI( title="ShelfMind AI", description="Smart Retail Shelf Intelligence API — YOLO26s + DINOv2 + FAISS", version="2.0.0", lifespan=lifespan, ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ── Helpers ─────────────────────────────────────────────── def image_to_base64(img: Image.Image, fmt="JPEG", quality=85) -> str: buf = io.BytesIO() img.save(buf, format=fmt, quality=quality) return base64.b64encode(buf.getvalue()).decode() def read_upload_image(contents: bytes) -> Image.Image: return Image.open(io.BytesIO(contents)).convert("RGB") # ============================================================================= # API ENDPOINTS # ============================================================================= @app.get("/") async def root(): return {"status": "ShelfMind AI Backend is running", "version": "2.0.0"} @app.get("/api/health") async def health(): return { "status": "healthy" if models.ready else "loading", "device": models.device, "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU", "models": { "yolo": models.yolo is not None, "dinov2": models.dinov2 is not None, "ocr": models.ocr_reader is not None, "rembg": models.rembg_session is not None, }, "products_registered": get_product_count(), "timestamp": datetime.now().isoformat(), } # ── Product Scanner: Single Product ────────────────────── @app.post("/api/scan/single") async def scan_single_product(image: UploadFile = File(...)): """Full single product scan: rembg crop → barcode → OCR → robust embedding.""" contents = await image.read() img = read_upload_image(contents) # 1. Barcode scan barcode = models.scan_barcode(img) barcode_info = {} if barcode: barcode_info = models.lookup_barcode(barcode) # 2. OCR ocr_text = models.read_ocr(img) # 3. Auto-crop with rembg cropped = models.auto_crop_rembg(img) # 4. Robust embedding (10 augmented views) embedding, view_names = models.get_robust_embedding(cropped) return { "barcode": barcode, "barcode_info": barcode_info, "ocr_text": ocr_text, "cropped_image": image_to_base64(cropped), "original_image": image_to_base64(img), "embedding": embedding.tolist(), "augmentation_views": view_names, "embedding_dim": len(embedding), } # ── Product Scanner: Bulk Shelf Scan ───────────────────── @app.post("/api/scan/bulk") async def scan_bulk_shelf(image: UploadFile = File(...), similarity_threshold: float = Form(0.82)): """Bulk shelf scan: YOLO detect → DINOv2 cluster unique products → OCR.""" contents = await image.read() img = read_upload_image(contents) # Step 1: Detect all products detections = models.detect(img, conf=0.3) # Step 2: Crop and embed each crops_data = [] for det in detections: x1, y1, x2, y2 = det["bbox"] crop = img.crop((max(0, x1), max(0, y1), x2, y2)) emb = models.get_embedding(crop) crops_data.append({ "crop": crop, "bbox": det["bbox"], "confidence": det["confidence"], "embedding": emb, }) # Step 3: Cluster unique products by cosine similarity unique_products = [] used = set() for i, cd in enumerate(crops_data): if i in used: continue cluster = [i] for j in range(i + 1, len(crops_data)): if j in used: continue sim = float(np.dot(cd["embedding"], crops_data[j]["embedding"])) if sim >= similarity_threshold: cluster.append(j) used.add(j) used.add(i) # OCR on representative crop ocr_text = models.read_ocr(cd["crop"]) unique_products.append({ "crop_image": image_to_base64(cd["crop"]), "bbox": cd["bbox"], "confidence": cd["confidence"], "embedding": cd["embedding"].tolist(), "ocr_text": ocr_text, "count": len(cluster), "geometry": { "width": cd["bbox"][2] - cd["bbox"][0], "height": cd["bbox"][3] - cd["bbox"][1], }, }) # Draw annotated image annotated = img.copy() draw = ImageDraw.Draw(annotated) for det in detections: x1, y1, x2, y2 = det["bbox"] draw.rectangle([x1, y1, x2, y2], outline="lime", width=2) return { "total_detected": len(detections), "unique_products": len(unique_products), "products": unique_products, "annotated_image": image_to_base64(annotated), } # ── Product CRUD ────────────────────────────────────────── @app.get("/api/products") async def list_products(): """List all registered products.""" products = get_products() result = [] for p in products: prod = dict(p) if hasattr(p, 'keys') else p # Convert embedding to flag (don't send full embedding) prod["has_embedding"] = prod.get("embedding") is not None and len(prod.get("embedding", [])) > 0 if "embedding" in prod: del prod["embedding"] # Add image URL if prod.get("image_path"): prod["image_url"] = f"/api/products/{prod['sku']}/image" result.append(prod) return {"products": result, "count": len(result)} @app.post("/api/products") async def create_product( name: str = Form(...), category: str = Form("Other"), price: float = Form(0.0), barcode: str = Form(""), embedding: str = Form(""), # JSON string image: Optional[UploadFile] = File(None), ): """Register a new product.""" next_id = get_next_product_id() sku = f"SKU_{next_id:04d}" # Save image img_filename = "" if image: contents = await image.read() img = read_upload_image(contents) img_filename = f"{sku}_{re.sub(r'[^a-z0-9_]', '', name.replace(' ', '_').lower())}.jpg" img_path = REF_IMG_DIR / img_filename img.save(str(img_path), "JPEG", quality=90) # Parse embedding emb = None if embedding: try: emb = json.loads(embedding) except json.JSONDecodeError: pass add_product( sku=sku, name=name, category=category, price=price, image_path=img_filename, embedding=emb, barcode=barcode if barcode else None, ) return {"status": "ok", "sku": sku, "name": name} @app.get("/api/products/{sku}/image") async def get_product_image(sku: str): """Serve product reference image.""" products = get_products() for p in products: if p["sku"] == sku and p.get("image_path"): img_path = REF_IMG_DIR / p["image_path"] if img_path.exists(): return FileResponse(str(img_path), media_type="image/jpeg") raise HTTPException(404, "Image not found") @app.delete("/api/products/{sku}") async def remove_product(sku: str): """Delete a product by SKU.""" delete_product(sku) return {"status": "ok"} @app.delete("/api/products") async def clear_products(): """Delete all products.""" clear_all_products() return {"status": "ok"} @app.post("/api/products/voice") async def voice_product(transcript: str = Form(...)): """Parse voice transcript into product fields.""" text = transcript.lower().strip() # Extract price price = 0.0 price_match = re.search(r'(?:price|rs|rupees?|₹)\s*(\d+(?:\.\d+)?)', text) if not price_match: price_match = re.search(r'(\d+(?:\.\d+)?)\s*(?:rupees?|rs)', text) if price_match: price = float(price_match.group(1)) # Extract category category = "Other" categories = { "snack": "Snacks", "biscuit": "Snacks", "chip": "Snacks", "namkeen": "Snacks", "drink": "Beverages", "juice": "Beverages", "water": "Beverages", "cola": "Beverages", "dairy": "Dairy", "milk": "Dairy", "curd": "Dairy", "paneer": "Dairy", "cheese": "Dairy", "soap": "Personal Care", "shampoo": "Personal Care", "cream": "Personal Care", "detergent": "Household", "cleaner": "Household", "wash": "Household", "rice": "Staples", "flour": "Staples", "atta": "Staples", "dal": "Staples", "chocolate": "Confectionery", "candy": "Confectionery", "spice": "Spices", "masala": "Spices", } for keyword, cat in categories.items(): if keyword in text: category = cat break # Name = transcript minus price/category mentions name_text = text for pattern in [r'(?:price|rs|rupees?|₹)\s*\d+(?:\.\d+)?', r'\d+(?:\.\d+)?\s*(?:rupees?|rs)', r'category\s+\w+']: name_text = re.sub(pattern, '', name_text, flags=re.I) name = name_text.strip().strip(',').strip().title() or "Unknown Product" return {"parsed": {"name": name, "category": category, "price": price}} # ── Planogram CRUD ──────────────────────────────────────── @app.post("/api/planogram/auto-detect") async def auto_detect_planogram(image: UploadFile = File(...), confidence: float = Form(0.25)): """Upload shelf image → detect products → identify → return planogram data.""" contents = await image.read() img = read_upload_image(contents) # Detect detections = models.detect(img, conf=confidence) boundaries, n_shelves = models.detect_shelf_levels(detections, img.height) shelf_assignments = models.assign_to_shelves(detections, boundaries) # Identify each product catalog_products = get_products() for det in detections: x1, y1, x2, y2 = det["bbox"] crop = img.crop((max(0, x1), max(0, y1), x2, y2)) emb = models.get_embedding(crop) match, score = models.search_product_faiss( emb, catalog_products, threshold=0.3, query_bbox=(x1, y1, x2, y2), ) if match: det["product_name"] = match["name"] det["product_sku"] = match["sku"] det["match_score"] = round(score, 3) else: det["product_name"] = "Unknown" det["product_sku"] = "UNKNOWN" det["match_score"] = round(score, 3) # Build planogram structure shelves = [] for shelf_id in sorted(shelf_assignments.keys()): shelf_dets = shelf_assignments[shelf_id] products_on_shelf = [] for pos, det in enumerate(shelf_dets): products_on_shelf.append({ "position": pos, "sku": det.get("product_sku", "UNKNOWN"), "name": det.get("product_name", "Unknown"), "confidence": det.get("match_score", 0), "bbox": det["bbox"], }) shelves.append({"level": shelf_id, "product_count": len(shelf_dets), "products": products_on_shelf}) # Annotated image annotated = models.draw_annotated(img, detections) return { "n_shelves": n_shelves, "total_products": len(detections), "shelves": shelves, "annotated_image": image_to_base64(annotated), } @app.post("/api/planograms") async def save_planogram( name: str = Form(...), data: str = Form(...), # JSON string of planogram data image: Optional[UploadFile] = File(None), ): """Save a planogram.""" planogram_data = json.loads(data) planogram_data["name"] = name planogram_data["created_at"] = datetime.now().isoformat() save_planogram_db(name, planogram_data) # Save reference image if image: contents = await image.read() img = read_upload_image(contents) img.save(str(PLANOGRAM_DIR / f"{name}_reference.jpg"), "JPEG", quality=90) return {"status": "ok", "name": name} @app.get("/api/planograms") async def list_planograms(): """List all planograms.""" planograms = get_planograms() return {"planograms": planograms} @app.delete("/api/planograms/{name}") async def remove_planogram(name: str): """Delete a planogram.""" delete_planogram(name) ref_img = PLANOGRAM_DIR / f"{name}_reference.jpg" if ref_img.exists(): ref_img.unlink() return {"status": "ok"} # ── Compliance Check ────────────────────────────────────── @app.post("/api/compliance/check") async def check_compliance( image: UploadFile = File(...), planogram_name: str = Form(...), confidence: float = Form(0.3), ): """Run compliance check: detect → identify → compare vs planogram.""" contents = await image.read() img = read_upload_image(contents) # Resize for performance max_dim = 640 if img.width > max_dim: ratio = max_dim / img.width img = img.resize((max_dim, int(img.height * ratio)), Image.LANCZOS) # Get planogram planograms = get_planograms() if planogram_name not in planograms: raise HTTPException(404, f"Planogram '{planogram_name}' not found") planogram = planograms[planogram_name] # Detect detections = models.detect(img, conf=confidence) boundaries, n_shelves = models.detect_shelf_levels(detections, img.height) shelf_assignments = models.assign_to_shelves(detections, boundaries) # Identify products catalog_products = get_products() all_expected = [] for ps in planogram.get("shelves", []): all_expected.extend(ps.get("products", [])) for det in detections: x1, y1, x2, y2 = det["bbox"] crop = img.crop((max(0, x1), max(0, y1), x2, y2)) emb = models.get_embedding(crop) match, score = models.search_product_faiss( emb, catalog_products, threshold=0.3, query_bbox=(x1, y1, x2, y2), expected_products=all_expected, size_weight=0.15, ) if match: det["product_name"] = match["name"] det["product_sku"] = match["sku"] det["product_price"] = match.get("price", 0) det["match_score"] = round(score, 3) det["status"] = "match" else: det["product_name"] = "Unknown" det["product_sku"] = "UNKNOWN" det["product_price"] = 0 det["match_score"] = round(score, 3) det["status"] = "unknown" # Compliance comparison all_alerts = [] shelf_compliance = {} total_expected = 0 total_matched = 0 for plan_shelf in planogram.get("shelves", []): shelf_id = plan_shelf["level"] expected_products = plan_shelf.get("products", []) detected_on_shelf = shelf_assignments.get(shelf_id, []) expected_counts = Counter(p["sku"] for p in expected_products if p["sku"] != "UNKNOWN") detected_counts = Counter( d.get("product_sku", "UNKNOWN") for d in detected_on_shelf if d.get("product_sku") != "UNKNOWN" ) issues = [] shelf_expected = len(expected_products) shelf_matched = 0 revenue_at_risk = 0 for sku, exp_count in expected_counts.items(): det_count = detected_counts.get(sku, 0) prod_name = next((p["name"] for p in expected_products if p["sku"] == sku), sku) prod_price = next((d.get("product_price", 0) for d in detected_on_shelf if d.get("product_sku") == sku), 0) if det_count == 0: issues.append({"type": "STOCKOUT", "product": prod_name, "expected": exp_count, "found": 0}) revenue_at_risk += prod_price * exp_count all_alerts.append({"type": "STOCKOUT", "shelf": shelf_id, "product": prod_name, "sku": sku, "expected": exp_count, "found": 0, "revenue": prod_price * exp_count, "priority": "CRITICAL"}) elif det_count < exp_count: missing = exp_count - det_count issues.append({"type": "LOW_STOCK", "product": prod_name, "expected": exp_count, "found": det_count}) revenue_at_risk += prod_price * missing all_alerts.append({"type": "LOW_STOCK", "shelf": shelf_id, "product": prod_name, "sku": sku, "expected": exp_count, "found": det_count, "revenue": prod_price * missing, "priority": "HIGH"}) shelf_matched += det_count else: shelf_matched += exp_count for sku, count in detected_counts.items(): if sku not in expected_counts: prod_name = next((d.get("product_name", sku) for d in detected_on_shelf if d.get("product_sku") == sku), sku) issues.append({"type": "UNAUTHORIZED", "product": prod_name}) all_alerts.append({"type": "UNAUTHORIZED", "shelf": shelf_id, "product": prod_name, "sku": sku, "priority": "MEDIUM"}) # Position check expected_order = [p["sku"] for p in expected_products if p["sku"] != "UNKNOWN"] detected_sorted = sorted( [d for d in detected_on_shelf if d.get("product_sku", "UNKNOWN") != "UNKNOWN"], key=lambda d: d["bbox"][0] ) detected_order = [d.get("product_sku") for d in detected_sorted] if expected_order and detected_order: for pos_idx in range(min(len(expected_order), len(detected_order))): if expected_order[pos_idx] != detected_order[pos_idx]: exp_name = next((p["name"] for p in expected_products if p["sku"] == expected_order[pos_idx]), "?") det_name = next((d.get("product_name", "?") for d in detected_sorted if d.get("product_sku") == detected_order[pos_idx]), "?") issues.append({"type": "MISPLACED", "product": det_name, "expected_product": exp_name, "position": pos_idx + 1}) all_alerts.append({"type": "MISPLACED", "shelf": shelf_id, "product": det_name, "priority": "HIGH"}) if shelf_matched > 0: shelf_matched -= 0.5 if not issues: issues.append({"type": "OK", "product": "All products in correct position"}) comp_pct = (shelf_matched / shelf_expected * 100) if shelf_expected > 0 else 100 shelf_compliance[shelf_id] = { "compliance": round(comp_pct, 1), "expected": shelf_expected, "detected": len(detected_on_shelf), "matched": shelf_matched, "issues": issues, "revenue_at_risk": revenue_at_risk, } total_expected += shelf_expected total_matched += shelf_matched overall = (total_matched / total_expected * 100) if total_expected > 0 else 100 total_revenue_risk = sum(s["revenue_at_risk"] for s in shelf_compliance.values()) # Draw annotated frame annotated = models.draw_annotated(img, detections) # Log to database try: log_id = log_compliance( planogram_name=planogram_name, compliance=round(overall, 1), detected=len(detections), expected=total_expected, revenue_risk=total_revenue_risk, alert_count=len(all_alerts), scan_number=0, ) for alert in all_alerts: log_alert( compliance_log_id=log_id, alert_type=alert["type"], shelf_id=alert.get("shelf"), product_name=alert["product"], product_sku=alert.get("sku", ""), priority=alert["priority"], expected_count=alert.get("expected"), found_count=alert.get("found"), revenue=alert.get("revenue", 0), ) except Exception: pass return { "overall_compliance": round(overall, 1), "total_detected": len(detections), "total_expected": total_expected, "revenue_at_risk": total_revenue_risk, "shelf_compliance": shelf_compliance, "alerts": all_alerts, "annotated_image": image_to_base64(annotated), } # ── Analytics ───────────────────────────────────────────── @app.get("/api/analytics/summary") async def analytics_summary(): try: return get_analytics_summary() except Exception: return {} @app.get("/api/compliance/history") async def compliance_history(): try: return {"logs": get_compliance_logs_as_list()} except Exception: return {"logs": []} @app.get("/api/alerts") async def alerts(): try: return {"alerts": get_alerts_history()} except Exception: return {"alerts": []} # ── Push Notifications ──────────────────────────────────── @app.post("/api/notify") async def send_notification( title: str = Form("ShelfMind Alert"), message: str = Form(...), topic: str = Form("shelfmind-alerts"), priority: str = Form("high"), ): """Send push notification via ntfy.sh.""" import requests as req try: clean_title = title.encode("ascii", "ignore").decode("ascii").strip() or "ShelfMind Alert" req.post( f"https://ntfy.sh/{topic}", data=message.encode("utf-8"), headers={"Title": clean_title, "Priority": priority, "Tags": "warning"}, timeout=5, ) return {"status": "sent"} except Exception as e: return {"status": "failed", "error": str(e)} # ── WebSocket Live Monitoring ───────────────────────────── @app.websocket("/ws/monitor") async def live_monitor(websocket: WebSocket): """Real-time shelf monitoring via WebSocket.""" await websocket.accept() logger.info("WebSocket client connected") try: while True: data = await websocket.receive_text() msg = json.loads(data) if msg.get("type") == "frame": img_data = base64.b64decode(msg["data"]) img = Image.open(io.BytesIO(img_data)).convert("RGB") # Resize for speed max_dim = 640 if img.width > max_dim: ratio = max_dim / img.width img = img.resize((max_dim, int(img.height * ratio)), Image.LANCZOS) planogram_name = msg.get("planogram", "") conf = msg.get("confidence", 0.3) # Run compliance if planogram specified if planogram_name: planograms = get_planograms() if planogram_name in planograms: # Full compliance (simplified for WebSocket speed) detections = models.detect(img, conf=conf) boundaries, n_shelves = models.detect_shelf_levels(detections, img.height) shelf_assignments = models.assign_to_shelves(detections, boundaries) catalog_products = get_products() for det in detections: x1, y1, x2, y2 = det["bbox"] crop = img.crop((max(0, x1), max(0, y1), x2, y2)) emb = models.get_embedding(crop) match, score = models.search_product_faiss(emb, catalog_products, threshold=0.3) if match: det["product_name"] = match["name"] det["product_sku"] = match["sku"] det["match_score"] = round(score, 3) det["status"] = "match" else: det["product_name"] = "Unknown" det["product_sku"] = "UNKNOWN" det["match_score"] = round(score, 3) det["status"] = "unknown" annotated = models.draw_annotated(img, detections) await websocket.send_json({ "type": "detection", "total_products": len(detections), "detections": detections, "annotated_image": image_to_base64(annotated), }) else: # Just detection without compliance detections = models.detect(img, conf=conf) annotated = models.draw_annotated(img, [ {**d, "product_name": f"P{i+1}", "match_score": d["confidence"], "status": "match"} for i, d in enumerate(detections) ]) await websocket.send_json({ "type": "detection", "total_products": len(detections), "annotated_image": image_to_base64(annotated), }) elif msg.get("type") == "ping": await websocket.send_json({"type": "pong"}) except WebSocketDisconnect: logger.info("WebSocket client disconnected") except Exception as e: logger.error(f"WebSocket error: {e}") # ============================================================================= # MAIN # ============================================================================= if __name__ == "__main__": port = int(os.environ.get("PORT", 7860)) logger.info(f"Starting ShelfMind AI on http://0.0.0.0:{port}") uvicorn.run( "api_server:app", host="0.0.0.0", port=port, reload=False, log_level="info", )