Update app.py

app.py

@@ -9,21 +9,24 @@ from sklearn.preprocessing import StandardScaler
 import joblib
 import os
 
-# Set device
+# Disable GPU for TensorFlow to avoid CUDA conflicts
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+# Set PyTorch 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 = models.vit_b_16(weights="DEFAULT")  # Fixed deprecated 'pretrained'
 vit_model.heads = torch.nn.Linear(in_features=768, out_features=2)  # Binary classification
 
-# Load ViT model weights
+# Load ViT model weights (if available)
 vit_model_path = "vit_bc.pth"
 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()
 
-# Define ViT image transformations
+# Define image transformations for ViT
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -35,7 +38,12 @@ class_names = ["Benign", "Malignant"]
 
 # Load trained Neural Network model (TensorFlow/Keras)
 nn_model_path = "my_NN_BC_model.keras"
-nn_model = tf.keras.models.load_model(nn_model_path) if os.path.exists(nn_model_path) else None
+nn_model = None
+if os.path.exists(nn_model_path):
+    try:
+        nn_model = tf.keras.models.load_model(nn_model_path)
+    except Exception as e:
+        print(f"Error loading NN model: {e}")
 
 # Load scaler for feature normalization
 scaler_path = "nn_bc_scaler.pkl"
@@ -97,13 +105,11 @@ with gr.Blocks() as demo:
 
     feature_inputs = [gr.Number(label=feature) for feature in feature_names]
 
-    # Dynamically arrange feature inputs into rows of 3 columns
+    # Organizing feature inputs into rows of 3 columns
     with gr.Row():
-        for i in range(0, len(feature_inputs), 3):
-            with gr.Column():
-                for j in range(3):
-                    if i + j < len(feature_inputs):
-                        feature_inputs[i + j].render()
+        with gr.Column():
+            for i in range(0, len(feature_inputs), 3):
+                gr.Row([feature_inputs[j] for j in range(i, min(i+3, len(feature_inputs)))])
 
     # Example buttons
     def fill_example(example):
@@ -116,7 +122,7 @@ with gr.Blocks() as demo:
 
     output_text = gr.Textbox(label="🔍 Model Prediction", interactive=False)
 
-    # Logic to toggle inputs based on model selection
+    # Toggle input fields based on model selection
     def toggle_inputs(choice):
         return gr.update(visible=(choice == "ViT")), gr.update(visible=(choice == "Neural Network"))