qonitasal commited on
Commit
b6a5674
·
verified ·
1 Parent(s): 152b780

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +32 -50
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
- draw = ImageDraw.Draw(final_image)
 
 
 
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" if score >= CONFIDENCE_THRESHOLD else "gray"
 
 
 
 
 
 
 
 
 
97
 
98
- draw.rectangle(box, outline=color, width=2)
99
- draw.text((box[0], box[1] - 30), label, fill=text_color, font=font)
100
- draw.text((box[0], box[1] - 10), f"Confidence: {score:.2f}", fill=text_color, font=font)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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