Update: Medical Image Segmentation (2026-01-27 15:43)

app.py

@@ -31,8 +31,8 @@ License: MIT
 
 @dataclass
 class Configs:
-    NUM_CLASSES: int = 4  # including background.
-    CLASSES: Tuple[str, ...] = ("Large bowel", "Small bowel", "Stomach")
+    NUM_CLASSES: int = 4  # bao gồm background
+    CLASSES: Tuple[str, ...] = ("Ruột già", "Ruột non", "Dạ dày")
     IMAGE_SIZE: Tuple[int, int] = (288, 288)  # W, H
     MEAN: Tuple[float, ...] = (0.485, 0.456, 0.406)
     STD: Tuple[float, ...] = (0.229, 0.224, 0.225)
@@ -43,17 +43,6 @@ class Configs:
 def get_model(*, model_path, num_classes):
     """
     Load pre-trained SegFormer model.
-    
-    Args:
-        model_path (str): Path to model directory containing config.json and pytorch_model.bin
-        num_classes (int): Number of segmentation classes
-    
-    Returns:
-        SegformerForSemanticSegmentation: Loaded model
-    
-    Raises:
-        FileNotFoundError: If model files not found
-        RuntimeError: If model loading fails
     """
     model = SegformerForSemanticSegmentation.from_pretrained(
         model_path, 
@@ -67,27 +56,10 @@ def get_model(*, model_path, num_classes):
 def predict(input_image, model=None, preprocess_fn=None, device="cpu"):
     """
     Perform semantic segmentation on input medical image.
-    
-    Args:
-        input_image (PIL.Image): Input medical image
-        model (SegformerForSemanticSegmentation): Trained segmentation model
-        preprocess_fn (callable): Image preprocessing function
-        device (str or torch.device): Device to run inference on ('cpu' or 'cuda')
-    
-    Returns:
-        Tuple[PIL.Image, str]: 
-            - Color-coded segmentation mask
-            - Text with confidence scores for each organ
-    
-    Raises:
-        ValueError: If input image is invalid
-        RuntimeError: If model inference fails
-    
-    Example:
-        >>> from PIL import Image
-        >>> img = Image.open('medical_scan.png')
-        >>> output, info = predict(img, model, preprocess_fn, device)
     """
+    if input_image is None:
+        return None, []
+
     shape_H_W = input_image.size[::-1]
     input_tensor = preprocess_fn(input_image)
     input_tensor = input_tensor.unsqueeze(0).to(device)
@@ -102,10 +74,19 @@ def predict(input_image, model=None, preprocess_fn=None, device="cpu"):
     confidence_map = probs.max(dim=1)[0].cpu().squeeze().numpy()
 
     # Create segmentation info with confidence
-    seg_info = [
-        (preds_argmax == idx, f"{class_name} (confidence: {confidence_map[preds_argmax == idx].mean():.2%})")
-        for idx, class_name in enumerate(Configs.CLASSES, 1)
-    ]
+    seg_info = []
+    
+    # Classes: 1=Ruột già, 2=Ruột non, 3=Dạ dày
+    for idx, class_name in enumerate(Configs.CLASSES, 1):
+        mask = preds_argmax == idx
+        if mask.sum() > 0:
+            # Chỉ tính confidence nếu có pixel được dự đoán
+            conf_score = confidence_map[mask].mean()
+            label = f"{class_name} ({conf_score:.1%})"
+            seg_info.append((mask, label))
+        else:
+            # Không hiển thị label nếu không phát hiện được
+            pass
 
     return (input_image, seg_info)
 
@@ -113,20 +94,13 @@ def predict(input_image, model=None, preprocess_fn=None, device="cpu"):
 if __name__ == "__main__":
     """
     Main application entry point.
-    
-    Initializes:
-    - Device selection (GPU/CPU)
-    - Model loading and setup
-    - Image preprocessing pipeline
-    - Gradio web interface
-    
-    The web interface allows users to:
-    - Upload medical images
-    - Generate segmentation predictions
-    - View color-coded organ detection
-    - See confidence scores
     """
-    class2hexcolor = {"Dạ dày": "#007fff", "Ruột non": "#009A17", "Ruột già": "#FF0000"}
+    # Mapping màu sắc cho hiển thị
+    class2hexcolor = {
+        "Dạ dày": "#007fff",      # Xanh dương
+        "Ruột non": "#009A17",    # Xanh lá
+        "Ruột già": "#FF0000"     # Đỏ
+    }
 
     DEVICE = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
 
@@ -140,11 +114,17 @@ if __name__ == "__main__":
         print(f"Error loading model: {e}")
         model_dir = "./segformer_trained_weights"
     
-    # Sử dụng đường dẫn từ W&B artifact hoặc mô hình cục bộ
+    # Load model
+    print(f"Loading model from: {model_dir}")
     model = get_model(model_path=model_dir, num_classes=Configs.NUM_CLASSES)
     model.to(DEVICE)
     model.eval()
-    _ = model(torch.randn(1, 3, *Configs.IMAGE_SIZE[::-1], device=DEVICE))
+    
+    # Warmup
+    try:
+        _ = model(torch.randn(1, 3, *Configs.IMAGE_SIZE[::-1], device=DEVICE))
+    except Exception as e:
+        print(f"Warmup warning: {e}")
 
     preprocess = TF.Compose(
         [
@@ -154,50 +134,57 @@ if __name__ == "__main__":
         ]
     )
 
-    with gr.Blocks(title="Medical Image Segmentation") as demo:
+    with gr.Blocks(title="Phân Đoạn Ảnh Y Tế") as demo:
         gr.Markdown("""
-        <h1><center>🏥 Medical Image Segmentation with UW-Madison GI Tract Dataset</center></h1>
-        <p><center>Phân đoạn tự động các cơ quan trong ảnh Y tế (Dạ dày, Ruột non, Ruột già)</center></p>
+        <h1><center>🏥 Phân Đoạn Ảnh Y Tế - Tập Dữ Liệu UW-Madison GI Tract</center></h1>
+        <p><center>Hệ thống tự động phát hiện và phân đoạn các cơ quan tiêu hóa: Dạ dày, Ruột non, Ruột già.</center></p>
         """)
         
         with gr.Row():
             with gr.Column():
-                gr.Markdown("### 📥 Input Image")
-                img_input = gr.Image(type="pil", height=360, width=360, label="Input image")
+                gr.Markdown("### 📥 Ảnh Đầu Vào")
+                img_input = gr.Image(type="pil", height=360, width=360, label="Tải ảnh lên")
             
             with gr.Column():
-                gr.Markdown("### 📊 Predictions")
-                img_output = gr.AnnotatedImage(label="Predictions", height=360, width=360, color_map=class2hexcolor)
-
-        section_btn = gr.Button("🎯 Generate Predictions", size="lg")
+                gr.Markdown("### 📊 Kết Quả Dự Đoán")
+                # AnnotatedImage hiển thị ảnh gốc + các lớp mask
+                img_output = gr.AnnotatedImage(
+                    label="Kết quả phân đoạn", 
+                    height=360, 
+                    width=360,
+                    color_map=class2hexcolor
+                )
+
+        section_btn = gr.Button("🎯 Chạy Phân Đoạn", size="lg", variant="primary")
         section_btn.click(partial(predict, model=model, preprocess_fn=preprocess, device=DEVICE), img_input, img_output)
 
         gr.Markdown("---")
-        gr.Markdown("### 📸 Sample Images")
+        gr.Markdown("### 📸 Ảnh Mẫu (Click để thử)")
         
         images_dir = glob(os.path.join(os.getcwd(), "samples") + os.sep + "*.png")
-        examples = [i for i in np.random.choice(images_dir, size=min(10, len(images_dir)), replace=False)]
-        
-        gr.Examples(
-            examples=examples,
-            inputs=img_input,
-            outputs=img_output,
-            fn=partial(predict, model=model, preprocess_fn=preprocess, device=DEVICE),
-            cache_examples=False,
-            label="Click to load example"
-        )
+        if len(images_dir) > 0:
+            examples = [i for i in np.random.choice(images_dir, size=min(10, len(images_dir)), replace=False)]
+            
+            gr.Examples(
+                examples=examples,
+                inputs=img_input,
+                outputs=img_output,
+                fn=partial(predict, model=model, preprocess_fn=preprocess, device=DEVICE),
+                cache_examples=False,
+                label="Thư viện ảnh mẫu"
+            )
         
         gr.Markdown("""
         ---
-        ### 🎨 Color Legend
-        - 🔵 **Blue (#007fff)**: Dạ dày (Stomach)
-        - 🟢 **Green (#009A17)**: Ruột non (Small bowel)
-        - 🔴 **Red (#FF0000)**: Ruột già (Large bowel)
+        ### 🎨 Chú Thích Màu Sắc
+        - 🔵 **Xanh Dương**: Dạ dày
+        - 🟢 **Xanh Lá**: Ruột non
+        - 🔴 **Đỏ**: Ruột già
         
-        ### ℹ️  Information
-        - Model: SegFormer (HuggingFace Transformers)
-        - Input size: 288 × 288 pixels
-        - Framework: PyTorch + Gradio
+        ### ℹ️  Thông Tin Hệ Thống
+        - **Mô hình**: SegFormer (mit-b0)
+        - **Kích thước đầu vào**: 288 × 288 pixels
+        - **Framework**: PyTorch + Gradio
         """)
 
     demo.launch()