|
from typing import Tuple |
|
from ultralytics import YOLO |
|
from ultralytics.engine.results import Boxes |
|
from ultralytics.utils.plotting import Annotator |
|
|
|
import gradio as gr |
|
|
|
cell_detector = YOLO("./weights/yolo_uninfected_cells.pt") |
|
yolo_detector = YOLO("./weights/yolo_infected_cells.pt") |
|
redetr_detector = YOLO("./weights/redetr_infected_cells.pt") |
|
|
|
models = {"Yolo V11": yolo_detector, "Real Time Detection Transformer": redetr_detector} |
|
|
|
|
|
|
|
def inference(image, model, conf) -> Tuple[str, str, str]: |
|
bboxes = [] |
|
labels = [] |
|
healthy_cell_count = 0 |
|
unhealthy_cell_count = 0 |
|
cells_results = cell_detector.predict(image, conf=0.4) |
|
selected_model_results = models[model].predict( |
|
image, conf=conf |
|
) |
|
|
|
for cell_result in cells_results: |
|
boxes: Boxes = cell_result.boxes |
|
healthy_cells_bboxes = boxes.xyxy.tolist() |
|
healthy_cell_count += len(healthy_cells_bboxes) |
|
bboxes.extend(healthy_cells_bboxes) |
|
labels.extend(["healthy"] * healthy_cell_count) |
|
|
|
for res in selected_model_results: |
|
boxes: Boxes = res.boxes |
|
unhealthy_cells_bboxes = boxes.xyxy.tolist() |
|
unhealthy_cell_count += len(unhealthy_cells_bboxes) |
|
bboxes.extend(unhealthy_cells_bboxes) |
|
labels.extend(["unhealthy"] * unhealthy_cell_count) |
|
|
|
annotator = Annotator(image, font_size=5, line_width=1) |
|
|
|
for box, label in zip(bboxes, labels): |
|
annotator.box_label(box, label) |
|
|
|
img = annotator.result() |
|
return (img, healthy_cell_count, unhealthy_cell_count) |
|
|
|
|
|
ifer = gr.Interface( |
|
fn=inference, |
|
inputs=[ |
|
gr.Image(label="Input Image", type="numpy"), |
|
gr.Dropdown( |
|
choices=["Yolo V11", "Real Time Detection Transformer"], multiselect=False |
|
), |
|
gr.Slider(minimum=0.01, maximum=1) |
|
], |
|
outputs=[ |
|
gr.Image(label="Output Image", type="numpy"), |
|
gr.Textbox(label="Healthy Cells Count"), |
|
gr.Textbox(label="Infected Cells Count"), |
|
], |
|
title="Blood Cancer Cell Detection and Counting" |
|
) |
|
|
|
ifer.launch(share=True) |