Update app.py

app.py

@@ -4,14 +4,13 @@ from ultralytics import YOLO
 import numpy as np
 import os
 
-# ===== Load model safely =====
+# === Load model ===
 model_path = "best.pt"
 if not os.path.exists(model_path):
-    raise FileNotFoundError(f"❌ File '{model_path}' not found. Upload best.pt to the root directory.")
+    raise FileNotFoundError(f"❌ File '{model_path}' not found. Upload best.pt to root directory.")
 model = YOLO(model_path)
-print("✅ Model loaded successfully!")
 
-# ===== Label and Color Maps (sesuai urutan data.yaml) =====
+# Label dan warna sesuai urutan
 label_map = {
     0: "coral",
     1: "pipeline",
@@ -25,87 +24,77 @@ color_map = {
     3: (0, 0, 255)      # blue
 }
 
-def predict_segmentation_with_legend(image):
-    if model is None:
-        return Image.new("RGB", image.size, color=(255, 255, 255))
-
+def predict_segmentation(image: Image.Image):
     image = image.convert("RGB")
     results = model.predict(image, conf=0.25, iou=0.5)
     result = results[0]
 
     if result.masks is None or len(result.boxes) == 0:
-        # No detection
-        print("⚠️ No detection.")
-        image_np = np.array(image)
-        final_image = Image.fromarray(image_np)
-        draw = ImageDraw.Draw(final_image)
-        try:
-            font = ImageFont.truetype("arial.ttf", 24)
-        except:
-            font = ImageFont.load_default()
+        draw = ImageDraw.Draw(image)
+        font = ImageFont.load_default()
         draw.text((10, 10), "No detection", fill="red", font=font)
-        return final_image
+        return image
 
+    # Convert to NumPy for blending mask
     image_np = np.array(image).copy()
-    draw = ImageDraw.Draw(image)
-    try:
-        font = ImageFont.truetype("arial.ttf", 24)
-    except:
-        font = ImageFont.load_default()
 
     masks = result.masks.data.cpu().numpy()
-    boxes = result.boxes.xyxy.cpu().numpy()
+    boxes = result.boxes.xyxy.cpu().numpy().astype(int)
     scores = result.boxes.conf.cpu().numpy()
     class_ids = result.boxes.cls.cpu().numpy().astype(int)
 
+    # Draw masks
     for i, mask in enumerate(masks):
         class_id = class_ids[i]
-        label = label_map.get(class_id, f"class_{class_id}")
         color = color_map.get(class_id, (255, 255, 255))
-        score = scores[i]
-        box = boxes[i].astype(int)
-
-        # Safety checks
-        if len(box) != 4 or mask.shape != image_np.shape[:2]:
-            continue
-
-        # Blend mask
-        color_mask = np.zeros_like(image_np)
         for c in range(3):
-            color_mask[:, :, c] = mask * color[c]
-        image_np = np.where(color_mask > 0, image_np * 0.5 + color_mask * 0.5, image_np)
+            image_np[:, :, c] = np.where(
+                mask > 0.5,
+                image_np[:, :, c] * 0.5 + color[c] * 0.5,
+                image_np[:, :, c]
+            )
+
+    # Convert back to PIL for drawing box/label
+    image_masked = Image.fromarray(image_np.astype(np.uint8))
+    draw = ImageDraw.Draw(image_masked)
+    try:
+        font = ImageFont.truetype("arial.ttf", 20)
+    except:
+        font = ImageFont.load_default()
 
-        # Draw box and label
-        draw = ImageDraw.Draw(Image.fromarray(image_np))
+    for i, box in enumerate(boxes):
+        class_id = class_ids[i]
+        label = label_map.get(class_id, str(class_id))
+        color = color_map.get(class_id, (255, 255, 255))
+        score = scores[i]
         draw.rectangle(box.tolist(), outline=color, width=3)
-        draw.text((box[0], box[1] - 10), f"{label}: {score:.2f}", fill="white", font=font)
-
-    # Convert final image
-    final_image = Image.fromarray(image_np.astype(np.uint8))
+        text = f"{label}: {score:.2f}"
+        draw.text((box[0], box[1] - 10), text, fill="white", font=font)
 
-    # Create legend
-    legend = Image.new("RGB", (400, 50), (255, 255, 255))
+    # Add legend at bottom
+    legend_height = 50
+    legend = Image.new("RGB", (image_masked.width, legend_height), (255, 255, 255))
     draw_legend = ImageDraw.Draw(legend)
     x = 10
     for cid, label in label_map.items():
-        draw_legend.rectangle([x, 10, x+20, 30], fill=color_map[cid])
-        draw_legend.text((x + 25, 10), label, fill="black", font=font)
-        x += 120
+        draw_legend.rectangle([x, 15, x + 20, 35], fill=color_map[cid])
+        draw_legend.text((x + 25, 15), label, fill="black", font=font)
+        x += 130
 
-    # Combine image and legend
-    combined = Image.new("RGB", (final_image.width, final_image.height + 50))
-    combined.paste(final_image, (0, 0))
-    combined.paste(legend, (10, final_image.height))
+    # Combine image + legend
+    combined = Image.new("RGB", (image_masked.width, image_masked.height + legend_height))
+    combined.paste(image_masked, (0, 0))
+    combined.paste(legend, (0, image_masked.height))
 
     return combined
 
-# === Gradio Interface ===
+# === Gradio interface ===
 iface = gr.Interface(
-    fn=predict_segmentation_with_legend,
+    fn=predict_segmentation,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
     title="YOLOv8 Segmentasi Sonar",
-    description="Upload citra Side Scan Sonar. Hasil segmentasi ditampilkan dengan mask warna, bounding box, confidence score, dan legenda."
+    description="Upload citra sonar. Hasil segmentasi akan menampilkan mask, bounding box, label, dan confidence score."
 )
 
 if __name__ == "__main__":