Update app.py

added model update capability based on known diagnoses from specilists

app.py

@@ -7,11 +7,18 @@ import plotly.graph_objects as go
 import torch
 import numpy as np
 from PIL import Image
-
 model_name = "./best_model"
 processor = ViTImageProcessor.from_pretrained(model_name)
 labels = ['Acne or Rosacea', 'Actinic Keratosis Basal Cell Carcinoma and other Malignant Lesions', 'Atopic Dermatitis', 'Bullous Disease', 'Cellulitis Impetigo and other Bacterial Infections', 'Contact Dermatitis', 'Eczema', 'Exanthems and Drug Eruptions', 'Hair Loss Photos Alopecia and other Hair Diseases', 'Herpes HPV and other STDs', 'Light Diseases and Disorders of Pigmentation', 'Lupus and other Connective Tissue diseases', 'Melanoma Skin Cancer Nevi and Moles', 'Nail Fungus and other Nail Disease', 'Psoriasis pictures Lichen Planus and related diseases', 'Scabies Lyme Disease and other Infestations and Bites', 'Seborrheic Keratoses and other Benign Tumors', 'Systemic Disease', 'Tinea Ringworm Candidiasis and other Fungal Infections', 'Urticaria Hives', 'Vascular Tumors', 'Vasculitis', 'Warts Molluscum and other Viral Infections']
 
+class ViTForImageClassificationWithAttention(ViTForImageClassification):
+    def forward(self, pixel_values):
+        outputs = super().forward(pixel_values)
+        attention = self.vit.encoder.layers[0].attention.attention_weights
+        return outputs, attention
+
+model = ViTForImageClassificationWithAttention.from_pretrained(model_name)
+
 class ViTForImageClassificationWithAttention(ViTForImageClassification):
     def forward(self, pixel_values, output_attentions=True):
         outputs = super().forward(pixel_values, output_attentions=output_attentions)
@@ -20,9 +27,11 @@ class ViTForImageClassificationWithAttention(ViTForImageClassification):
 
 model = ViTForImageClassificationWithAttention.from_pretrained(model_name,attn_implementation="eager")
 
+i_count = 0
 def classify_image(image):
-    img = Image.open(image)
-    inputs = processor(images=img, return_tensors="pt")
+    model_name = "best_model.pth"
+    model.load_state_dict(torch.load(model_name))
+    inputs = processor(images=image, return_tensors="pt")
     outputs, attention = model(**inputs, output_attentions=True)
     logits = outputs.logits
     probs = torch.nn.functional.softmax(logits, dim=1)
@@ -44,13 +53,11 @@ def classify_image(image):
             data=[go.Heatmap(z=attention[0][0, 0, :, :].detach().numpy(), colorscale='Viridis', showscale=False)])
         fig_heatmap.update_layout(title="Attention Heatmap")
     else:
-        fig_heatmap = None
+        fig_heatmap = go.Figure()  # Return an empty plot
 
-    # Overlay the attention heatmap on the input image
-    # Overlay the attention heatmap on the input image
     # Overlay the attention heatmap on the input image
     if attention is not None:
-        img_array = np.array(img)
+        img_array = np.array(image)
         heatmap = np.array(attention[0][0, 0, :, :].detach().numpy())
         heatmap = np.resize(heatmap, (img_array.shape[0], img_array.shape[1]))
         heatmap = heatmap / heatmap.max() * 255  # Normalize heatmap to [0, 255]
@@ -60,26 +67,80 @@ def classify_image(image):
         heatmap_color[:, :, 1] = heatmap  # Green channel
         heatmap_color[:, :, 2] = 0  # Blue channel
         attention_overlay = (img_array * 0.5 + heatmap_color * 0.5).astype(np.uint8)
-        img = Image.fromarray(attention_overlay)
-        img.save("attention_overlay.png")
+        attention_overlay = Image.fromarray(attention_overlay)
+        attention_overlay.save("attention_overlay.png")
         attention_overlay = gr.Image("attention_overlay.png")
     else:
-        attention_overlay = None
+        attention_overlay = gr.Image()  # Return an empty image
 
     # Return the predicted label, the bar chart, and the heatmap
-    return predicted_class_label, fig_bar, attention_overlay
-
-demo = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="filepath"),
-    outputs=[
-        gr.Label(label="Predicted Class Label"),
-        gr.Plot(label="Top 5 Predicted Labels with Probabilities"),
-        gr.Image(label="Attention Overlay")
+    return predicted_class_label, fig_bar, fig_heatmap, attention_overlay
+
+
+def update_model(image, label):
+    # Convert the label to an integer
+    label_idx = labels.index(label)
+    labels_tensor = torch.tensor([label_idx])
+
+    inputs = processor(images=image, return_tensors="pt")
+    loss_fn = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+    # Zero the gradients
+    optimizer.zero_grad()
+
+    # Forward pass
+    outputs, attention = model(**inputs)
+    loss = loss_fn(outputs.logits, labels_tensor)
+
+    # Backward pass
+    loss.backward()
+
+    # Update the model parameters
+    optimizer.step()
+
+    # Save the updated model
+    torch.save(model.state_dict(), "best_model.pth")
+
+    return "Model updated successfully"
+
+
+demo = gr.TabbedInterface(
+    [
+        gr.Interface(
+            fn=classify_image,
+            inputs=[
+                gr.Image(type="pil", label="Image")
+            ],
+            outputs=[
+                gr.Label(label="Predicted Class Label"),
+                gr.Plot(label="Top 5 Predicted Labels with Probabilities")
+            ],
+            title="Dermatological Image Classification Demo",
+            description="Upload an image to see the predicted class label, top 5 predicted labels with probabilities, and attention heatmap",
+            allow_flagging=False
+        ),
+        gr.Interface(
+            fn=update_model,
+            inputs=[
+                gr.Image(type="pil", label="Image"),
+                gr.Radio(
+                    choices=labels,
+                    type="value",
+                    label="Label",
+                    value=labels[0]
+                )
+            ],
+            outputs=[
+                gr.Textbox(label="Model Update Status")
+            ],
+            title="Train Model",
+            description="Upload an image and label to update the model",
+            allow_flagging=False
+        )
     ],
-    title="Dermatological Image Classification Demo",
-    description="Upload an image to see the predicted class label, top 5 predicted labels with probabilities, and attention overlay",
+    title="Dermatological Image Classification and Training"
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)