Update app.py

app.py

@@ -3,6 +3,8 @@ 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
@@ -35,7 +37,7 @@ foods = {
     }
 }
 
-# Function to calculate bounding box area and class names
+# Function to calculate bounding box area
 def calculate_bounding_box_area(results):
     areas = []
     class_names = []
@@ -46,10 +48,10 @@ def calculate_bounding_box_area(results):
         area = (x2 - x1) * (y2 - y1)  # Calculate area
         areas.append(area)
         
-        # 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
+        # Get the class ID for the current bounding box
+        class_ids = box.cls.tolist()  # This returns a list of class IDs
 
-        # Add all detected class names for this box
+        # Convert class ID(s) to class name(s)
         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
@@ -75,10 +77,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(result_img)  # Convert back to PIL
+    result_pil = Image.fromarray(cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB))  # Convert back to PIL
+    print(detected_classes, bounding_areas)
     return result_pil, detected_classes, bounding_areas
 
-
 # Streamlit UI
 st.title("Nutri-AI")
 
@@ -92,6 +94,7 @@ 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)