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"  # Update with uploaded model path
+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)
@@ -87,35 +87,46 @@ def classify(model_choice, image=None, *features):
 
 # Gradio UI
 model_selector = gr.Radio(["ViT", "Neural Network"], label="Choose Model")
-image_input = gr.Image(type="pil", label="Upload Mammogram Image")
+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 = []
-with gr.Blocks():  # Ensures proper layout management
-    for i in range(0, len(feature_names), num_columns):
-        with gr.Row():
-            for feature in feature_names[i:i+num_columns]:
-                feature_inputs.append(gr.Number(label=feature, scale=0.5))
-
-
-# Example buttons
-def fill_example(example):
-    return {feature_inputs[i]: example[i] for i in range(len(feature_inputs))}
-
-examples = [
-    ["Neural Network", None] + benign_example,
-    ["Neural Network", None] + malignant_example
-]
-
-iface = gr.Interface(
-    fn=classify,
-    inputs=[model_selector, image_input] + feature_inputs,
-    outputs="text",
-    title="Breast Cancer Classification",
-    description="Choose between ViT (image-based) and Neural Network (feature-based) classification.",
-    live=True
-)
-
-iface.launch()
+# 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
+    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))
+
+    # Example buttons
+    def fill_example(example):
+        return {feature_inputs[i]: example[i] for i in range(len(feature_inputs))}
+
+    examples = [
+        ["Neural Network", None] + benign_example,
+        ["Neural Network", None] + malignant_example
+    ]
+
+    gr.Interface(
+        fn=classify,
+        inputs=[model_selector, image_input] + feature_inputs,
+        outputs="text",
+        examples=examples,
+        live=True
+    )
+
+demo.launch()