Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor | |
| from PIL import Image | |
| import torch | |
| import numpy as np | |
| # --- Documentation Strings --- | |
| USAGE_GUIDELINES = """ | |
| ## 1. Quick Start Guide: HemaScan Pro (Binary Mask) | |
| HemaScan Pro generates a high-contrast black & white segmentation mask. | |
| 1. Upload a blood smear image (JPG/PNG). | |
| 2. Click "Run Segmentation". | |
| 3. View the generated binary mask. | |
| """ | |
| INPUT_EXPLANATION = """ | |
| ## 2. Expected Inputs | |
| | Field | Requirement | | |
| |-------|------------| | |
| | Upload Image | JPG / PNG blood smear image | | |
| ✔ Automatically resized to 512×512. | |
| """ | |
| OUTPUT_EXPLANATION = """ | |
| ## 3. Output Description (Black & White Mask) | |
| • Background = White | |
| • Detected Regions = Black | |
| • Enlarged by 400% (4×) for clarity | |
| • Clean binary medical-style visualization | |
| """ | |
| # -------------------- | |
| # Model | |
| # -------------------- | |
| processor = SegformerImageProcessor(do_reduce_labels=False) | |
| model = SegformerForSemanticSegmentation.from_pretrained( | |
| "nvidia/segformer-b0-finetuned-ade-512-512" | |
| ) | |
| model.eval() | |
| def segment_image(input_image): | |
| if input_image is None: | |
| 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 | |
| pred_mask = torch.argmax(logits, dim=1)[0].cpu().numpy() | |
| # Convert to binary mask (object vs background) | |
| binary_mask = np.where(pred_mask == 0, 255, 0).astype(np.uint8) | |
| output_image = Image.fromarray(binary_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) | |
| # -------------------- | |
| # UI | |
| # -------------------- | |
| with gr.Blocks(title="Malaria Cell Segmentation Tool") as demo: | |
| gr.Markdown("<h1 style='text-align:center; background:linear-gradient(90deg,#4facfe,#00f2fe); color:white; padding:10px;'>HemaScan Pro - Binary Segmentation</h1>") | |
| with gr.Accordion(" Documentation", open=False): | |
| gr.Markdown(USAGE_GUIDELINES) | |
| gr.Markdown("---") | |
| gr.Markdown(INPUT_EXPLANATION) | |
| gr.Markdown("---") | |
| gr.Markdown(OUTPUT_EXPLANATION) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("## Step 1: Upload Blood Smear Image") | |
| # Define Input component directly inside the column (No .render() needed) | |
| input_image = gr.Image(type="pil", label="Step 1: Upload Blood Smear Image", width=600, height=600) | |
| gr.Markdown("## Step 2: Click Submit for Segmentation") | |
| with gr.Row(): | |
| submit_button = gr.Button("Submit for Segmentation", variant="primary") | |
| with gr.Column(scale=1): | |
| gr.Markdown("## Output") | |
| # Define Output component directly inside the column (No .render() needed) | |
| output_image = gr.Image(type="pil", label="Step 3: Predicted Masks", width=600, height=600) | |
| gr.Markdown("---") | |
| gr.Markdown("## Example Images") | |
| gr.Examples( | |
| examples=["data/1.png", "data/2.png", "data/3.png"], | |
| inputs=input_image, | |
| outputs=output_image, | |
| fn=segment_image, | |
| cache_examples=False, | |
| ) | |
| submit_button.click(segment_image, input_image, output_image) | |
| if __name__ == "__main__": | |
| demo.launch() | |