Update app.py

app.py

@@ -3,8 +3,6 @@ from ultralytics import YOLO
 from huggingface_hub import hf_hub_download
 from PIL import Image
 import io
-import cv2
-import numpy as np
 
 # Load YOLOv8 model from Hugging Face
 repo_id = "Norphel/nutri-ai-n2"  # Replace with your Hugging Face model repo
@@ -37,7 +35,7 @@ foods = {
     }
 }
 
-# Function to calculate bounding box area
+# Function to calculate bounding box area and class names
 def calculate_bounding_box_area(results):
     areas = []
     class_names = []
@@ -48,10 +46,10 @@ def calculate_bounding_box_area(results):
         area = (x2 - x1) * (y2 - y1)  # Calculate area
         areas.append(area)
         
-        # Get the class ID for the current bounding box
-        class_ids = box.cls.tolist()  # This returns a list of class IDs
+        # Get class IDs for the current bounding box (it can have multiple class IDs)
+        class_ids = box.cls.tolist()  # This may be a list if multiple classes detected
 
-        # Convert class ID(s) to class name(s)
+        # Add all detected class names for this box
         for class_id in class_ids:
             class_name = model.names[int(class_id)]  # Convert class ID to class name
             class_names.append(class_name)  # Add class name to list
@@ -77,10 +75,10 @@ def run_yolo(image):
 
     # Draw bounding boxes on the image and return the processed image
     result_img = results[0].plot()  # Draw bounding boxes on the image
-    result_pil = Image.fromarray(cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB))  # Convert back to PIL
-    print(detected_classes, bounding_areas)
+    result_pil = Image.fromarray(result_img)  # Convert back to PIL
     return result_pil, detected_classes, bounding_areas
 
+
 # Streamlit UI
 st.title("Nutri-AI")
 
@@ -94,7 +92,6 @@ if uploaded_file:
         st.write("🔍 Running YOLOv8 detection...")
 
         detected_image, detected_classes, bounding_areas = run_yolo(resized_image)
-        print(detected_classes, bounding_areas)
 
         # Display detected image
         st.image(detected_image, caption="Detected Objects", use_container_width=True)