Update app.py

app.py

@@ -9,11 +9,14 @@ from sklearn.preprocessing import StandardScaler
 import joblib
 import os
 
+# Force CPU usage for TensorFlow to avoid CUDA issues
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
 # 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 = models.vit_b_16(weights=None)
 vit_model.heads = torch.nn.Linear(in_features=768, out_features=2)  # Binary classification
 
 # Load ViT model weights
@@ -34,11 +37,15 @@ transform = transforms.Compose([
 class_names = ["Benign", "Malignant"]
 
 # Load trained Neural Network model (TensorFlow/Keras)
-nn_model_path = "my_NN_BC_model.keras"  # Update with uploaded model path
+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
 
+# Load trained CNN model (TensorFlow/Keras)
+cnn_model_path = "my_CNN_BC_model.keras"
+cnn_model = tf.keras.models.load_model(cnn_model_path) if os.path.exists(cnn_model_path) else None
+
 # Load scaler for feature normalization
-scaler_path = "nn_bc_scaler.pkl"  # Update with uploaded model path
+scaler_path = "nn_bc_scaler.pkl"
 scaler = joblib.load(scaler_path) if os.path.exists(scaler_path) else None
 
 # Feature names
@@ -52,7 +59,7 @@ feature_names = [
 ]
 
 def classify(model_choice, image=None, file=None, *features):
-    """Classify using ViT (image) or NN (features from manual input or file)."""
+    """Classify using ViT (image), NN, or CNN (features from manual input or file)."""
     if model_choice == "ViT":
         if image is None:
             return "Please upload an image for ViT classification."
@@ -65,7 +72,7 @@ def classify(model_choice, image=None, file=None, *features):
 
         return class_names[predicted_class]
 
-    elif model_choice == "Neural Network":
+    elif model_choice in ["Neural Network", "CNN"]:
         if file is not None:
             try:
                 file_data = file.read().decode("utf-8").strip().split(',')
@@ -74,42 +81,34 @@ def classify(model_choice, image=None, file=None, *features):
                 return f"Error reading file: {e}"
         
         if len(features) != len(feature_names):
-            return "Please provide all 30 numerical features."
+            return "Incorrect number of features provided. Please provide exactly 30 numerical values."
 
         input_data = np.array(features).reshape(1, -1)
         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]]
+        
+        if model_choice == "Neural Network":
+            prediction = nn_model.predict(input_data_std) if nn_model else [[0, 1]]
+        else:  # CNN
+            prediction = cnn_model.predict(input_data_std) if cnn_model else [[0, 1]]
+        
         predicted_class = np.argmax(prediction)
-
         return class_names[predicted_class]
 
+# Gradio UI
 with gr.Blocks() as iface:
     gr.Markdown("# Breast Cancer Classification")
-    gr.Markdown("Choose between ViT (image-based) and Neural Network (feature-based) classification.")
-
-    with gr.Row():
-        model_selector = gr.Radio(["ViT", "Neural Network"], label="Choose Model", interactive=True)
-
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(type="pil", label="Upload Image")
-            file_input = gr.File(label="Upload Text File (for NN model)")
-
-    with gr.Row():
-        gr.Markdown("### Enter Features (For Neural Network Model)")
+    model_selector = gr.Radio(["ViT", "Neural Network", "CNN"], label="Choose Model")
     
-    num_cols = 3
-    feature_inputs = [gr.Number(label=feature) for feature in feature_names]
-    feature_inputs_split = [feature_inputs[i::num_cols] for i in range(num_cols)]
-
     with gr.Row():
-        for col in feature_inputs_split:
-            with gr.Column():
-                for feature in col:
-                    feature.render()
-
-    output = gr.Textbox(label="Prediction Result")
-    classify_button = gr.Button("Classify")
-    classify_button.click(classify, inputs=[model_selector, image_input, file_input] + feature_inputs, outputs=output)
+        image_input = gr.Image(type="pil", label="Upload Mammogram Image")
+        file_input = gr.File(label="Upload Text File (for NN/CNN model)")
+    
+    with gr.Column():
+        feature_inputs = [gr.Number(label=feature) for feature in feature_names]
+    
+    submit_button = gr.Button("Classify")
+    output_text = gr.Textbox(label="Prediction Result")
+    
+    submit_button.click(fn=classify, inputs=[model_selector, image_input, file_input] + feature_inputs, outputs=output_text)
 
 iface.launch()