Update app.py

app.py

@@ -7,6 +7,15 @@ from ultralytics import YOLO
 # Initialize YOLO model
 model = YOLO('./best.pt')
 
+# Class to color mapping
+ClassToColorMapping = {
+    0: "black",
+    1: "gray",
+    2: "green",
+    3: "purple",
+    4: "red"
+}
+
 def check_ffmpeg():
     """Check if ffmpeg is installed."""
     if shutil.which("ffmpeg") is None:
@@ -46,25 +55,28 @@ def process_video(video_path, output_option):
     if not matched_mp4_files:
         raise Exception(f"No .mp4 video files found in directory {output_dir} after conversion for {video_name}.")
 
-    # Initialize object count
-    object_count = 0
+    # Initialize object count dictionary based on colors
+    color_counts = {color: 0 for color in ClassToColorMapping.values()}
 
-    # Calculate object count based on results
+    # Calculate object count based on results and map to colors
     for result in results:
         if hasattr(result, 'boxes'):
-            object_count += len(result.boxes)
+            for box in result.boxes:
+                class_id = int(box.cls)  # Assuming `cls` is the class index
+                color = ClassToColorMapping.get(class_id, "unknown")
+                if color in color_counts:
+                    color_counts[color] += 1
 
     # Set default outputs
     video_output = matched_mp4_files[0] if matched_mp4_files else None
-    count_output = object_count
 
     # Determine the output based on user's choice
     if output_option == "Count":
-        return count_output, None
+        return color_counts, None
     elif output_option == "Video":
         return None, video_output
     elif output_option == "Both":
-        return count_output, video_output
+        return color_counts, video_output
 
 # Define Gradio inputs
 video_input = gr.Video()
@@ -78,14 +90,14 @@ examples = [
 video_interface = gr.Interface(
     fn=process_video,
     inputs=[video_input, output_option_input],
-    outputs=[gr.Textbox(label="Object Count"), "video"],  # Ensure two outputs are defined
-    title="YOLO Video Tracking Application",
-    description="A simple application to track objects in a video using YOLO model. Upload your own video, or click one of the examples to load them.",
+    outputs=[gr.JSON(label="Color Counts"), "video"],  # Ensure two outputs are defined, using JSON for color counts
+    title="YOLO Video Tracking Application with Color Counting",
+    description="A simple application to track objects in a video using YOLO model and count objects by color. Upload your own video, or click one of the examples to load them.",
     article="""<div>
-                <p style="text-align: center">Upload a video file and select the type of output you want: object count, processed video, or both. Then, hit submit to process the video.</p>
+                <p style="text-align: center">Upload a video file and select the type of output you want: object color counts, processed video, or both. Then, hit submit to process the video.</p>
                </div>""",
     examples=examples,
-    cache_examples=True  # Disable caching to speed up launch
+    cache_examples=False  # Disable caching to speed up launch
 )
 
 # Deploy the interface with share enabled