MrNavi/MultiClassification_Model

1import gradio as gr2import torch3from ultralytics import YOLO4from PIL import Image5 6# Load the trained model (kidney trained model)7model = YOLO("best.pt")8 9# Disease information for kidney10disease_info = {11    "Cyst": {12        "description": "Fluid-filled sacs that can develop on the kidneys.",13        "causes": "Genetic factors, age.",14        "treatment": "Monitoring, surgical removal if symptomatic."15    },16    "Normal": {17        "description": "Normal kidney tissue without abnormalities.",18        "causes": "N/A",19        "treatment": "N/A"20    },21    "Stone": {22        "description": "Mineral deposits that form in the kidneys.",23        "causes": "Dehydration, diet, obesity.",24        "treatment": "Increased fluid intake, medication, surgery for larger stones."25    },26    "Tumor": {27        "description": "Abnormal growth in the kidney tissue.",28        "causes": "Genetic mutations, smoking, obesity.",29        "treatment": "Surgery, chemotherapy, targeted therapy."30    },31    "No Disease": {32        "description": "No abnormalities detected.",33        "causes": "N/A",34        "treatment": "N/A"35    }36}37 38# Function for the first step: Image prediction39def predict(image):40    # Perform inference on the image41    results = model(image)42    43    # Get the top predicted class and confidence44    class_index = results[0].probs.top145    class_name = results[0].names[class_index]46    confidence = results[0].probs.top1conf.item()47 48    # Annotate the image with the prediction49    annotated_image = results[0].plot()50 51    return annotated_image, f"Detected: {class_name}, Confidence: {confidence:.2f}", class_name52 53 54# Function for the second step: Disease selection55def disease_details(selected_disease):56    """Return details, causes, and treatment based on user selection."""57    info = disease_info.get(selected_disease, {58        "description": "No information available.",59        "causes": "N/A",60        "treatment": "N/A"61    })62    return (63        f"Description: {info['description']}\n\n"64        f"Causes: {info['causes']}\n\n"65        f"Treatment: {info['treatment']}"66        )67 68# Gradio UI components for kidney classification69with gr.Blocks() as interface:70    gr.Markdown("<h1 style='text-align: center; color: #4CAF50;'>🩺 Kidney Image Classification System</h1>")71    72    # Step 1: Image Upload and Prediction73    with gr.Row():74        image_input = gr.Image(type="pil", label="📂 Upload MRI Image")75        submit_btn1 = gr.Button("🔍 Analyze Image")76    77    # Step 1 Outputs78    with gr.Row():79        output_image = gr.Image(label="🧠 Annotated MRI Image")80        output_text = gr.Textbox(label="🔬 Prediction Info")81 82    # Step 2: Disease Selection based on Kidney Prediction83    gr.Markdown("<h2 style='text-align: center;'>Get More Information for Kidney Disease</h2>")84    with gr.Row():85        disease_dropdown = gr.Radio(label="🩺 Select Detected Kidney Disease for More Info", choices=["Cyst", "Normal", "Stone", "Tumor"])86        submit_btn2 = gr.Button("Get Kidney Disease Details")87    88    # Step 2 Output for Kidney89    with gr.Row():90        disease_info_output = gr.Textbox(label="💡 Kidney Disease Info")91 92    # Button functionalities93    submit_btn1.click(fn=predict, inputs=image_input, outputs=[output_image, output_text, disease_dropdown])94    submit_btn2.click(fn=disease_details, inputs=disease_dropdown, outputs=disease_info_output)95 96# Launch the interface97interface.launch(share=True)98