fix gradio

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
 import os
+import matplotlib.pyplot as plt
 from ultralytics import YOLO
 
 # Coral AI model files hosted in the Hugging Face model repository
@@ -28,14 +29,39 @@ for model_name in model_names:
 model = YOLO(os.path.join(model_dir, current_model_name))
 
 
+def generate_coverage_graph(results):
+    """
+    Generates a coverage percentage graph for the output classes.
+    """
+    # Extract segmentation masks and class labels
+    masks = results[0].masks.data  # Segmentation masks
+    class_ids = results[0].masks.cls  # Class IDs
+
+    # Calculate coverage percentages for each class
+    total_pixels = masks.sum()
+    class_percentages = {int(cls): (mask.sum() / total_pixels) * 100 for cls, mask in zip(class_ids, masks)}
+
+    # Generate the graph
+    plt.figure(figsize=(8, 6))
+    plt.bar(class_percentages.keys(), class_percentages.values())
+    plt.xlabel("Class ID")
+    plt.ylabel("Coverage Percentage")
+    plt.title("Class Coverage in Segmentation")
+    plt.xticks(list(class_percentages.keys()))
+    graph_path = "class_coverage.png"
+    plt.savefig(graph_path)
+    plt.close()
+    return graph_path
+
+
 def coral_ai_inference(image: str, model_name: str):
     """
-    Coral AI inference function
+    Coral AI inference function with class coverage graph.
     Args:
         image: Input image filepath
         model_name: Name of the model
     Returns:
-        Rendered image
+        Rendered image and class coverage graph
     """
     global model
     global current_model_name
@@ -46,9 +72,12 @@ def coral_ai_inference(image: str, model_name: str):
     # Perform inference
     results = model.predict(image, return_outputs=True)
 
-    # Render results (placeholder logic)
-    rendered_image = results[0].plot()  # Visualization of predictions
-    return rendered_image
+    # Rendered image
+    rendered_image = results[0].plot()
+
+    # Generate class coverage graph
+    graph_path = generate_coverage_graph(results)
+    return rendered_image, graph_path
 
 
 # Define Gradio inputs and outputs
@@ -61,7 +90,11 @@ inputs = [
     ),
 ]
 
-outputs = gr.Image(type="filepath", label="Output Image")
+outputs = [
+    gr.Image(type="filepath", label="Segmented Image"),
+    gr.Image(type="filepath", label="Class Coverage Graph"),
+]
+
 title = "Coral AI YOLOv8 Segmentation Demo"
 
 examples = [

requirements.txt

@@ -1,3 +1,4 @@
 gradio
 torch
-requests
+ultralytics
+matplotlib