|
import gradio as gr |
|
from ultralytics import YOLO |
|
import cv2 |
|
import numpy as np |
|
from PIL import Image, ImageDraw, ImageFont |
|
import base64 |
|
from io import BytesIO |
|
import tempfile |
|
import os |
|
from pathlib import Path |
|
import shutil |
|
from openpyxl import Workbook, load_workbook |
|
|
|
|
|
model = YOLO("best.pt") |
|
|
|
|
|
uploaded_folder = Path('Uploaded_Picture') |
|
predicted_folder = Path('Predicted_Picture') |
|
uploaded_folder.mkdir(parents=True, exist_ok=True) |
|
predicted_folder.mkdir(parents=True, exist_ok=True) |
|
|
|
|
|
xlsx_db_file = Path('patient_predictions.xlsx') |
|
|
|
|
|
if not xlsx_db_file.exists(): |
|
workbook = Workbook() |
|
sheet = workbook.active |
|
sheet.title = "Predictions" |
|
sheet.append(["Name", "Age", "Medical Record", "Sex", "Result", "Image Path"]) |
|
workbook.save(xlsx_db_file) |
|
|
|
def predict_image(input_image, name, age, medical_record, sex): |
|
if input_image is None: |
|
return None, "Please Input The Image" |
|
|
|
|
|
image_np = np.array(input_image) |
|
|
|
|
|
if len(image_np.shape) == 2: |
|
image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB) |
|
elif image_np.shape[2] == 4: |
|
image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB) |
|
|
|
|
|
results = model(image_np) |
|
|
|
|
|
image_with_boxes = image_np.copy() |
|
raw_predictions = [] |
|
|
|
if results[0].boxes: |
|
|
|
highest_confidence_result = max(results[0].boxes, key=lambda x: x.conf.item()) |
|
|
|
|
|
class_index = highest_confidence_result.cls.item() |
|
if class_index == 0: |
|
label = "Cataract" |
|
color = (255, 0, 0) |
|
elif class_index == 1: |
|
label = "Normal" |
|
color = (0, 255, 0) |
|
|
|
confidence = highest_confidence_result.conf.item() |
|
xmin, ymin, xmax, ymax = map(int, highest_confidence_result.xyxy[0]) |
|
|
|
|
|
box_width = xmax - xmin |
|
box_height = ymax - ymin |
|
avg_dimension = (box_width + box_height) / 2 |
|
|
|
|
|
radius = int(avg_dimension / 12) |
|
|
|
|
|
center_x = int((xmin + xmax) / 2) |
|
center_y = int((ymin + ymax) / 2) |
|
|
|
|
|
cv2.circle(image_with_boxes, (center_x, center_y), radius, color, 2) |
|
|
|
|
|
font_scale = 1.0 |
|
thickness = 2 |
|
|
|
|
|
(text_width, text_height), baseline = cv2.getTextSize(f'{label} {confidence:.2f}', cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness) |
|
cv2.rectangle(image_with_boxes, (xmin, ymin - text_height - baseline), (xmin + text_width, ymin), (0, 0, 0), cv2.FILLED) |
|
|
|
|
|
cv2.putText(image_with_boxes, f'{label} {confidence:.2f}', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness) |
|
|
|
raw_predictions.append(f"Label: {label}, Confidence: {confidence:.2f}, Circle Center: [{center_x}, {center_y}], Radius: {radius}") |
|
|
|
raw_predictions_str = "\n".join(raw_predictions) |
|
|
|
|
|
pil_image_with_boxes = Image.fromarray(image_with_boxes) |
|
|
|
|
|
pil_image_with_boxes = add_text_and_watermark(pil_image_with_boxes, name, age, medical_record, sex, label) |
|
|
|
|
|
image_name = f"{name}-{age}-{sex}-{medical_record}.png" |
|
input_image.save(uploaded_folder / image_name) |
|
pil_image_with_boxes.save(predicted_folder / image_name) |
|
|
|
|
|
buffered = BytesIO() |
|
pil_image_with_boxes.save(buffered, format="PNG") |
|
predicted_image_base64 = base64.b64encode(buffered.getvalue()).decode() |
|
|
|
|
|
append_patient_info_to_xlsx(name, age, medical_record, sex, label, image_name) |
|
|
|
return pil_image_with_boxes, raw_predictions_str |
|
|
|
def add_watermark(image): |
|
try: |
|
logo = Image.open('image-logo.png').convert("RGBA") |
|
image = image.convert("RGBA") |
|
basewidth = 100 |
|
wpercent = (basewidth / float(logo.size[0])) |
|
hsize = int((float(wpercent) * logo.size[1])) |
|
logo = logo.resize((basewidth, hsize), Image.LANCZOS) |
|
position = (image.width - logo.width - 10, image.height - logo.height - 10) |
|
transparent = Image.new('RGBA', (image.width, image.height), (0, 0, 0, 0)) |
|
transparent.paste(image, (0, 0)) |
|
transparent.paste(logo, position, mask=logo) |
|
return transparent.convert("RGB") |
|
except Exception as e: |
|
print(f"Error adding watermark: {e}") |
|
return image |
|
|
|
def add_text_and_watermark(image, name, age, medical_record, sex, label): |
|
draw = ImageDraw.Draw(image) |
|
font_size = 24 |
|
try: |
|
font = ImageFont.truetype("font.ttf", size=font_size) |
|
except IOError: |
|
font = ImageFont.load_default() |
|
print("Error: cannot open resource, using default font.") |
|
|
|
text = f"Name: {name}, Age: {age}, Medical Record: {medical_record}, Sex: {sex}, Result: {label}" |
|
text_bbox = draw.textbbox((0, 0), text, font=font) |
|
text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1] |
|
text_x = 20 |
|
text_y = 40 |
|
padding = 10 |
|
draw.rectangle([text_x - padding, text_y - padding, text_x + text_width + padding, text_y + text_height + padding], fill="black") |
|
draw.text((text_x, text_y), text, fill=(255, 255, 255, 255), font=font) |
|
|
|
image_with_watermark = add_watermark(image) |
|
return image_with_watermark |
|
|
|
def append_patient_info_to_xlsx(name, age, medical_record, sex, result, image_path): |
|
if not xlsx_db_file.exists(): |
|
workbook = Workbook() |
|
sheet = workbook.active |
|
sheet.title = "Predictions" |
|
sheet.append(["Name", "Age", "Medical Record", "Sex", "Result", "Image Path"]) |
|
workbook.save(xlsx_db_file) |
|
|
|
workbook = load_workbook(xlsx_db_file) |
|
sheet = workbook.active |
|
sheet.append([name, age, medical_record, sex, result, str(image_path)]) |
|
workbook.save(xlsx_db_file) |
|
|
|
return str(xlsx_db_file) |
|
|
|
def download_folder(folder): |
|
zip_path = os.path.join(tempfile.gettempdir(), f"{folder}.zip") |
|
shutil.make_archive(zip_path.replace('.zip', ''), 'zip', folder) |
|
return zip_path |
|
|
|
def interface(name, age, medical_record, sex, input_image): |
|
if input_image is None: |
|
return None, "Please upload an image.", None |
|
|
|
output_image, raw_result = predict_image(input_image, name, age, medical_record, sex) |
|
|
|
return output_image, raw_result, str(xlsx_db_file) |
|
|
|
def download_predicted_folder(): |
|
return download_folder(predicted_folder) |
|
|
|
def download_uploaded_folder(): |
|
return download_folder(uploaded_folder) |
|
|
|
with gr.Blocks() as demo: |
|
with gr.Column(): |
|
gr.Markdown("# Cataract Detection System") |
|
gr.Markdown("Upload an image to detect cataract and add patient details.") |
|
gr.Markdown("This application uses YOLOv10 with mAP=0.994") |
|
|
|
with gr.Column(): |
|
name = gr.Textbox(label="Name") |
|
age = gr.Number(label="Age") |
|
medical_record = gr.Number(label="Medical Record") |
|
sex = gr.Radio(["Male", "Female"], label="Sex") |
|
input_image = gr.Image(type="pil", label="Upload an Image", image_mode="RGB") |
|
|
|
with gr.Column(): |
|
submit_btn = gr.Button("Submit") |
|
output_image = gr.Image(type="pil", label="Predicted Image") |
|
|
|
with gr.Row(): |
|
raw_result = gr.Textbox(label="Prediction Result") |
|
|
|
with gr.Row(): |
|
download_xlsx_btn = gr.Button("Download Patient Information (XLSX)") |
|
download_uploaded_btn = gr.Button("Download Uploaded Images") |
|
download_predicted_btn = gr.Button("Download Predicted Images") |
|
|
|
xlsx_file = gr.File(label="Patient Information XLSX File") |
|
uploaded_folder_file = gr.File(label="Uploaded Images Zip File") |
|
predicted_folder_file = gr.File(label="Predicted Images Zip File") |
|
|
|
submit_btn.click(fn=interface, inputs=[name, age, medical_record, sex, input_image], outputs=[output_image, raw_result, xlsx_file]) |
|
download_xlsx_btn.click(fn=lambda: str(xlsx_db_file), outputs=xlsx_file) |
|
download_uploaded_btn.click(fn=download_uploaded_folder, outputs=uploaded_folder_file) |
|
download_predicted_btn.click(fn=download_predicted_folder, outputs=predicted_folder_file) |
|
|
|
demo.launch() |