fix demo graph

app.py

@@ -30,20 +30,20 @@ def compute_class_areas(predictions, image_shape):
     """
     Compute the area percentage covered by each class for the prediction.
     """
-    total_pixels = image_shape[0] * image_shape[1]
+    total_pixels = image_shape[0] * image_shape[1]  # Total pixels in the image
     class_areas = {}
 
     for mask, cls_id in zip(predictions.masks.data, predictions.boxes.cls):
-        cls_id = int(cls_id.item())
-        mask_area = mask.sum().item()
+        class_name = predictions.names[int(cls_id.item())]  # Get the class name
+        mask_area = mask.sum().item()  # Count non-zero pixels in the mask
 
-        if cls_id not in class_areas:
-            class_areas[cls_id] = 0
-        class_areas[cls_id] += mask_area
+        if class_name not in class_areas:
+            class_areas[class_name] = 0
+        class_areas[class_name] += mask_area
 
-    # Convert to percentages
-    for cls_id in class_areas:
-        class_areas[cls_id] = (class_areas[cls_id] / total_pixels) * 100
+    # Convert areas to percentages relative to the total image
+    for class_name in class_areas:
+        class_areas[class_name] = (class_areas[class_name] / total_pixels) * 100
 
     return class_areas
 
@@ -56,9 +56,9 @@ def generate_coverage_graph(class_areas):
     classes = list(class_areas.keys())
     coverage = list(class_areas.values())
     plt.bar(classes, coverage, color="skyblue")
-    plt.xlabel("Class ID")
+    plt.xlabel("Class")
     plt.ylabel("Coverage Percentage")
-    plt.title("Class Coverage")
+    plt.title("Class Coverage as % of Total Image")
     graph_path = "class_coverage.png"
     plt.savefig(graph_path)
     plt.close()