Update app.py

app.py

@@ -7,54 +7,36 @@ Original file is located at
     https://colab.research.google.com/drive/1sjyLFLqBccpUzaUi4eyyP3NYE3gDtHfs
 """
 
-import gradio as gr
-from fastai.vision.all import load_learner
-from PIL import Image
+import streamlit as st
+from streamlit_option_menu import option_menu
+import tensorflow as tf
+from tensorflow import keras
+import numpy as np
+import pandas as pd
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torchvision import transforms
+from PIL import Image
+import os
+import time
+from efficientnet_pytorch import EfficientNet
+from fastai.vision.all import load_learner
 
-# Define the custom model class (MelanomaModel in this case)
-class MelanomaModel(nn.Module):
-    def __init__(self):
-        super(MelanomaModel, self).__init__()
-        # Define your model layers here (example)
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)
-        self.fc1 = nn.Linear(64*224*224, 10)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = x.view(x.size(0), -1)
-        x = self.fc1(x)
-        return x
-
-# Model paths for all disease types
-model_path_skin_disease = 'multi_weight.pth'  # Skin Disease Model
-
-model_path_alzheimers = 'alzheimers_model.pkl'
-model_path_eye_disease = 'eye_disease_model.pkl'
-
-# Load models with CPU mapping
-skin_disease_model = torch.load(model_path_skin_disease, map_location=torch.device('cpu'))  # For Skin Disease model
-
-alzheimers_model = load_learner(model_path_alzheimers, cpu=True)
-eye_disease_model = load_learner(model_path_eye_disease, cpu=True)
+# Set up environment variables for GPU handling
+os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
 
-# Diagnosis Map for Skin Disease Model
-DIAGNOSIS_MAP = {
-    0: 'Melanoma',
-    1: 'Melanocytic nevus',
-    2: 'Basal cell carcinoma',
-    3: 'Actinic keratosis',
-    4: 'Benign keratosis',
-    5: 'Dermatofibroma',
-    6: 'Vascular lesion',
-    7: 'Squamous cell carcinoma',
-    8: 'Unknown'
-}
+# Cache to load models efficiently
+@st.cache_resource
+def load_skin_model():
+    model = MelanomaModel(out_size=9)
+    model_path = "multi_weight.pth"
+    checkpoint = torch.load(model_path, map_location=torch.device('cpu'), weights_only=False)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+    return model
 
-# Image Preprocessing for Skin Disease Model
+# Transform for skin lesion images
 transform = transforms.Compose([
     transforms.Resize(256),
     transforms.CenterCrop(224),
@@ -62,82 +44,141 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-# Skin Disease Prediction Function
-def predict_skin_disease(img: Image.Image):
-    img_tensor = transform(img).unsqueeze(0)
-    with torch.no_grad():
-        outputs = skin_disease_model(img_tensor)
-        probs = F.softmax(outputs, dim=1)
-        top_probs, top_idxs = torch.topk(probs, 3, dim=1)  # top 3 predictions
-
-    predictions = []
-    for prob, idx in zip(top_probs[0], top_idxs[0]):
-        label = DIAGNOSIS_MAP.get(idx.item(), "Unknown")
-        confidence = prob.item() * 100
-        predictions.append(f"{label}: {confidence:.2f}%")
-
-    return "\n".join(predictions)
-
-# Brain Tumor Prediction Function
-def predict_brain_tumor(image):
-    pred, _, prob = brain_tumor_model.predict(image)
-    return f"Prediction: {pred}, Probability: {prob.max():.2f}"
+# Diagnosis map for skin disease model
+DIAGNOSIS_MAP = {
+    0: 'Melanoma', 1: 'Melanocytic nevus', 2: 'Basal cell carcinoma', 3: 'Actinic keratosis', 
+    4: 'Benign keratosis', 5: 'Dermatofibroma', 6: 'Vascular lesion', 7: 'Squamous cell carcinoma', 8: 'Unknown'
+}
 
-# Alzheimer's Prediction Function
-def predict_alzheimers(image):
-    pred, _, prob = alzheimers_model.predict(image)
-    return f"Prediction: {pred}, Probability: {prob.max():.2f}"
+# Model for skin lesion classification
+class MelanomaModel(nn.Module):
+    def __init__(self, out_size, dropout_prob=0.5):
+        super(MelanomaModel, self).__init__()
+        self.efficient_net = EfficientNet.from_pretrained('efficientnet-b0')
+        self.efficient_net._fc = nn.Identity()
+        self.fc1 = nn.Linear(1280, 512)
+        self.fc2 = nn.Linear(512, 256)
+        self.fc3 = nn.Linear(256, out_size)
+        self.dropout = nn.Dropout(dropout_prob)
 
-# Eye Disease Prediction Function
-def predict_eye_disease(image):
-    pred, _, prob = eye_disease_model.predict(image)
-    return f"Prediction: {pred}, Probability: {prob.max():.2f}"
+    def forward(self, x):
+        x = self.efficient_net(x)
+        x = x.view(x.size(0), -1)
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = F.relu(self.fc2(x))
+        x = self.dropout(x)
+        x = self.fc3(x)
+        return x
 
-# Gradio Interface Function
+# Alzheimer's Prediction
+@st.cache_resource
+def load_alzheimer_model():
+    return keras.models.load_model('alzheimer_99.5.h5')
+
+# Brain Tumor Prediction
+@st.cache_resource
+def load_brain_tumor_model(classes):
+    if classes == '44 Classes':
+        return keras.models.load_model('44class_96.5.h5')
+    elif classes == '17 Classes':
+        return keras.models.load_model('17class_98.1.h5')
+    elif classes == '15 Classes':
+        return keras.models.load_model('15class_99.8.h5')
+    else:  # Default 2 classes
+        return keras.models.load_model('2calss_lagre_dataset_99.1.h5')
+
+# Prediction for Skin Disease
+def predict_skin_lesion(img: Image.Image, model: nn.Module):
+    img_tensor = transform(img).unsqueeze(0)
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probs = F.softmax(outputs, dim=1)
+        top_probs, top_idxs = torch.topk(probs, 3, dim=1)
+    predictions = [(DIAGNOSIS_MAP.get(idx.item(), "Unknown"), prob.item() * 100) for prob, idx in zip(top_probs[0], top_idxs[0])]
+    return predictions
+
+# Prediction for Brain Tumor and Alzheimer
+def predict(img_path, model, result_classes):
+    img = tf.keras.utils.load_img(img_path, target_size=(224, 224))
+    img_array = np.array(img).reshape(-1, 224, 224, 3)
+    pred = model.predict(img_array)
+    return result_classes[np.argmax(pred, axis=1)[0]]
+
+# Sidebar for Disease Categories
+def spr_sidebar():
+    menu = option_menu(
+        menu_title="Navigation",
+        options=["Brain Tumor", "Alzheimer", "Skin Disease", "Eye Disease", "About"],
+        icons=["house", "brain", "microscope", "eye", "info-square"],
+        menu_icon="cast",
+        default_index=0,
+        orientation="horizontal"
+    )
+    return menu
+
+# Home Page Content
+def home_page():
+    st.title("Disease Detection Web App")
+    
+    uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+    
+    if uploaded_file is not None:
+        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+
+        if st.button("Classify"):
+            if selected_category == "Brain Tumor":
+                classes = st.selectbox("Select Number of Classes", ['44 Classes', '17 Classes', '15 Classes', '2 Classes'])
+                model = load_brain_tumor_model(classes)
+                result_classes = ['Astrocitoma', 'Carcinoma', 'Ependimoma', '_NORMAL', 'etc...']  # Define all the classes
+                result = predict(uploaded_file, model, result_classes)
+                st.success(f"Prediction: {result}")
+
+            elif selected_category == "Alzheimer":
+                model = load_alzheimer_model()
+                result_classes = ['Mild_Demented', 'Moderate_Demented', 'Non_Demented']
+                result = predict(uploaded_file, model, result_classes)
+                st.success(f"Prediction: {result}")
+
+            elif selected_category == "Skin Disease":
+                model = load_skin_model()
+                img = Image.open(uploaded_file)
+                predictions = predict_skin_lesion(img, model)
+                for idx, (label, confidence) in enumerate(predictions, 1):
+                    st.write(f"{idx}. {label}: {confidence:.2f}%")
+
+            elif selected_category == "Eye Disease":
+                # Implement Eye Disease prediction (similar to others)
+                pass
+
+# About Page Content
+def about_page():
+    st.header('About the Project')
+    st.write("""
+        This web app detects different diseases using machine learning models.
+        - **Brain Tumor Detection**: Classifies different types of brain tumors.
+        - **Alzheimer Detection**: Classifies stages of Alzheimer.
+        - **Skin Disease Classification**: Classifies types of skin lesions.
+        - **Eye Disease Classification**: TBD.
+    """)
+
+# Main Function to Run the App
 def main():
-    # Image input component
-    image_input = gr.inputs.Image(shape=(224, 224), image_mode='RGB')
-
-    # Dropdown to choose disease type
-    model_choice = gr.inputs.Dropdown(choices=[
-        "Skin Disease", "Brain Tumor", "Alzheimer's Detection", "Eye Disease"],
-        label="Select Disease Type")
-
-    # Gradio tabs for each category
-    with gr.Blocks() as demo:
-        gr.Markdown("# Medical Image Classifier Dashboard")
-
-        with gr.Tab("Skin Disease Prediction"):
-            with gr.Column():
-                gr.Markdown("Upload a skin lesion image for diagnosis prediction.")
-                image_input_skin = gr.Image(type="pil", label="Upload Skin Lesion Image")
-                output_skin = gr.Textbox(label="Prediction Results")
-                image_input_skin.change(predict_skin_disease, inputs=image_input_skin, outputs=output_skin)
-
-        with gr.Tab("Brain Tumor Prediction"):
-            with gr.Column():
-                gr.Markdown("Upload a brain scan image for tumor classification.")
-                image_input_brain = gr.Image(type="pil", label="Upload Brain Scan Image")
-                output_brain = gr.Textbox(label="Prediction Results")
-                image_input_brain.change(predict_brain_tumor, inputs=image_input_brain, outputs=output_brain)
-
-        with gr.Tab("Alzheimer's Prediction"):
-            with gr.Column():
-                gr.Markdown("Upload a brain image for Alzheimer's detection.")
-                image_input_alz = gr.Image(type="pil", label="Upload Alzheimer's Image")
-                output_alz = gr.Textbox(label="Prediction Results")
-                image_input_alz.change(predict_alzheimers, inputs=image_input_alz, outputs=output_alz)
-
-        with gr.Tab("Eye Disease Prediction"):
-            with gr.Column():
-                gr.Markdown("Upload an image for eye disease classification.")
-                image_input_eye = gr.Image(type="pil", label="Upload Eye Disease Image")
-                output_eye = gr.Textbox(label="Prediction Results")
-                image_input_eye.change(predict_eye_disease, inputs=image_input_eye, outputs=output_eye)
-
-    demo.launch()
-
-# Run the Gradio app
-if __name__ == "__main__":
+    selected_category = spr_sidebar()
+    
+    if selected_category == "Brain Tumor":
+        home_page()
+    elif selected_category == "Alzheimer":
+        home_page()
+    elif selected_category == "Skin Disease":
+        home_page()
+    elif selected_category == "Eye Disease":
+        home_page()  # Eye Disease detection page
+    elif selected_category == "About":
+        about_page()
+
+if __name__ == '__main__':
     main()
 
+
+