Update app.py

app.py

@@ -1,4 +1,3 @@
-import os
 import torch
 import torchvision.transforms as transforms
 import torchvision.models as models
@@ -8,43 +7,51 @@ import tensorflow as tf
 from PIL import Image
 from sklearn.preprocessing import StandardScaler
 import joblib
+import os
 
-# Disable TensorFlow GPU to avoid CUDA conflicts
+# 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 ViT model
-vit_model = models.vit_b_16(weights="DEFAULT")
-vit_model.heads = torch.nn.Linear(in_features=768, out_features=2)
+# Load trained ViT model (PyTorch)
+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 (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).eval()
+vit_model.to(device)
+vit_model.eval()
+
+# Define image transformations for ViT
+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 Neural Net model
+# Class labels
+class_names = ["Benign", "Malignant"]
+
+# Load trained Neural Network model (TensorFlow/Keras)
 nn_model_path = "my_NN_BC_model.keras"
-nn_model = None
+
+nn_model = tf.keras.models.load_model(nn_model_path)
+
 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}")
+        print(f"Error loading NN model: {e}")
 
-# Load scaler
+# Load scaler for feature normalization
 scaler_path = "nn_bc_scaler.pkl"
 scaler = joblib.load(scaler_path) if os.path.exists(scaler_path) else None
 
-# Transforms
-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])
-])
-
-# Labels
-class_names = ["Benign", "Malignant"]
-
-# Feature names (30)
+# 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",
@@ -54,88 +61,117 @@ feature_names = [
     "Worst Compactness", "Worst Concavity", "Worst Concave Points", "Worst Symmetry", "Worst Fractal Dimension"
 ]
 
-# Sample Inputs
-benign_example = [9.504, 12.44, 60.34, 273.9, 0.1024, 0.06492, 0.02956, 0.02076, 0.1815, 0.06905, 
-                  0.2773, 0.9768, 1.909, 15.7, 0.009606, 0.01432, 0.01985, 0.01421, 0.02027, 0.002968,
-                  10.23, 15.66, 65.13, 314.9, 0.1324, 0.1148, 0.08867, 0.06227, 0.245, 0.07773]
+# Example inputs
+benign_example = [9.504,12.44,60.34,273.9,0.1024,0.06492,0.02956,0.02076,0.1815,0.06905,0.2773,0.9768,
+                  1.909,15.7,0.009606,0.01432,0.01985,0.01421,0.02027,0.002968,10.23,15.66,65.13,314.9,
+                  0.1324,0.1148,0.08867,0.06227,0.245,0.07773]
 
-malignant_example = [11.42, 20.38, 77.58, 386.1, 0.1425, 0.2839, 0.2414, 0.1052, 0.2597, 0.09744,
-                     0.4956, 1.156, 3.445, 27.23, 0.00911, 0.07458, 0.05661, 0.01867, 0.05963, 0.009208,
-                     14.91, 26.5, 98.87, 567.7, 0.2098, 0.8663, 0.6869, 0.2575, 0.6638, 0.173]
+malignant_example = [11.42,20.38,77.58,386.1,0.1425,0.2839,0.2414,0.1052,0.2597,0.09744,0.4956,1.156,
+                     3.445,27.23,0.00911,0.07458,0.05661,0.01867,0.05963,0.009208,14.91,26.5,98.87,567.7,
+                     0.2098,0.8663,0.6869,0.2575,0.6638,0.173]
 
-
-# --- Classification Function ---
 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")
         input_tensor = transform(image).unsqueeze(0).to(device)
+
         with torch.no_grad():
             output = vit_model(input_tensor)
-        pred_class = torch.argmax(output, dim=1).item()
-        return class_names[pred_class]
+            predicted_class = torch.argmax(output, dim=1).item()
+
+        return class_names[predicted_class]
 
     elif model_choice == "Neural Network":
         if any(f is None for f in features):
-            return "❌ All 30 numerical features are required."
+            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
         prediction = nn_model.predict(input_data_std) if nn_model else [[0, 1]]
-        pred_class = np.argmax(prediction)
-        return class_names[pred_class]
+        predicted_class = np.argmax(prediction)
 
-# --- File parser for NN features ---
-def extract_features_from_file(file):
-    try:
-        content = file.read().decode("utf-8").strip()
-        values = [float(x) for x in content.replace(",", " ").split()]
-        if len(values) != 30:
-            raise ValueError("Expected 30 values, got {}".format(len(values)))
-        return values  # must return a list of 30 floats
-    except Exception as e:
-        print(f"Error reading file: {e}")
-        return [0.0] * 30  # return empty values for safety
+        return class_names[predicted_class]
 
-# --- UI with Gradio ---
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Breast Cancer Classifier (ViT & Neural Net)")
+    gr.Markdown("## 🩺 Breast Cancer Classification Model")
+    gr.Markdown("Select a model and provide input data to classify breast cancer as **Benign** or **Malignant**.")
 
-    model_selector = gr.Radio(["ViT", "Neural Network"], label="Choose Model", value="ViT")
+    with gr.Row():
+        model_selector = gr.Radio(["ViT", "Neural Network"], label="🔬 Choose Model", value="ViT")
 
-    with gr.Column(visible=True) as vit_section:
-        image_input = gr.Image(label="Upload Image", type="pil")
-        # vit_example = gr.Examples(
-        #     examples=["images/benign (1)_aug_0.png", "images/malignant (1)_aug_0.png"],  # Add scan images
-        #     inputs=[image_input],
-        # )
+    image_input = gr.Image(type="pil", label="📷 Upload Image (for ViT)", visible=True)
 
-    with gr.Column(visible=False) as nn_section:
-        file_input = gr.File(label="📂 Upload 30-feature TXT file", file_types=[".txt"])
-        feature_inputs = [gr.Number(label=name, visible=True, elem_id=f"feature_{i}") for i, name in enumerate(feature_names)]
+    feature_inputs = [gr.Number(label=feature) for feature in feature_names]
 
-        with gr.Row():
-            nn_example_btn_1 = gr.Button("🔴 Malignant Example")
-            nn_example_btn_2 = gr.Button("🔵 Benign Example")
+    # Organizing feature inputs into rows of 3 columns
+    with gr.Row():
+        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)))])
 
-    output = gr.Textbox(label="Prediction")
+    # Example buttons
+    def fill_example(example):
+        """Pre-fills example inputs."""
+        return {feature_inputs[i]: example[i] for i in range(len(feature_inputs))}
 
-    classify_btn = gr.Button("🚀 Classify")
+    with gr.Row():
+        example_btn_1 = gr.Button("🔴 Malignant Example")
+        example_btn_2 = gr.Button("🔵 Benign Example")
 
-    # Event logic
+    output_text = gr.Textbox(label="🔍 Model Prediction", interactive=False)
+
+    def extract_features_from_file(file):
+        """Reads a text file and extracts numerical features."""
+        if file is None:
+            return "❌ Please upload a valid feature file."
+    
+        try:
+            # Read and process file contents
+            content = file.read().decode("utf-8").strip()
+            values = [float(x) for x in content.replace(",", " ").split()]
+        
+            # Check if we have exactly 30 features
+            if len(values) != 30:
+                return "❌ The file must contain exactly 30 numerical values."
+
+            return {feature_inputs[i]: values[i] for i in range(30)}
+
+        except Exception as e:
+            return f"❌ Error processing file: {e}"
+
+    # Add file upload component
+    file_input = gr.File(label="📂 Upload Feature File (for NN)", type="binary", visible=False)
+
+    # Update UI logic to show file input for NN model
     def toggle_inputs(choice):
-        return (
-            gr.update(visible=choice == "ViT"),
-            gr.update(visible=choice == "Neural Network"),
-        )
+        image_visibility = choice == "ViT"
+        feature_visibility = choice == "Neural Network"
+        file_visibility = choice == "Neural Network"
+        return [gr.update(visible=image_visibility)] + [gr.update(visible=feature_visibility)] * len(feature_inputs) + [gr.update(visible=file_visibility)]
 
-    model_selector.change(toggle_inputs, model_selector, [vit_section, nn_section])
+    model_selector.change(toggle_inputs, model_selector, [image_input, *feature_inputs, file_input])
+
+    # Process uploaded file and populate feature fields
+    file_input.change(extract_features_from_file, file_input, feature_inputs)
+
+
+    # Toggle input fields based on model selection
+    """Toggle visibility of inputs based on model selection."""
+    def toggle_inputs(choice):
+        image_visibility = choice == "ViT"
+        feature_visibility = choice == "Neural Network"
+        return [gr.update(visible=image_visibility)] + [gr.update(visible=feature_visibility)] * len(feature_inputs)
 
-    nn_example_btn_1.click(lambda: malignant_example, None, feature_inputs)
-    nn_example_btn_2.click(lambda: benign_example, None, feature_inputs)
+    model_selector.change(toggle_inputs, model_selector, [image_input, *feature_inputs])
 
-    file_input.change(extract_features_from_file, file_input, outputs=feature_inputs)
+    example_btn_1.click(lambda: fill_example(benign_example), None, feature_inputs)
+    example_btn_2.click(lambda: fill_example(malignant_example), None, feature_inputs)
 
-    classify_btn.click(classify, inputs=[model_selector, image_input] + feature_inputs, outputs=output)
+    classify_button = gr.Button("🚀 Classify")
+    classify_button.click(classify, [model_selector, image_input] + feature_inputs, output_text)
 
 demo.launch()