Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -1,35 +1,3 @@
|
|
1 |
-
import gradio as gr
|
2 |
-
from PIL import Image, ImageDraw, ImageFont
|
3 |
-
from ultralytics import YOLO
|
4 |
-
import numpy as np
|
5 |
-
import os
|
6 |
-
|
7 |
-
# Load model
|
8 |
-
try:
|
9 |
-
model_path = "best.pt"
|
10 |
-
if not os.path.exists(model_path):
|
11 |
-
raise FileNotFoundError(f"❌ File {model_path} tidak ditemukan. Upload 'best.pt' ke root.")
|
12 |
-
model = YOLO(model_path)
|
13 |
-
print("✅ Model loaded successfully!")
|
14 |
-
except Exception as e:
|
15 |
-
print("❌ Gagal load model:", e)
|
16 |
-
model = None
|
17 |
-
|
18 |
-
# Label dan Warna untuk tiap class
|
19 |
-
label_map = {
|
20 |
-
0: "coral or rock",
|
21 |
-
1: "pipeline",
|
22 |
-
2: "ripple marks",
|
23 |
-
3: "shipwreck"
|
24 |
-
}
|
25 |
-
|
26 |
-
color_map = {
|
27 |
-
0: (0, 255, 0), # coral or rock - hijau
|
28 |
-
1: (255, 0, 0), # pipeline - merah
|
29 |
-
2: (255, 165, 0), # ripple marks - oranye
|
30 |
-
3: (0, 0, 255) # shipwreck - biru
|
31 |
-
}
|
32 |
-
|
33 |
def predict_segmentation_with_legend(image):
|
34 |
try:
|
35 |
if model is None:
|
@@ -83,8 +51,11 @@ def predict_segmentation_with_legend(image):
|
|
83 |
image_np = np.where(mask_resized[..., None] > 0.5,
|
84 |
image_np * 0.5 + color_mask * 0.5, image_np)
|
85 |
|
86 |
-
final_image = Image.fromarray(image_np.astype(np.uint8))
|
87 |
-
|
|
|
|
|
|
|
88 |
|
89 |
# Tampilkan box & label
|
90 |
for i in indices_to_use:
|
@@ -93,11 +64,34 @@ def predict_segmentation_with_legend(image):
|
|
93 |
class_id = class_ids[i]
|
94 |
label = label_map.get(class_id, str(class_id))
|
95 |
color = color_map.get(class_id, (255, 255, 0))
|
96 |
-
text_color = "white"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
97 |
|
98 |
-
|
99 |
-
|
100 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
101 |
|
102 |
# Buat legenda
|
103 |
legend = Image.new("RGB", (500, 50), (255, 255, 255))
|
@@ -118,15 +112,3 @@ def predict_segmentation_with_legend(image):
|
|
118 |
except Exception as e:
|
119 |
print("❌ Error saat segmentasi:", e)
|
120 |
return image
|
121 |
-
|
122 |
-
iface = gr.Interface(
|
123 |
-
fn=predict_segmentation_with_legend,
|
124 |
-
inputs=gr.Image(type="pil", label="Upload Citra Side Scan Sonar"),
|
125 |
-
outputs=gr.Image(type="pil", label="Hasil Segmentasi"),
|
126 |
-
title="YOLOv11 Segmentasi Citra Sonar",
|
127 |
-
description="Upload citra sonar dan dapatkan hasil segmentasi lengkap dengan warna mask, label, confidence, dan legenda warna.",
|
128 |
-
allow_flagging="never"
|
129 |
-
)
|
130 |
-
|
131 |
-
if __name__ == "__main__":
|
132 |
-
iface.launch(share=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
def predict_segmentation_with_legend(image):
|
2 |
try:
|
3 |
if model is None:
|
|
|
51 |
image_np = np.where(mask_resized[..., None] > 0.5,
|
52 |
image_np * 0.5 + color_mask * 0.5, image_np)
|
53 |
|
54 |
+
final_image = Image.fromarray(image_np.astype(np.uint8)).convert("RGBA")
|
55 |
+
|
56 |
+
# Layer teks transparan
|
57 |
+
overlay = Image.new("RGBA", final_image.size, (255, 255, 255, 0))
|
58 |
+
draw_overlay = ImageDraw.Draw(overlay)
|
59 |
|
60 |
# Tampilkan box & label
|
61 |
for i in indices_to_use:
|
|
|
64 |
class_id = class_ids[i]
|
65 |
label = label_map.get(class_id, str(class_id))
|
66 |
color = color_map.get(class_id, (255, 255, 0))
|
67 |
+
text_color = "white"
|
68 |
+
|
69 |
+
# Gambar kotak
|
70 |
+
draw_overlay.rectangle(box, outline=color + (255,), width=2)
|
71 |
+
|
72 |
+
# Teks dan ukurannya
|
73 |
+
label_text = label
|
74 |
+
score_text = f"Conf: {score:.2f}"
|
75 |
+
label_size = draw_overlay.textbbox((0, 0), label_text, font=font)
|
76 |
+
score_size = draw_overlay.textbbox((0, 0), score_text, font=font)
|
77 |
|
78 |
+
text_x = max(box[0], 0)
|
79 |
+
text_y = max(box[1] - (label_size[3] + score_size[3] + 8), 0)
|
80 |
+
|
81 |
+
# Background teks semi-transparan
|
82 |
+
draw_overlay.rectangle(
|
83 |
+
[text_x - 2, text_y - 2,
|
84 |
+
text_x + max(label_size[2], score_size[2]) + 4,
|
85 |
+
text_y + label_size[3] + score_size[3] + 6],
|
86 |
+
fill=(0, 0, 0, 160)
|
87 |
+
)
|
88 |
+
|
89 |
+
# Tulis teks
|
90 |
+
draw_overlay.text((text_x, text_y), label_text, fill=text_color, font=font)
|
91 |
+
draw_overlay.text((text_x, text_y + label_size[3] + 2), score_text, fill=text_color, font=font)
|
92 |
+
|
93 |
+
# Gabungkan overlay ke gambar
|
94 |
+
final_image = Image.alpha_composite(final_image, overlay).convert("RGB")
|
95 |
|
96 |
# Buat legenda
|
97 |
legend = Image.new("RGB", (500, 50), (255, 255, 255))
|
|
|
112 |
except Exception as e:
|
113 |
print("❌ Error saat segmentasi:", e)
|
114 |
return image
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|