Update app.py

app.py

@@ -3,14 +3,16 @@ import torchvision.transforms as transforms
 import torchvision.models as models
 import torch.nn as nn
 from joblib import load
-from gradio import File
 from PIL import Image
-import gradio as gr
 import matplotlib.pyplot as plt
 import io
+import numpy as np
+import gradio as gr
 
-# Transformation and device setup
+# Device configuration
 device = torch.device("cpu")
+
+# Transformation for the input images
 data_transforms = transforms.Compose([
     transforms.Resize(224),
     transforms.CenterCrop(224),
@@ -19,57 +21,77 @@ data_transforms = transforms.Compose([
 ])
 
 # Load the Isolation Forest model
-clf = load('Models/Anomaly_MSI_MSS_Isolation_Forest_model.joblib')
-
-# Load feature extractor
-feature_extractor_path = 'Models/feature_extractor.pth'
-feature_extractor = models.resnet50(weights=None)
-feature_extractor.fc = nn.Sequential()
-feature_extractor.load_state_dict(torch.load(feature_extractor_path, map_location=device))
-feature_extractor.to(device)
-feature_extractor.eval()
-
-# Load gastric classification model
-GASTRIC_MODEL_PATH = 'Gastric_Models/the_resnet_50_model.pth'
-model_ft = torch.load(GASTRIC_MODEL_PATH, map_location=device)
-model_ft.to(device)
-model_ft.eval()
+def load_isolation_forest():
+    path = 'Models/Anomaly_MSI_MSS_Isolation_Forest_model.joblib'
+    return load(path)
 
-# Anomaly detection function
-def is_anomaly(clf, feature_extractor, input_image):
+# Load the feature extractor
+def load_feature_extractor():
+    feature_extractor_path = 'Models/feature_extractor.pth'
+    feature_extractor = models.resnet50(weights=None)
+    feature_extractor.fc = nn.Sequential()
+    feature_extractor.load_state_dict(torch.load(feature_extractor_path, map_location=device))
     feature_extractor.to(device)
-    with torch.no_grad():
-        image_features = feature_extractor(input_image)
-
-    is_outlier = clf.predict(image_features.cpu().numpy().reshape(1, -1))
-    return is_outlier[0] == -1
+    feature_extractor.eval()
+    return feature_extractor
 
-# Anomaly detection and classification function
-def classify_image(uploaded_image):
-    image = Image.open(uploaded_image).convert('RGB')
-    input_image = data_transforms(image).unsqueeze(0).to(device)
-
-    # Anomaly detection
-    if is_anomaly(clf, feature_extractor, input_image):
-        return "Anomaly detected. Image will not be classified.", None
+# Anomaly detection function
+def is_anomaly(clf, feature_extractor, image):
+    with torch.no_grad():
+        image_features = feature_extractor(image)
+    return clf.predict(image_features.cpu().numpy().reshape(1, -1))[0] == -1
 
-    # Classification
+# Classification function
+def classify_image(model, image):
     with torch.no_grad():
-        outputs = model_ft(input_image)
+        outputs = model(image)
         probabilities = torch.nn.functional.softmax(outputs, dim=1)
         _, predicted = torch.max(outputs, 1)
 
-    predicted_class_index = predicted.item()
     class_names = ['abnormal', 'normal']
+    predicted_class_index = predicted.item()
     predicted_class_name = class_names[predicted_class_index]
     predicted_probability = probabilities[0][predicted_class_index].item() * 100
 
-    return f"Class: {predicted_class_name}, Probability: {predicted_probability:.2f}%", None
+    return predicted_class_name, predicted_probability
+
+# Load the classification model
+def load_classification_model():
+    model_path = 'Gastric_Models/the_resnet_50_model.pth'
+    model = torch.load(model_path, map_location=device)
+    model.to(device)
+    model.eval()
+    return model
+
+# Function to process the image and get results
+def process_image(image_path):
+    # Convert to PIL and apply transforms
+    image = Image.open(io.BytesIO(image_path.read())).convert('RGB')
+    input_image = data_transforms(image).unsqueeze(0).to(device)
+
+    # Load models
+    clf = load_isolation_forest()
+    feature_extractor = load_feature_extractor()
+    classification_model = load_classification_model()
+
+    # Check for anomaly
+    if is_anomaly(clf, feature_extractor, input_image):
+        return "Anomaly detected. Image will not be classified.", None, None
+
+    # Classify image
+    predicted_class, probability = classify_image(classification_model, input_image)
+    result = f"The predicted class is: {predicted_class} with a probability of {probability:.2f}%"
+
+    # Further processing for heatmap or additional features can be added here
+
+    return result, None, None  # Returning placeholders for additional outputs if needed
 
+# Gradio interface
+    
 iface = gr.Interface(
-    fn=classify_image,
+    fn=process_image,
     inputs=File(type="filepath"),
-    outputs=gr.Image(),
+    outputs=[gr.Textbox(label="Result"), gr.Image(label="Heatmap"), gr.Image(label="Additional Output")],
     title="GastroHub AI Gastric Image Classifier",
     description="Upload an image to classify it as normal or abnormal.",
     article="Above is a sample image to test the results of the model. Click it to see the results.",
@@ -79,5 +101,4 @@ iface = gr.Interface(
     allow_flagging="never",
 )
 
-# Run the Gradio app
 iface.launch()