Spaces:
Sleeping
Sleeping
| 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) |