Update app.py

app.py

@@ -17,7 +17,7 @@ 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
-vit_model_path = "vit_bc.pth"  
+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)
@@ -34,11 +34,11 @@ 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 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
@@ -85,35 +85,26 @@ def classify(model_choice, image=None, *features):
 
         return class_names[predicted_class]
 
-# Gradio UI
-model_selector = gr.Radio(["ViT", "Neural Network"], label="Choose Model")
-image_input = gr.Image(type="pil", label="Upload Image")
-
-# feature_inputs = [gr.Number(label=feature, scale=0.5) for feature in feature_names]
-num_columns = 3  # Change to 4 for a 4-column layout
-
-feature_inputs = []
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# Breast Cancer Classification")
     gr.Markdown("Choose between ViT (image-based) and Neural Network (feature-based) classification.")
 
     model_selector = gr.Radio(["ViT", "Neural Network"], label="Choose Model")
-
     image_input = gr.Image(type="pil", label="Upload Image")
 
     # Arrange feature inputs in a matrix layout (3 columns)
-    num_columns = 3  # Change to 4 for a 4-column layout
+    num_columns = 3
     feature_inputs = []
 
     with gr.Row():
-        for i in range(0, len(feature_names), num_columns):
-            with gr.Column():
-                for feature in feature_names[i:i+num_columns]:
-                    feature_inputs.append(gr.Number(label=feature, scale=1))
+        columns = [gr.Column() for _ in range(num_columns)]
+        for i, feature in enumerate(feature_names):
+            with columns[i % num_columns]:
+                feature_inputs.append(gr.Number(label=feature, scale=1))
 
-    # Example buttons
     def fill_example(example):
+        """Pre-fills example inputs."""
         return {feature_inputs[i]: example[i] for i in range(len(feature_inputs))}
 
     examples = [
@@ -127,6 +118,6 @@ with gr.Blocks() as demo:
         outputs="text",
         examples=examples,
         live=True
-    )
+    ).render()
 
 demo.launch()