MedAI / app.py
samH98's picture
Update app.py
aa88fe1
import gradio as gr
import cv2
import torch
# Load YOLOv7 model
model = torch.hub.load('models', 'custom', 'models/100epoch.pt', force_reload=True, source='local', trust_repo=True)
model.eval()
# Create a mapping for detected labels
label_mapping = {
'A': 'Adenocarcinoma',
'B': 'Small Cell Carcinoma',
'E': 'Large Cell Carcinoma',
'G': 'Squamous Cell Carcinoma'
}
def process_image(input_image):
# Perform inference on the input image
results = model(input_image)
img = results.render()[0]
if results.pred is not None and len(results.pred[0]) > 0:
detection = results.pred[0]
class_index = int(detection[0, -1])
print(class_index)
class_label = model.names[class_index]
mapped_label = label_mapping.get(class_label, "Unknown")
else:
mapped_label = "No cancer detected"
return img, mapped_label
iface = gr.Interface(
fn=process_image,
inputs=gr.components.Image(type='pil', label="Input Image").style(height=280),
outputs=[gr.components.Image(type='pil', label="Processed Image").style(height=280), gr.components.Textbox(label="Detected Cancer Type")],
live=True,
title="Lung Cancer Detector ⚕️",
description="The AI model was trained to detect the following types of lung cancer:\n"
"1. Adenocarcinoma (A)\n"
"2. Small Cell Carcinoma (B)\n"
"3. Large Cell Carcinoma (E)\n"
"4. Squamous Cell Carcinoma (G)\n\n"
"How to Use :\n"
"1. Upload a CT scan image of a patient's lungs.\n"
"2. The app will display the predicted type of lung cancer.",
theme=gr.themes.Monochrome(font=[gr.themes.GoogleFont("Noto Serif"), "Preahvihear", "sans-serif"])
)
if __name__ == '__main__':
iface.launch()