Update app.py

app.py

@@ -1,25 +1,44 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation
 from PIL import Image
+import torch
+import numpy as np
 
-# โหลด pipeline สำหรับ semantic segmentation
-segmenter = pipeline("image-segmentation", model="nvidia/segformer-b0-finetuned-ade-512-512")
+# โหลดโมเดล SegFormer
+processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+model = AutoModelForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 
-def segment_image(image):
-    segments = segmenter(image)
-    # สร้างภาพรวมของ segmentation mask ด้วยสีที่ต่างกัน
-    mask = Image.new("RGBA", image.size)
-    for segment in segments:
-        mask_part = Image.new("RGBA", image.size, segment["color"])
-        mask_part.putalpha(int(255 * segment["score"]))
-        mask = Image.alpha_composite(mask, mask_part)
-    return mask
+def segment(image):
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits  # (1, num_labels, H, W)
+
+    upsampled_logits = torch.nn.functional.interpolate(
+        logits,
+        size=image.size[::-1],  # (H, W)
+        mode="bilinear",
+        align_corners=False,
+    )[0]
+
+    predicted = upsampled_logits.argmax(0).numpy()
+
+    # สร้างภาพ segmentation mask
+    colored_mask = Image.fromarray(segmentation_to_color(predicted))
+
+    return colored_mask
+
+# แปลง mask เป็นสี (แบบง่าย)
+def segmentation_to_color(segmentation):
+    num_classes = np.max(segmentation) + 1
+    colors = np.random.randint(0, 255, size=(num_classes, 3), dtype=np.uint8)
+    return colors[segmentation]
 
 # Gradio UI
 gr.Interface(
-    fn=segment_image,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs="image",
+    fn=segment,
+    inputs=gr.Image(type="pil", label="Upload an image"),
+    outputs=gr.Image(type="pil", label="Segmentation Mask"),
     title="Semantic Segmentation with SegFormer",
-    description="ใช้โมเดล nvidia/segformer-b0-finetuned-ade-512-512 ตรวจจับพื้นที่ของวัตถุต่าง ๆ ในภาพ"
+    description="ใช้โมเดล NVIDIA SegFormer สำหรับ Semantic Segmentation"
 ).launch()