Update app.py

app.py

@@ -1,37 +1,40 @@
 import torch
-from torchvision import models, transforms
 from torch import nn
+from torchvision import models, transforms
 from PIL import Image
 import gradio as gr
 import numpy as np
 
-# Load the pre-trained model (ResNet50)
+# Define model (pretrained ResNet50)
 model = models.resnet50(pretrained=True)
-model.fc = nn.Linear(model.fc.in_features, 2)  # For binary classification (Thyroid: Positive/Negative)
+model.fc = nn.Linear(model.fc.in_features, 3)  # 3 output classes: Normal, Hypothyroidism, Hyperthyroidism
 
-# Define image transformation for preprocessing
+# Define image transformations (resizing and normalization)
 transform = transforms.Compose([
-    transforms.Resize((224, 224)),
+    transforms.Resize((224, 224)),  # Resize image for input
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Function to classify the thyroid condition based on image input
-def classify_thyroid_image(image):
-    # Convert the numpy array (from Gradio) to a Pillow image
-    image = Image.fromarray(image.astype('uint8'), 'RGB')
-    
-    # Apply transformation (resize, normalize, etc.)
+# Example classification function
+def classify_thyroid_condition(image):
+    image = Image.fromarray(image.astype('uint8'), 'RGB')  # Convert numpy array to Pillow Image
     image = transform(image).unsqueeze(0)  # Add batch dimension
-    
-    model.eval()  # Set the model to evaluation mode
 
+    model.eval()  # Set the model to evaluation mode
     with torch.no_grad():
         output = model(image)
         _, predicted = torch.max(output, 1)
 
-    diagnosis = "Thyroid Disease Detected" if predicted.item() == 1 else "No Thyroid Disease"
+    # Map prediction to class labels
+    if predicted.item() == 0:
+        diagnosis = "Normal"
+    elif predicted.item() == 1:
+        diagnosis = "Hypothyroidism"
+    else:
+        diagnosis = "Hyperthyroidism"
+
     return diagnosis
 
 # Create Gradio interface for image input
-gr.Interface(fn=classify_thyroid_image, inputs="image", outputs="text").launch()
+gr.Interface(fn=classify_thyroid_condition, inputs="image", outputs="text").launch()