resize the images

app.py

@@ -3,72 +3,111 @@ from transformers import SegformerForSemanticSegmentation, SegformerImageProcess
 from PIL import Image
 import torch
 import numpy as np
-import os
 
 # --- Documentation Strings ---
 
 USAGE_GUIDELINES = """
-## Quick Start: HemaScan Segmentation
-Upload an image of a blood smear to generate a segmentation mask.
+## 1. Quick Start Guide: HemaScan Pro
+
+HemaScan Pro uses an advanced semantic segmentation AI model to detect structural regions in microscopic blood smear images.
+
+1. **Upload Image** – Select a JPG or PNG blood smear image.
+2. **Try Samples** – Click on example images for quick testing.
+3. **Run Analysis** – Press the "Run Segmentation" button.
+4. **View Results** – A high-visibility color segmentation mask will appear.
+"""
+
+INPUT_EXPLANATION = """
+## 2. Expected Inputs
+
+| Input Field | Description | Format |
+|-------------|------------|--------|
+| Upload Image | Microscopic blood smear image | JPG / PNG |
+
+✔ Image is automatically resized to 512×512 for processing.
 """
 
-INPUT_EXPLANATION = "Upload a JPG or PNG blood smear image (512x512 auto-resized)."
-OUTPUT_EXPLANATION = "Predicted grayscale mask highlighting detected objects (scaled 4× for clarity)."
+OUTPUT_EXPLANATION = """
+## 3. Expected Outputs (Color Segmentation Mask)
+
+The output is a **color-coded segmentation mask**:
+
+• Each detected object category is assigned a distinct color.  
+• Enhances boundary clarity between cells and background.  
+• Mask is enlarged by **400% (4×)** for improved visibility.  
+
+### Example Testing
+Click any example image below to automatically run segmentation.
+"""
 
 # --------------------
-# Core Pipeline Functions
+# Model
 # --------------------
 processor = SegformerImageProcessor(do_reduce_labels=False)
-model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+model = SegformerForSemanticSegmentation.from_pretrained(
+    "nvidia/segformer-b0-finetuned-ade-512-512"
+)
 model.eval()
 
+# Create vibrant color palette
+def create_color_palette(num_classes=150):
+    np.random.seed(42)
+    palette = np.random.randint(0, 255, size=(num_classes, 3))
+    palette[0] = [0, 0, 0]
+    return palette
+
+palette = create_color_palette()
+
 def segment_image(input_image):
     if input_image is None:
-        gr.Warning("Upload an image.")
+        gr.Warning("Please upload an image.")
         return None
-        
+
     inputs = processor(images=input_image, return_tensors="pt")
+
     with torch.no_grad():
         outputs = model(**inputs)
-        logits = outputs.logits
 
+    logits = outputs.logits
     pred_mask = torch.argmax(logits, dim=1)[0].cpu().numpy()
-    num_classes = logits.shape[1]
-    # normalized_mask = (pred_mask * (255 // num_classes)).astype(np.uint8)
-    #make the mask color to whit and the background to black
-    normalized_mask = np.where(pred_mask > 0, 255, 0).astype(np.uint8)
-    output_image = Image.fromarray(normalized_mask)
 
-    # Bigger mask (4x)
+    # Apply color palette
+    colored_mask = palette[pred_mask]
+    colored_mask = colored_mask.astype(np.uint8)
+
+    output_image = Image.fromarray(colored_mask)
+
+    # Scale 4x
     scale_factor = 4
     new_size = (output_image.width * scale_factor, output_image.height * scale_factor)
     return output_image.resize(new_size, resample=Image.NEAREST)
 
 # --------------------
-# Gradio UI
+# UI
 # --------------------
-with gr.Blocks(title="HemaScan Segmentation Tool") as demo:
-    gr.Markdown("<h1 style='text-align:center; background: linear-gradient(90deg, #4facfe 0%, #00f2fe 100%); padding: 10px; color:white;'>HemaScan Segmentation Tool</h1>")
-    gr.Markdown("Analyze blood smear images and generate segmentation masks.")
+with gr.Blocks(title="HemaScan Pro") as demo:
+    gr.Markdown("<h1 style='text-align:center; background:linear-gradient(90deg,#4facfe,#00f2fe); color:white; padding:10px;'>HemaScan Pro - Blood Smear Segmentation</h1>")
 
-    with gr.Accordion("Tips & Guidelines", open=False):
+    with gr.Accordion("📘 Documentation & Usage", open=False):
         gr.Markdown(USAGE_GUIDELINES)
+        gr.Markdown("---")
         gr.Markdown(INPUT_EXPLANATION)
+        gr.Markdown("---")
         gr.Markdown(OUTPUT_EXPLANATION)
 
-    input_image = gr.Image(type="pil", label="Upload Blood Smear Image")
-    submit_button = gr.Button("Segment Image", variant="primary")
-    output_image = gr.Image(type="pil", label="Predicted Mask (4x)")
+    input_image = gr.Image(type="pil", label="Upload Blood Smear Image", width=512, height=512)
+    submit_button = gr.Button("Run Segmentation", variant="primary")
+    output_image = gr.Image(type="pil", label="Color Segmentation Mask (4x)", width=512, height=512)
 
     gr.Examples(
         examples=["data/1.png", "data/2.png", "data/3.png", "data/211.png"],
-        inputs=[input_image],
-        outputs=[output_image],
+        inputs=input_image,
+        outputs=output_image,
         fn=segment_image,
-        cache_examples=False
+        cache_examples=False,
     )
 
-    submit_button.click(fn=segment_image, inputs=input_image, outputs=output_image)
+    submit_button.click(segment_image, input_image, output_image)
 
 if __name__ == "__main__":
     demo.launch()