fix model

app.py

@@ -6,12 +6,9 @@ from ultralytics import YOLO
 
 # Coral AI model files hosted in the Hugging Face model repository
 model_names = [
-    "yolov8_xlarge_latest.pt",
-    # "yolov8_xlarge_v1.pt",
-    # "yolov8_xlarge_v2.pt",
+    "yolov8_xlarge_latest.pt"
 ]
 
-# Set the initial model
 current_model_name = "yolov8_xlarge_latest.pt"
 model_dir = "models"
 os.makedirs(model_dir, exist_ok=True)
@@ -29,39 +26,48 @@ for model_name in model_names:
 model = YOLO(os.path.join(model_dir, current_model_name))
 
 
-def generate_coverage_graph(results):
+def compute_class_areas(predictions, image_shape):
     """
-    Generates a coverage percentage graph for the output classes.
+    Compute the area percentage covered by each class for the prediction.
     """
-    # Extract segmentation masks and class labels
-    masks = results[0].masks.data  # Segmentation masks
-    class_ids = results[0].masks.cls  # Class IDs
+    total_pixels = image_shape[0] * image_shape[1]
+    class_areas = {}
+    merged_masks = {}
 
-    # 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)}
+    for mask, cls in zip(predictions.masks.data, predictions.masks.cls):
+        cls_id = int(cls.item())
+        if cls_id not in merged_masks:
+            merged_masks[cls_id] = mask
+        else:
+            merged_masks[cls_id] = torch.logical_or(merged_masks[cls_id], mask)
 
-    # Generate the graph
+    for cls_id, mask in merged_masks.items():
+        mask_area = mask.sum().item()
+        class_areas[cls_id] = (mask_area / total_pixels) * 100
+
+    return class_areas
+
+
+def generate_coverage_graph(class_areas):
+    """
+    Generates a graph for class coverage percentages.
+    """
     plt.figure(figsize=(8, 6))
-    plt.bar(class_percentages.keys(), class_percentages.values())
+    classes = list(class_areas.keys())
+    coverage = list(class_areas.values())
+    plt.bar(classes, coverage, color="skyblue")
     plt.xlabel("Class ID")
     plt.ylabel("Coverage Percentage")
-    plt.title("Class Coverage in Segmentation")
-    plt.xticks(list(class_percentages.keys()))
+    plt.title("Class Coverage")
     graph_path = "class_coverage.png"
     plt.savefig(graph_path)
     plt.close()
     return graph_path
 
 
-def coral_ai_inference(image: str, model_name: str):
+def coral_ai_inference(image, model_name):
     """
-    Coral AI inference function with class coverage graph.
-    Args:
-        image: Input image filepath
-        model_name: Name of the model
-    Returns:
-        Rendered image and class coverage graph
+    Perform inference and generate class coverage data.
     """
     global model
     global current_model_name
@@ -69,29 +75,26 @@ def coral_ai_inference(image: str, model_name: str):
         model = YOLO(os.path.join(model_dir, model_name))
         current_model_name = model_name
 
-    # Perform inference
-    results = model.predict(image, return_outputs=True)
+    results = model(image)  # Perform inference
+
+    # Generate class coverage
+    class_areas = compute_class_areas(results[0], image.shape[:2])
+    graph_path = generate_coverage_graph(class_areas)
 
-    # Rendered image
+    # Render the prediction
     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
+# Define Gradio interface
 inputs = [
-    gr.Image(type="filepath", label="Input Image"),
-    gr.Dropdown(
-        model_names,
-        value=current_model_name,
-        label="Model Type",
-    ),
+    gr.Image(type="numpy", label="Input Image"),
+    gr.Dropdown(model_names, value=current_model_name, label="Model Type"),
 ]
 
 outputs = [
-    gr.Image(type="filepath", label="Segmented Image"),
+    gr.Image(type="numpy", label="Segmented Image"),
     gr.Image(type="filepath", label="Class Coverage Graph"),
 ]
 
@@ -103,7 +106,6 @@ examples = [
     # ["examples/coral_image3.jpg", "yolov8_xlarge_latest.pt"],
 ]
 
-# Create and launch the Gradio interface
 demo_app = gr.Interface(
     fn=coral_ai_inference,
     inputs=inputs,
@@ -111,6 +113,5 @@ demo_app = gr.Interface(
     title=title,
     examples=examples,
     cache_examples=True,
-    theme="default",
 )
 demo_app.queue().launch(debug=True)