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()