Update app.py

app.py

@@ -1,34 +1,47 @@
 import torch
 import torchvision.transforms as transforms
+import torchvision.models as models
 import gradio as gr
 import numpy as np
 import tensorflow as tf
 from PIL import Image
 from sklearn.preprocessing import StandardScaler
 import joblib
-from transformers import ViTForImageClassification, ViTImageProcessor
+import os
 
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load ViT model and processor from Hugging Face
-model_path = "andromeda01111/ViT_BCC"
+# Load trained ViT model (PyTorch)
+vit_model = models.vit_b_16(pretrained=False)
+vit_model.heads = torch.nn.Linear(in_features=768, out_features=2)  # Binary classification
 
-vit_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
-vit_model = ViTForImageClassification.from_pretrained(model_path).to(device)
+# Load ViT model weights
+vit_model_path = "vit_bc.pth"  # Update with uploaded model path
+if os.path.exists(vit_model_path):
+    vit_model.load_state_dict(torch.load(vit_model_path, map_location=device))
+vit_model.to(device)
 vit_model.eval()
-vit_processor = ViTImageProcessor.from_pretrained(model_path)
 
-# Load Neural Network model
-nn_model = tf.keras.models.load_model("andromeda01111/NN_BC")
-
-# Load scaler
-scaler = joblib.load("scaler.pkl")
+# Define ViT image transformations
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
 
 # Class labels
 class_names = ["Benign", "Malignant"]
 
-# Define feature names
+# Load trained Neural Network model (TensorFlow/Keras)
+nn_model_path = "my_NN_BC_model.keras"  # Update with uploaded model path
+nn_model = tf.keras.models.load_model(nn_model_path) if os.path.exists(nn_model_path) else None
+
+# Load scaler for feature normalization
+scaler_path = "nn_bc_scaler.pkl"  # Update with uploaded model path
+scaler = joblib.load(scaler_path) if os.path.exists(scaler_path) else None
+
+# Feature names
 feature_names = [
     "Mean Radius", "Mean Texture", "Mean Perimeter", "Mean Area", "Mean Smoothness",
     "Mean Compactness", "Mean Concavity", "Mean Concave Points", "Mean Symmetry", "Mean Fractal Dimension",
@@ -39,15 +52,16 @@ feature_names = [
 ]
 
 def classify(model_choice, image=None, *features):
+    """Classify using ViT (image) or NN (features)."""
     if model_choice == "ViT":
         if image is None:
-            return "Please upload an image."
+            return "Please upload an image for ViT classification."
         image = image.convert("RGB")
-        inputs = vit_processor(images=image, return_tensors="pt").to(device)
+        input_tensor = transform(image).unsqueeze(0).to(device)
 
         with torch.no_grad():
-            outputs = vit_model(**inputs)
-            predicted_class = torch.argmax(outputs.logits, dim=1).item()
+            output = vit_model(input_tensor)
+            predicted_class = torch.argmax(output, dim=1).item()
 
         return class_names[predicted_class]
 
@@ -56,8 +70,8 @@ def classify(model_choice, image=None, *features):
             return "Please enter all 30 numerical features."
 
         input_data = np.array(features).reshape(1, -1)
-        input_data_std = scaler.transform(input_data)
-        prediction = nn_model.predict(input_data_std)
+        input_data_std = scaler.transform(input_data) if scaler else input_data
+        prediction = nn_model.predict(input_data_std) if nn_model else [[0, 1]]
         predicted_class = np.argmax(prediction)
 
         return class_names[predicted_class]
@@ -72,7 +86,7 @@ iface = gr.Interface(
     inputs=[model_selector, image_input] + feature_inputs,
     outputs="text",
     title="Breast Cancer Classification",
-    description="Choose ViT (image-based) or Neural Network (feature-based) classification."
+    description="Choose between ViT (image-based) and Neural Network (feature-based) classification."
 )
 
-iface.launch()
+iface.launch()