support bins count

app.py

@@ -4,13 +4,13 @@ import requests
 import os
 import random
 from ultralytics import YOLO
+import numpy as np
+from collections import defaultdict
 
-# 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}
+# Import the supervision library
+import supervision as sv
 
+# --- File Downloading ---
 # File URLs for sample images and video
 file_urls = [
     'https://huggingface.co/spaces/iamsuman/waste-detection/resolve/main/samples/mix2.jpg?download=true',
@@ -18,119 +18,232 @@ file_urls = [
     'https://huggingface.co/spaces/iamsuman/waste-detection/resolve/main/samples/sample_waste.mp4?download=true',
 ]
 
-# Function to download files (always overwrites existing ones)
 def download_file(url, save_name):
-    print(f"Downloading from: {url}")  # Log the URL
+    """Downloads a file from a URL, overwriting if it exists."""
+    print(f"Downloading from: {url}")
     try:
         response = requests.get(url, stream=True)
         response.raise_for_status()  # Check for HTTP errors
         with open(save_name, 'wb') as file:
             for chunk in response.iter_content(1024):
                 file.write(chunk)
-        print(f"Downloaded and overwritten: {save_name}")
+        print(f"Downloaded and overwrote: {save_name}")
     except requests.exceptions.RequestException as e:
         print(f"Error downloading {url}: {e}")
 
-# Download images and video
+# Download sample images and video for the examples
 for i, url in enumerate(file_urls):
-    print(i, url)
-    if 'mp4' in file_urls[i]:
-        download_file(file_urls[i], f"video.mp4")
+    if 'mp4' in url:
+        download_file(url, "video.mp4")
     else:
-        download_file(file_urls[i], f"image_{i}.jpg")
+        download_file(url, f"image_{i}.jpg")
 
-# Load YOLO model
+# --- Model and Class Configuration ---
+# Load your custom YOLO model
+# IMPORTANT: Replace 'best.pt' with the path to your model trained on the 12 classes.
 model = YOLO('best.pt')
 
-# Sample paths
-path = [['image_0.jpg'], ['image_1.jpg']]
-video_path = [['video.mp4']]
+# Get class names and generate colors dynamically from the loaded model
+# This is the best practice as it ensures names and colors match the model's output.
+class_names = model.model.names
+class_colors = {
+    name: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+    for name in class_names.values()
+}
+
+# Define paths for Gradio examples
+image_example_paths = [['image_0.jpg'], ['image_1.jpg']]
+video_example_path = [['video.mp4']]
+
 
-# Function to process and display predictions on images
+# --- Image Processing Function ---
 def show_preds_image(image_path):
+    """Processes a single image and overlays YOLO predictions."""
     image = cv2.imread(image_path)
-    outputs = model.predict(source=image_path)
+    outputs = model.predict(source=image_path, verbose=False)
     results = outputs[0].cpu().numpy()
 
-    boxes = results.boxes
-    names = model.model.names
+    # Convert to supervision Detections object for easier handling
+    detections = sv.Detections.from_ultralytics(outputs[0])
 
-    for box, conf, cls in zip(boxes.xyxy, boxes.conf, boxes.cls):
+    # Annotate the image with bounding boxes and labels
+    for i, (box, conf, cls) in enumerate(zip(detections.xyxy, detections.confidence, detections.class_id)):
         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
-
+        class_name = class_names[cls]
+        color = class_colors[class_name]
+        
         # 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, 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
-
-        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
+        # Create and display label
+        label = f"{class_name}: {conf:.2f}"
+        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA)
 
-            # Draw bounding box
-            cv2.rectangle(frame_copy, (x1, y1), (x2, y2), color=color, thickness=2, lineType=cv2.LINE_AA)
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
-            # 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)
+# --- Video Processing Function (with Supervision) ---
+def process_video_with_two_side_bins(video_path):
 
-    cap.release()
+    if video_path is None:
+        return
+    
+    generator = sv.get_video_frames_generator(video_path)
 
-# Gradio Image Interface
-inputs_image = [gr.Image(type="filepath", label="Input Image")]
-outputs_image = [gr.Image(type="numpy", label="Output Image")]
+    try:
+        first_frame = next(generator)
+    except StopIteration:
+        print("No frames found in the provided video input.")
+        # Option 1: Return or yield a blank frame or error image
+        # For example, yield a blank black image of fixed size:
+        blank_frame = np.zeros((480, 640, 3), dtype=np.uint8)
+        yield cv2.cvtColor(blank_frame, cv2.COLOR_BGR2RGB)
+        return
+    
+    first_frame = next(generator)
+    frame_height, frame_width, _ = first_frame.shape
+
+    # Define two bins: recyle and trash sides
+
+    bins = [
+        {
+            "name": "Recycle Bin",
+            "coords": (
+                int(frame_width * 0.05),
+                int(frame_height * 0.5),
+                int(frame_width * 0.25),
+                int(frame_height * 0.95),
+            ),
+            "color": (200, 16, 46),  # Blue-ish
+        },
+        {
+            "name": "Trash Bin",
+            "coords": (
+                int(frame_width * 0.75),
+                int(frame_height * 0.5),
+                int(frame_width * 0.95),
+                int(frame_height * 0.95),
+            ),
+            "color": (50, 50, 50),  # Red-ish
+        },
+    ]
+
+    box_annotator = sv.BoxAnnotator(thickness=2)
+    label_annotator = sv.LabelAnnotator(
+        text_scale=1.2,  # bigger text size
+        text_thickness=3,
+        text_position=sv.Position.TOP_LEFT,
+    )
+
+    tracker = sv.ByteTrack()
+
+    items_in_bins = {bin_["name"]: set() for bin_ in bins}
+    class_counts_per_bin = {bin_["name"]: defaultdict(int) for bin_ in bins}
+
+    for i, frame in enumerate(generator):
+        results = model(frame, verbose=False)[0]
+        detections = sv.Detections.from_ultralytics(results)
+        tracked_detections = tracker.update_with_detections(detections)
+
+        annotated_frame = frame.copy()
+
+        # Draw bins and bigger labels
+        for bin_ in bins:
+            x1, y1, x2, y2 = bin_["coords"]
+            color = bin_["color"]
+            cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color=color, thickness=3)
+            cv2.putText(
+                annotated_frame,
+                bin_["name"],
+                (x1 + 5, y1 - 15),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1.5,  # bigger font
+                color,
+                3,
+                cv2.LINE_AA,
+            )
+
+        if tracked_detections.tracker_id is None:
+            yield cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
+            continue
+
+        for box, track_id, class_id in zip(
+            tracked_detections.xyxy,
+            tracked_detections.tracker_id,
+            tracked_detections.class_id,
+        ):
+            x1, y1, x2, y2 = map(int, box)
+            cx = (x1 + x2) // 2
+            cy = (y1 + y2) // 2
+
+            for bin_ in bins:
+                bx1, by1, bx2, by2 = bin_["coords"]
+                if (bx1 <= cx <= bx2) and (by1 <= cy <= by2):
+                    if track_id not in items_in_bins[bin_["name"]]:
+                        items_in_bins[bin_["name"]].add(track_id)
+                        class_name = class_names[class_id]
+                        class_counts_per_bin[bin_["name"]][class_name] += 1
+
+        labels = [
+            f"#{tid} {class_names[cid]}"
+            for cid, tid in zip(tracked_detections.class_id, tracked_detections.tracker_id)
+        ]
+
+        annotated_frame = box_annotator.annotate(
+            scene=annotated_frame, detections=tracked_detections
+        )
+        annotated_frame = label_annotator.annotate(
+            scene=annotated_frame, detections=tracked_detections, labels=labels
+        )
+
+        # Show counts per bin with bigger font
+        y_pos = 50
+        for bin_name, class_count_dict in class_counts_per_bin.items():
+            text = (
+                f"{bin_name}: "
+                + ", ".join(f"{cls}={count}" for cls, count in class_count_dict.items())
+            )
+            cv2.putText(
+                annotated_frame,
+                text,
+                (30, y_pos),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1.1,  # bigger font for counts
+                (255, 255, 255),
+                3,
+                cv2.LINE_AA,
+            )
+            y_pos += 40
+
+        yield cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
+
+
+
+# --- Gradio Interface Setup ---
+# Gradio Interface for Image Processing
 interface_image = gr.Interface(
     fn=show_preds_image,
-    inputs=inputs_image,
-    outputs=outputs_image,
-    title="Waste Detection",
-    examples=path,
+    inputs=gr.Image(type="filepath", label="Input Image"),
+    outputs=gr.Image(type="numpy", label="Output Image"),
+    title="Waste Detection (Image)",
+    description="Upload an image to see waste detection results.",
+    examples=image_example_paths,
     cache_examples=False,
 )
 
-# Gradio Video Interface
-inputs_video = [gr.Video(label="Input Video")]
-outputs_video = [gr.Image(type="numpy", label="Output Image")]
+# Gradio Interface for Video Processing
 interface_video = gr.Interface(
-    fn=show_preds_video,
-    inputs=inputs_video,
-    outputs=outputs_video,
-    title="Waste Detection",
-    examples=video_path,
+    fn=process_video_with_two_side_bins,
+    inputs=gr.Video(label="Input Video"),
+    outputs=gr.Image(type="numpy", label="Output Video Stream"),
+    title="Waste Tracking and Counting (Video)",
+    description="Upload a video to see real-time object tracking and counting.",
+    examples=video_example_path,
     cache_examples=False,
 )
 
-# Launch Gradio App
+# Launch the Gradio App with separate tabs for each interface
 gr.TabbedInterface(
     [interface_image, interface_video],
     tab_names=['Image Inference', 'Video Inference']
-).queue().launch()
+).queue().launch(debug=True)

best.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6ba9e2e924f6acd401c8ef4c4d2b10fe80049dcbcf46597be2196e03986271d6
-size 6211875
+oid sha256:ea3a57f09eaca3340d951c0f5fcf90604f9544cf0a2643f949a182616e936af1
+size 6215843

requirements.txt

@@ -1,47 +1,7 @@
-# Ultralytics requirements
-# Usage: pip install -r requirements.txt
-
-# Base ----------------------------------------
-hydra-core>=1.2.0
-matplotlib>=3.2.2
-numpy>=1.18.5
-opencv-python>=4.1.1
-Pillow>=7.1.2
-PyYAML>=5.3.1
-requests>=2.23.0
-scipy>=1.4.1
-torch>=1.7.0
-torchvision>=0.8.1
-tqdm>=4.64.0
+gradio==4.10.0
 ultralytics
-
-# Logging -------------------------------------
-tensorboard>=2.4.1
-# clearml
-# comet
-
-# Plotting ------------------------------------
-pandas>=1.1.4
-seaborn>=0.11.0
-
-# Export --------------------------------------
-# coremltools>=6.0  # CoreML export
-# onnx>=1.12.0  # ONNX export
-# onnx-simplifier>=0.4.1  # ONNX simplifier
-# nvidia-pyindex  # TensorRT export
-# nvidia-tensorrt  # TensorRT export
-# scikit-learn==0.19.2  # CoreML quantization
-# tensorflow>=2.4.1  # TF exports (-cpu, -aarch64, -macos)
-# tensorflowjs>=3.9.0  # TF.js export
-# openvino-dev  # OpenVINO export
-
-# Extras --------------------------------------
-ipython  # interactive notebook
-psutil  # system utilization
-thop>=0.1.1  # FLOPs computation
-# albumentations>=1.0.3
-# pycocotools>=2.0.6  # COCO mAP
-# roboflow
-
-# HUB -----------------------------------------
-GitPython>=3.1.24
+Pillow
+pydantic==2.8.2
+pydantic-core==2.20.1
+fastapi==0.112.4
+supervision>=0.26.1  # YOLOv8 inference