updated app file

app.py

@@ -2,167 +2,127 @@ import gradio as gr
 import cv2
 import requests
 import os
-
+import random
 from ultralytics import YOLO
 
+# Define class names based on YOLO labels
+class_names = {0: 'AluCan', 1: 'Glass', 2: 'PET', 3: 'HDPEM'}
+
+# Generate random colors for each class
+class_colors = {cls: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for cls in class_names}
+
+# File URLs for sample images and video
 file_urls = [
-    'https://huggingface.co/spaces/iamsuman/ripe-and-unripe-tomatoes-detection/resolve/main/samples/riped_tomato_93.jpeg?download=true',
-    'https://huggingface.co/spaces/iamsuman/ripe-and-unripe-tomatoes-detection/resolve/main/samples/unriped_tomato_18.jpeg?download=true',
-    'https://huggingface.co/spaces/iamsuman/ripe-and-unripe-tomatoes-detection/resolve/main/samples/tomatoes.mp4?download=true',
+    'https://huggingface.co/spaces/iamsuman/ripe-and-unripe-tomatoes-detection/resolve/main/samples/AluCan1,000.jpg?download=true',
+    'https://huggingface.co/spaces/iamsuman/ripe-and-unripe-tomatoes-detection/resolve/main/samples/Glass847.jpg?download=true',
+    'https://huggingface.co/spaces/iamsuman/ripe-and-unripe-tomatoes-detection/resolve/main/samples/sample_waste.mp4?download=true',
 ]
 
+# Function to download files
 def download_file(url, save_name):
-    url = url
     if not os.path.exists(save_name):
         file = requests.get(url)
         open(save_name, 'wb').write(file.content)
 
+# Download images and video
 for i, url in enumerate(file_urls):
     if 'mp4' in file_urls[i]:
-        download_file(
-            file_urls[i],
-            f"video.mp4"
-        )
+        download_file(file_urls[i], "video.mp4")
     else:
-        download_file(
-            file_urls[i],
-            f"image_{i}.jpg"
-        )
+        download_file(file_urls[i], f"image_{i}.jpg")
 
+# Load YOLO model
 model = YOLO('best.pt')
-path  = [['image_0.jpg'], ['image_1.jpg']]
-video_path = [['video.mp4']]
-
-
 
+# Sample paths
+path = [['image_0.jpg'], ['image_1.jpg']]
+video_path = [['video.mp4']]
 
+# Function to process and display predictions on images
 def show_preds_image(image_path):
     image = cv2.imread(image_path)
     outputs = model.predict(source=image_path)
     results = outputs[0].cpu().numpy()
 
-    # Print the detected objects' information (class, coordinates, and probability)
-    box = results[0].boxes
-    names = model.model.names
     boxes = results.boxes
+    names = model.model.names
 
     for box, conf, cls in zip(boxes.xyxy, boxes.conf, boxes.cls):
-
         x1, y1, x2, y2 = map(int, box)
 
         class_name = names[int(cls)]
-        print(class_name, "class_name", class_name.lower() == 'ripe')
-        if class_name.lower() == 'ripe':
-            color = (0, 0, 255)  # Red for ripe
-        else:
-            color = (0, 255, 0)  # Green for unripe
-
-        # Draw rectangle around object
-        cv2.rectangle(
-            image,
-            (x1, y1),
-            (x2, y2),
-            color=color,
-            thickness=2,
-            lineType=cv2.LINE_AA
-        )
-
-        # Display class label on top of rectangle
+        color = class_colors.get(int(cls), (255, 255, 255))  # Default to white if class is unknown
+
+        # Draw bounding box
+        cv2.rectangle(image, (x1, y1), (x2, y2), color=color, thickness=2, lineType=cv2.LINE_AA)
+
+        # Display class label
         label = f"{class_name.capitalize()}: {conf:.2f}"
-        cv2.putText(image, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color,  # Use the same color as the rectangle
-            2,
-            cv2.LINE_AA)
-        
-    # Convert image to RGB (Gradio expects RGB format)
+        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2, cv2.LINE_AA)
+
     return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+# Function to process and display predictions on video
+def show_preds_video(video_path):
+    cap = cv2.VideoCapture(video_path)
     
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        frame_copy = frame.copy()
+        outputs = model.predict(source=frame)
+        results = outputs[0].cpu().numpy()
+        
+        boxes = results.boxes
+        confidences = boxes.conf
+        classes = boxes.cls
+        names = model.model.names
 
-inputs_image = [
-    gr.components.Image(type="filepath", label="Input Image"),
-]
-outputs_image = [
-    gr.components.Image(type="numpy", label="Output Image"),
-]
+        for box, conf, cls in zip(boxes.xyxy, confidences, classes):
+            x1, y1, x2, y2 = map(int, box)
+
+            class_name = names[int(cls)]
+            color = class_colors.get(int(cls), (255, 255, 255))  # Default to white if class is unknown
+
+            # Draw bounding box
+            cv2.rectangle(frame_copy, (x1, y1), (x2, y2), color=color, thickness=2, lineType=cv2.LINE_AA)
+
+            # Display class label
+            label = f"{class_name.capitalize()}: {conf:.2f}"
+            cv2.putText(frame_copy, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)
+
+        yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
+
+    cap.release()
+
+# Gradio Image Interface
+inputs_image = [gr.Image(type="filepath", label="Input Image")]
+outputs_image = [gr.Image(type="numpy", label="Output Image")]
 interface_image = gr.Interface(
     fn=show_preds_image,
     inputs=inputs_image,
     outputs=outputs_image,
-    title="Ripe And Unripe Tomatoes Detection",
+    title="Waste Detection",
     examples=path,
     cache_examples=False,
 )
 
-def show_preds_video(video_path):
-    cap = cv2.VideoCapture(video_path)
-    while(cap.isOpened()):
-        ret, frame = cap.read()
-        if ret:
-            frame_copy = frame.copy()
-            outputs = model.predict(source=frame)
-            results = outputs[0].cpu().numpy()
-            
-            boxes = results.boxes
-            confidences = boxes.conf
-            classes = boxes.cls
-            names = model.model.names
-
-            for box, conf, cls in zip(boxes.xyxy, confidences, classes):
-                x1, y1, x2, y2 = map(int, box)
-
-                # Determine color based on class
-                class_name = names[int(cls)]
-                if class_name.lower() == 'ripe':
-                    color = (0, 0, 255)  # Red for ripe
-                else:
-                    color = (0, 255, 0)  # Green for unripe
-
-                # Draw rectangle around object
-                cv2.rectangle(
-                    frame_copy,
-                    (x1, y1),
-                    (x2, y2),
-                    color=color,
-                    thickness=2,
-                    lineType=cv2.LINE_AA
-                )
-
-                # Display class label on top of rectangle with capitalized class name
-                label = f"{class_name.capitalize()}: {conf:.2f}"
-                cv2.putText(
-                    frame_copy,
-                    label,
-                    (x1, y1 - 10),  # Position slightly above the top of the rectangle
-                    cv2.FONT_HERSHEY_SIMPLEX,
-                    0.5,
-                    color,  # Use the same color as the rectangle
-                    1,
-                    cv2.LINE_AA
-                )
-
-            # Convert frame to RGB (Gradio expects RGB format)
-            yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
-        else:
-            break
-
-    cap.release()
-
-inputs_video = [
-    gr.components.Video(label="Input Video"),
-
-]
-outputs_video = [
-    gr.components.Image(type="numpy", label="Output Image"),
-]
+# Gradio Video Interface
+inputs_video = [gr.Video(label="Input Video")]
+outputs_video = [gr.Image(type="numpy", label="Output Image")]
 interface_video = gr.Interface(
     fn=show_preds_video,
     inputs=inputs_video,
     outputs=outputs_video,
-    title="Ripe And Unripe Tomatoes Detection",
+    title="Waste Detection",
     examples=video_path,
     cache_examples=False,
 )
 
+# Launch Gradio App
 gr.TabbedInterface(
     [interface_image, interface_video],
-    tab_names=['Image inference', 'Video inference']
-).queue().launch()
+    tab_names=['Image Inference', 'Video Inference']
+).queue().launch()