Update app.py

app.py

@@ -64,7 +64,7 @@ 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."
+            return "❌ Please upload an image for ViT classification."
         image = image.convert("RGB")
         input_tensor = transform(image).unsqueeze(0).to(device)
 
@@ -72,52 +72,59 @@ def classify(model_choice, image=None, *features):
             output = vit_model(input_tensor)
             predicted_class = torch.argmax(output, dim=1).item()
 
-        return class_names[predicted_class]
+        return f"🔍 **Prediction:** {class_names[predicted_class]}"
 
     elif model_choice == "Neural Network":
         if any(f is None for f in features):
-            return "Please enter all 30 numerical 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
         prediction = nn_model.predict(input_data_std) if nn_model else [[0, 1]]
         predicted_class = np.argmax(prediction)
 
-        return class_names[predicted_class]
+        return f"🔍 **Prediction:** {class_names[predicted_class]}"
 
 # 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.")
+    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")
-    image_input = gr.Image(type="pil", label="Upload Image")
+    with gr.Row():
+        model_selector = gr.Radio(["ViT", "Neural Network"], label="🔬 Choose Model", value="ViT")
+
+    image_input = gr.Image(type="pil", label="📷 Upload Image (for ViT)", visible=True)
 
-    # Arrange feature inputs in a matrix layout (3 columns)
-    num_columns = 3
-    feature_inputs = []
+    feature_inputs = [gr.Number(label=feature) for feature in feature_names]
 
+    # Dynamically arrange feature inputs into rows of 3 columns
     with gr.Row():
-        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))
+        for i in range(0, len(feature_inputs), 3):
+            with gr.Column():
+                for j in range(3):
+                    if i + j < len(feature_inputs):
+                        feature_inputs[i + j].render()
 
+    # Example buttons
     def fill_example(example):
         """Pre-fills example inputs."""
         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
-    ).render()
+    with gr.Row():
+        example_btn_1 = gr.Button("🔵 Benign Example")
+        example_btn_2 = gr.Button("🔴 Malignant Example")
+
+    output_text = gr.Textbox(label="🔍 Model Prediction", interactive=False)
+
+    # Logic to toggle inputs based on model selection
+    def toggle_inputs(choice):
+        return gr.update(visible=(choice == "ViT")), gr.update(visible=(choice == "Neural Network"))
+
+    model_selector.change(toggle_inputs, model_selector, [image_input, *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_button = gr.Button("🚀 Classify")
+    classify_button.click(classify, [model_selector, image_input] + feature_inputs, output_text)
 
 demo.launch()