Update app.py

app.py

@@ -1,10 +1,10 @@
 import torch
-import torch
 import torch.nn as nn
 from transformers import ViTImageProcessor, ViTModel, BertTokenizerFast, BertModel
 from PIL import Image
 import gradio as gr
 
+# Model definition and setup
 class VisionLanguageModel(nn.Module):
     def __init__(self):
         super(VisionLanguageModel, self).__init__()
@@ -33,7 +33,6 @@ class VisionLanguageModel(nn.Module):
         logits = self.classifier(combined_features)
         return logits
 
-# Load the model checkpoint with safer loading
 model = VisionLanguageModel()
 model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu'), weights_only=True))
 model.eval()
@@ -42,10 +41,7 @@ tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
 feature_extractor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k')
 
 def predict(image, text_input):
-    # Preprocess the image
     image = feature_extractor(images=image, return_tensors="pt").pixel_values
-
-    # Preprocess the text
     encoding = tokenizer(
         text_input,
         add_special_tokens=True,
@@ -54,8 +50,6 @@ def predict(image, text_input):
         truncation=True,
         return_tensors='pt'
     )
-
-    # Make a prediction
     with torch.no_grad():
         outputs = model(
             input_ids=encoding['input_ids'],
@@ -63,19 +57,44 @@ def predict(image, text_input):
             pixel_values=image
         )
     _, prediction = torch.max(outputs, dim=1)
-    return "Malignant" if prediction.item() == 1 else "Benign"
-
-# Define Gradio interface with updated component syntax
-iface = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Image(type="pil", label="Upload Skin Lesion Image"),
-        gr.Textbox(label="Clinical Information (e.g., patient age, symptoms)")
-    ],
-    outputs="text",
-    title="Skin Lesion Classification Demo",
-    description="This model classifies skin lesions as benign or malignant based on an image and clinical information."
-)
-
-iface.launch()
+    return prediction.item()  # 1 for Malignant, 0 for Benign
+
+# Enhanced UI with color-coded prediction display
+with gr.Blocks(css="""
+    .benign {background-color: white; border: 1px solid lightgray; padding: 10px; border-radius: 5px;}
+    .malignant {background-color: white; border: 1px solid lightgray; padding: 10px; border-radius: 5px;}
+    .benign.correct {background-color: lightgreen;}
+    .malignant.correct {background-color: lightgreen;}
+""") as demo:
+    gr.Markdown(
+        """
+        # 🩺 SKIN LESION CLASSIFICATION 
+        Upload an image of a skin lesion and provide clinical details to get a prediction of benign or malignant.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="Upload Skin Lesion Image")
+            text_input = gr.Textbox(label="Clinical Information (e.g., patient age, symptoms)")
+
+        with gr.Column(scale=1):
+            benign_output = gr.HTML("<div class='benign'>Benign</div>")
+            malignant_output = gr.HTML("<div class='malignant'>Malignant</div>")
+            gr.Markdown("## Example:")
+            example_image = gr.Image(value="skin_cancer_detection/Unknown-4.png")  # Provide path to an example image
+            example_text = gr.Textbox(value="consistent with resolving/involuting keratoacanthoma 67", interactive=False)
+
+    def display_prediction(image, text_input):
+        prediction = predict(image, text_input)
+        benign_html = "<div class='benign{}'>Benign</div>".format(" correct" if prediction == 0 else "")
+        malignant_html = "<div class='malignant{}'>Malignant</div>".format(" correct" if prediction == 1 else "")
+        return benign_html, malignant_html
+
+    # Submit button and prediction outputs
+    submit_btn = gr.Button("Get Prediction")
+    submit_btn.click(display_prediction, inputs=[image_input, text_input], outputs=[benign_output, malignant_output])
+
+demo.launch()
+