Update app.py

app.py

@@ -9,14 +9,11 @@ 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(weights=None)
+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 weights
@@ -37,15 +34,11 @@ transform = transforms.Compose([
 class_names = ["Benign", "Malignant"]
 
 # Load trained Neural Network model (TensorFlow/Keras)
-nn_model_path = "my_NN_BC_model.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 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"
+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
@@ -58,8 +51,17 @@ feature_names = [
     "Worst Compactness", "Worst Concavity", "Worst Concave Points", "Worst Symmetry", "Worst Fractal Dimension"
 ]
 
-def classify(model_choice, image=None, file=None, *features):
-    """Classify using ViT (image), NN, or CNN (features from manual input or file)."""
+# Example inputs
+benign_example = [13.54, 14.36, 87.46, 566.3, 0.09779, 0.08129, 0.06664, 0.04781, 0.1885, 0.05766,
+                  0.2699, 0.7886, 2.058, 23.56, 0.008462, 0.0146, 0.02387, 0.01315, 0.0198, 0.0023,
+                  15.11, 19.26, 99.7, 711.2, 0.144, 0.1773, 0.239, 0.1288, 0.2977, 0.07259]
+
+malignant_example = [17.99, 10.38, 122.8, 1001.0, 0.1184, 0.2776, 0.3001, 0.1471, 0.2419, 0.07871,
+                     1.095, 0.9053, 8.589, 153.4, 0.006399, 0.04904, 0.05373, 0.01587, 0.03003, 0.006193,
+                     25.38, 17.33, 184.6, 2019.0, 0.1622, 0.6656, 0.7119, 0.2654, 0.4601, 0.1189]
+
+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."
@@ -72,43 +74,37 @@ def classify(model_choice, image=None, file=None, *features):
 
         return class_names[predicted_class]
 
-    elif model_choice in ["Neural Network", "CNN"]:
-        if file is not None:
-            try:
-                file_data = file.read().decode("utf-8").strip().split(',')
-                features = [float(x) for x in file_data]
-            except Exception as e:
-                return f"Error reading file: {e}"
-        
-        if len(features) != len(feature_names):
-            return "Incorrect number of features provided. Please provide exactly 30 numerical values."
+    elif model_choice == "Neural Network":
+        if any(f is None for f in features):
+            return "Please enter all 30 numerical features."
 
         input_data = np.array(features).reshape(1, -1)
         input_data_std = scaler.transform(input_data) if scaler else input_data
-        
-        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]]
-        
+        prediction = nn_model.predict(input_data_std) if nn_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")
-    model_selector = gr.Radio(["ViT", "Neural Network", "CNN"], label="Choose Model")
-    
-    with gr.Row():
-        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)
+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]
+
+# Example buttons
+def fill_example(example):
+    return {feature_inputs[i]: example[i] for i in range(len(feature_inputs))}
+
+example_buttons = gr.Radio(["Benign Example", "Malignant Example"], label="Select Example Input")
+def update_example(choice):
+    return fill_example(benign_example if choice == "Benign Example" else malignant_example)
+
+iface = gr.Interface(
+    fn=classify,
+    inputs=[model_selector, image_input, example_buttons] + feature_inputs,
+    outputs="text",
+    title="Breast Cancer Classification",
+    description="Choose between ViT (image-based) and Neural Network (feature-based) classification.",
+    live=True
+)
 
 iface.launch()