app.py

from typing import Tuple
from ultralytics import YOLO
from ultralytics.engine.results import Boxes
from ultralytics.utils.plotting import Annotator
import gradio as gr
import os

# --- Model Loading ---
try:
    cell_detector = YOLO("./weights/yolo_uninfected_cells.pt")
    yolo_detector = YOLO("./weights/yolo_infected_cells.pt")
    redetr_detector = YOLO("./weights/redetr_infected_cells.pt")
except Exception as e:
    print(f"Warning: Model loading failed. Ensure weights files are in ./weights/ directory. Error: {e}")
    # Define placeholder models if real models fail to load (for UI development)
    class DummyYOLO:
        def predict(self, image, conf=0.5):
            # Return dummy results structure
            class DummyBoxes:
                xyxy = []
            class DummyResult:
                boxes = DummyBoxes()
            return [DummyResult()]
    cell_detector = DummyYOLO()
    yolo_detector = DummyYOLO()
    redetr_detector = DummyYOLO()

models = {"Yolo V11": yolo_detector, "Real Time Detection Transformer": redetr_detector}

# --- Documentation Strings ---

USAGE_GUIDELINES = """
## 1. Quick Start Guide: Cell Detection and Counting
This application uses two specialized Artificial Intelligence models to analyze a blood smear image, simultaneously detecting both healthy and potentially infected (unhealthy) cells.

1.  **Upload**: Upload a clear blood smear image (JPG or PNG) using the 'Input Image' box.
2.  **Select Model**: Choose between the two detection models: `Yolo V11` (often fast and accurate for common objects) or `Real Time Detection Transformer`.
3.  **Adjust Confidence**: Use the slider to set the **Confidence Threshold**. (A higher value means the model must be more certain of a detection.)
4.  **Run**: Click the **"Submit"** button.
5.  **Review**: The output image will show bounding boxes around detected cells (colors based on model configuration), and the counts will be displayed below.

### Key Requirement:
*   The system uses **two independent models**: one strictly for **Healthy Cells**, and one (the selected model) for **Infected Cells**.
"""

INPUT_EXPLANATION = """
## 2. Expected Inputs

| Parameter | Purpose | Range/Options | Guidance for Non-Tech Users |
| :--- | :--- | :--- | :--- |
| **Input Image** | The microscopic blood smear image to be analyzed. | JPG, PNG format. | Ensure the image is clear and focused. |
| **Model Selection** | Chooses the AI architecture used for detecting **Infected Cells**. | Yolo V11, Real Time Detection Transformer | Start with the default (`Yolo V11`) unless specific performance is required. |
| **Confidence Threshold** | The minimum probability required for a detection box to be shown. | 0.01 to 1.00 | Setting this too low (e.g., 0.1) may show many false positives. Setting it too high (e.g., 0.9) may miss real cells. Start around 0.5. |
"""

OUTPUT_EXPLANATION = """
## 3. Expected Outputs

| Output Field | Description | Interpretation |
| :--- | :--- | :--- |
| **Output Image** | The input image with colored bounding boxes drawn around every detected cell. | Visually confirms the location and classification of each cell. |
| **Healthy Cells Count** | The total number of cells detected by the dedicated *uninfected* cell model. | Provides a baseline count of normal cells. |
| **Infected Cells Count** | The total number of cells detected by the *selected* model (Yolo V11 or RT DETR). | This represents the count of potentially cancerous/abnormal cells. |

"""

# --- Example Data Setup ---
SAMPLE_EXAMPLES = [
    ["./blood_smear_1.jpg", "Yolo V11", 0.5],
    ["./blood_smear_2.jpg", "Real Time Detection Transformer", 0.45],
]

# ----------------- Core Inference Function -----------------

def inference(image, model, conf) -> Tuple[str, str, str]:
    if image is None:
        gr.Error("Please upload an image.")
        return None, "0", "0"
    if model not in models:
        gr.Error(f"Selected model '{model}' is not available.")
        return None, "0", "0"

    bboxes = []
    labels = []
    
    # Use lists to store counts that will be incremented
    healthy_cell_count_list = [0]
    unhealthy_cell_count_list = [0]
    
    # 1. Healthy Cell Detection (Fixed model and fixed confidence 0.4)
    cells_results = cell_detector.predict(image, conf=0.4)
    for cell_result in cells_results:
        boxes: Boxes = cell_result.boxes
        healthy_cells_bboxes = boxes.xyxy.tolist()
        healthy_cell_count_list[0] += len(healthy_cells_bboxes)
        bboxes.extend(healthy_cells_bboxes)
        # Note: YOLO classes start at 0. Here we use custom labels 'healthy'
        labels.extend(["healthy"] * len(healthy_cells_bboxes)) 

    # 2. Infected Cell Detection (Selected model and user-defined confidence)
    selected_model_results = models[model].predict(image, conf=conf)

    for res in selected_model_results:
        boxes: Boxes = res.boxes
        unhealthy_cells_bboxes = boxes.xyxy.tolist()
        unhealthy_cell_count_list[0] += len(unhealthy_cells_bboxes)
        bboxes.extend(unhealthy_cells_bboxes)
        # Note: Use 'unhealthy' label for the selected model's output
        labels.extend(["unhealthy"] * len(unhealthy_cells_bboxes)) 

    # 3. Annotation
    annotator = Annotator(image, font_size=30, line_width=4, pil=True) # Increased font/width for visibility

    # Define colors based on label
    color_map = {"healthy": (0, 255, 0), "unhealthy": (255, 0, 0)} # Green for healthy, Red for unhealthy

    for box, label in zip(bboxes, labels):
        # Annotator expects a list of 4 float coords and an optional label string
        annotator.box_label(box, label, color=color_map.get(label, (255, 255, 255)))

    img = annotator.result()
    
    # Return results as strings for the Textbox components
    return (img, str(healthy_cell_count_list[0]), str(unhealthy_cell_count_list[0]))

# ----------------- Gradio Interface (Blocks) -----------------

with gr.Blocks(title="Blood Cell Detection") as ifer:
    gr.Markdown("<h1 style='text-align: center;'>  Blood Cell Cancer Detection and Counting </h1>")
    gr.Markdown("Uses specialized object detection models to count healthy and infected cells in blood smear images.")

    # 1. Documentation
    with gr.Accordion(" Tips & Guidelines ", open=False):
        gr.Markdown(USAGE_GUIDELINES)
        gr.Markdown("---")
        gr.Markdown(INPUT_EXPLANATION)
        gr.Markdown("---")
        gr.Markdown(OUTPUT_EXPLANATION)

    # 2. Interface Inputs
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Step 1: Upload Image ")
            image_input = gr.Image(label="Input Image", type="pil")
        with gr.Column():
            gr.Markdown("## Step 2: Set Parameters")
            model_selection = gr.Dropdown(
                label="Select Detection Model (for Infected Cells)",
                choices=["Yolo V11", "Real Time Detection Transformer"], 
                multiselect=False, 
                value="Yolo V11"
            )
            conf_slider = gr.Slider(
                minimum=0.01, 
                maximum=1, 
                value=0.5, 
                step=0.01, 
                label="Confidence Threshold (Min. certainty required)"
            )
    
    gr.Markdown("## Step 3: Click Analyze Image")
    with gr.Row():
        submit_button = gr.Button("Analyze Image", variant="primary")

    # 3. Interface Outputs
    gr.Markdown("## Results")
    output_image = gr.Image(label="Output Image (Detected Cells)", type="numpy")
    
    with gr.Row():
        healthy_count = gr.Textbox(label="Healthy Cells Count")
        unhealthy_count = gr.Textbox(label="Infected Cells Count")
        
    # 4. Examples
    gr.Markdown("---")
    gr.Markdown("## Example Inputs")
    gr.Examples(
        examples=SAMPLE_EXAMPLES,
        inputs=[image_input, model_selection, conf_slider],
        outputs=[output_image, healthy_count, unhealthy_count],
        fn=inference,
        cache_examples=False,
        label="Click a row to load the image and parameters"
    )

    # Event Handler
    submit_button.click(
        fn=inference,
        inputs=[image_input, model_selection, conf_slider],
        outputs=[output_image, healthy_count, unhealthy_count]
    )

if __name__ == "__main__":
    ifer.launch(share=True)