Create app.py

app.py

@@ -0,0 +1,42 @@
+import gradio as gr
+from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
+from PIL import Image
+import torch
+import numpy as np
+
+# Load pretrained model
+processor = SegformerImageProcessor(do_reduce_labels=False)
+model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+model.eval()
+
+# Prediction function
+def segment_image(input_image):
+    inputs = processor(images=input_image, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+
+    pred_mask = torch.argmax(logits, dim=1)[0].cpu().numpy()
+    normalized_mask = (pred_mask * (255 // logits.shape[1])).astype(np.uint8)
+    output_image = Image.fromarray(normalized_mask)
+
+    # Bigger mask (3x)
+    scale_factor = 3
+    new_size = (output_image.width * scale_factor, output_image.height * scale_factor)
+    bigger_output = output_image.resize(new_size, resample=Image.NEAREST)
+
+    return bigger_output
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=segment_image,
+    inputs=gr.Image(type="pil", label="Upload Blood Smear Image"),
+    outputs=gr.Image(type="pil", label="Predicted Grayscale Mask"),
+    title="Malaria Blood Smear Segmentation (SegFormer - Pretrained)",
+    description="Upload a blood smear image to segment it using a pretrained SegFormer model (ADE20K 150 classes).",
+    examples=["examples/sample1.jpg", "examples/sample2.jpg"]  # Example images
+)
+
+if __name__ == "__main__":
+    demo.launch()