Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from ultralyticsplus import YOLO
+from ultralytics import YOLO
 from PIL import Image, ImageDraw
 import io
 from fpdf import FPDF
@@ -7,9 +7,8 @@ import tempfile
 import cv2
 import os
 
-# Load the pre-trained model for shelf product and empty spot detection
-model = YOLO('foduucom/product-detection-in-shelf-yolov8')
-
+# Load the pre-trained YOLOv8n model
+model = YOLO('yolov8n.pt')
 model.overrides['conf'] = 0.25  # Confidence threshold
 model.overrides['iou'] = 0.45   # IoU threshold
 model.overrides['agnostic_nms'] = False
@@ -25,7 +24,6 @@ def process_input(input_file):
         cap.release()
         if not ret:
             raise ValueError("Could not read video frame.")
-        # Save temp frame
         temp_frame_path = tempfile.NamedTemporaryFile(suffix='.jpg', delete=False).name
         cv2.imwrite(temp_frame_path, frame)
         process_img_path = temp_frame_path
@@ -37,24 +35,25 @@ def process_input(input_file):
     boxes = results[0].boxes
     class_names = results[0].names
 
-    # Count present bottles (products) and empty spots
+    # Count bottles (assuming class 'bottle' from COCO dataset)
     num_present = 0
-    num_empty = 0
+    num_empty = 0  # Note: yolov8n.pt doesn't detect empty shelves
     empty_bboxes = []
     for box in boxes:
         cls_id = int(box.cls.item())
         cls_name = class_names[cls_id]
-        if cls_name == 'Magical Products':
+        if cls_name == 'bottle':
             num_present += 1
-        elif cls_name == 'Empty Shelves':
-            num_empty += 1
-            empty_bboxes.append(box.xyxy[0].cpu().numpy())  # [x1, y1, x2, y2]
+        # Empty spots not detected; you can add logic to infer empty spaces later
 
-    # Load image for annotation
+    # Load image for annotation (highlight bottles instead of empty spots for now)
     img = Image.open(process_img_path)
     draw = ImageDraw.Draw(img)
-    for bbox in empty_bboxes:
-        draw.rectangle(((bbox[0], bbox[1]), (bbox[2], bbox[3])), outline="red", width=3)
+    for box in boxes:
+        cls_id = int(box.cls.item())
+        if class_names[cls_id] == 'bottle':
+            bbox = box.xyxy[0].cpu().numpy()
+            draw.rectangle(((bbox[0], bbox[1]), (bbox[2], bbox[3])), outline="green", width=3)
 
     # Generate PDF
     pdf = FPDF()
@@ -62,31 +61,33 @@ def process_input(input_file):
     pdf.set_font("Arial", size=12)
     pdf.cell(200, 10, txt="Wine Shop Inventory Report", ln=1, align='C')
     pdf.ln(5)
-    pdf.cell(200, 10, txt=f"Number of bottles present (full/stock): {num_present}", ln=1)
-    pdf.cell(200, 10, txt=f"Number of empty spots (need to fill): {num_empty}", ln=1)
+    pdf.cell(200, 10, txt=f"Number of bottles detected: {num_present}", ln=1)
+    pdf.cell(200, 10, txt="Note: Empty spot detection not supported in this model.", ln=1)
     pdf.ln(10)
 
-    # Add annotated full image (scaled for UX)
+    # Add annotated full image
     with tempfile.NamedTemporaryFile(suffix='.png') as tmp_annotated:
         img.save(tmp_annotated.name)
-        pdf.image(tmp_annotated.name, x=10, y=pdf.get_y(), w=180)  # Full width for readability
+        pdf.image(tmp_annotated.name, x=10, y=pdf.get_y(), w=180)
     pdf.add_page()
-    pdf.cell(200, 10, txt="Screenshots of Empty Spots", ln=1, align='C')
+    pdf.cell(200, 10, txt="Detected Bottles", ln=1, align='C')
     pdf.ln(5)
     y_pos = pdf.get_y()
-    for i, bbox in enumerate(empty_bboxes):
-        cropped_img = img.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
-        with tempfile.NamedTemporaryFile(suffix='.png') as tmp_crop:
-            cropped_img.save(tmp_crop.name)
-            pdf.image(tmp_crop.name, x=10, y=y_pos, w=90)  # Medium size for clear view
-            y_pos += 100  # Space between images
-            pdf.cell(200, 10, txt=f"Empty Spot {i+1} (Location: x1={int(bbox[0])}, y1={int(bbox[1])}, x2={int(bbox[2])}, y2={int(bbox[3])})", ln=1)
-            pdf.ln(5)
-            if y_pos > 200:
-                pdf.add_page()
-                y_pos = 10
+    for i, box in enumerate(boxes):
+        if class_names[int(box.cls.item())] == 'bottle':
+            bbox = box.xyxy[0].cpu().numpy()
+            cropped_img = img.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
+            with tempfile.NamedTemporaryFile(suffix='.png') as tmp_crop:
+                cropped_img.save(tmp_crop.name)
+                pdf.image(tmp_crop.name, x=10, y=y_pos, w=90)
+                y_pos += 100
+                pdf.cell(200, 10, txt=f"Bottle {i+1} (Location: x1={int(bbox[0])}, y1={int(bbox[1])}, x2={int(bbox[2])}, y2={int(bbox[3])})", ln=1)
+                pdf.ln(5)
+                if y_pos > 200:
+                    pdf.add_page()
+                    y_pos = 10
 
-    # Output PDF as bytes
+    # Output PDF
     pdf_bytes = io.BytesIO()
     pdf.output(pdf_bytes)
     pdf_bytes.seek(0)