Update app.py

app.py

@@ -5,41 +5,27 @@ import gradio as gr
 import numpy as np
 import tensorflow as tf
 from PIL import Image
-from sklearn.preprocessing import StandardScaler  # Required for feature scaling
-import joblib  # To load the scaler
+from sklearn.preprocessing import StandardScaler
+import joblib
+from transformers import ViTForImageClassification, ViTImageProcessor
 
-# Set device for ViT model (PyTorch)
+# Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 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
+# Load ViT model from Hugging Face
+vit_model = ViTForImageClassification.from_pretrained("andromeda01111/ViT_BCC").to(device)
+vit_processor = ViTImageProcessor.from_pretrained("andromeda01111/ViT_BCC")
 
-# Load ViT model weights
-vit_model_path = "andromeda01111/ViT_BCC"
-vit_model.load_state_dict(torch.load(vit_model_path, map_location=device))
-vit_model.to(device)
-vit_model.eval()
+# Load Neural Network model from Hugging Face
+nn_model = tf.keras.models.load_model("andromeda01111/NN_BC")
 
-# 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])
-])
+# Load scaler (ensure it's uploaded in the Hugging Face repo)
+scaler = joblib.load("scaler.pkl")
 
 # Class labels
 class_names = ["Benign", "Malignant"]
 
-# Load trained Neural Network model (TensorFlow/Keras)
-nn_model_path = "andromeda01111/NN_BC/my_NN_BC_model.keras"  # Ensure the correct path
-nn_model = tf.keras.models.load_model(nn_model_path)
-
-# Load scaler for feature normalization
-scaler_path = "andromeda01111/NN_BC/nn_bc_scaler.pkl"  # Update path
-scaler = joblib.load(scaler_path)  # Load pre-fitted scaler
-
-# Define feature names for NN model
+# Define 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",
@@ -50,16 +36,15 @@ 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 for ViT classification."
-        image = image.convert("RGB")  # Ensure RGB format
-        input_tensor = transform(image).unsqueeze(0).to(device)  # Preprocess image
+            return "Please upload an image."
+        image = image.convert("RGB")
+        inputs = vit_processor(images=image, return_tensors="pt").to(device)
 
         with torch.no_grad():
-            output = vit_model(input_tensor)
-            predicted_class = torch.argmax(output, dim=1).item()
+            outputs = vit_model(**inputs)
+            predicted_class = torch.argmax(outputs.logits, dim=1).item()
 
         return class_names[predicted_class]
 
@@ -67,33 +52,24 @@ def classify(model_choice, image=None, *features):
         if any(f is None for f in features):
             return "Please enter all 30 numerical features."
 
-        # Convert input features to NumPy array
         input_data = np.array(features).reshape(1, -1)
-
-        # Standardize using pre-trained scaler
         input_data_std = scaler.transform(input_data)
-
-        # Run prediction using TensorFlow model
         prediction = nn_model.predict(input_data_std)
         predicted_class = np.argmax(prediction)
 
         return class_names[predicted_class]
 
-# Define Gradio UI components
+# Gradio UI
 model_selector = gr.Radio(["ViT", "Neural Network"], label="Choose Model")
 image_input = gr.Image(type="pil", label="Upload Mammogram Image")
+feature_inputs = [gr.Number(label=feature) for feature in feature_names]
 
-# Feature inputs labeled correctly
-feature_inputs = [gr.Number(label=feature_names[i]) for i in range(30)]
-
-# Gradio Interface
 iface = gr.Interface(
     fn=classify,
-    inputs=[model_selector, image_input] + feature_inputs,  # Image + Feature inputs
+    inputs=[model_selector, image_input] + feature_inputs,
     outputs="text",
     title="Breast Cancer Classification",
-    description="Choose between ViT (image-based) and Neural Network (feature-based) classification."
+    description="Choose ViT (image-based) or Neural Network (feature-based) classification."
 )
 
-# launch app
 iface.launch()