Update app.py

app.py

@@ -55,8 +55,91 @@ def get_ai_explanation(diagnosis, probabilities):
     
     return completion.choices[0].message.content
 
-# Rest of your existing functions (classify_and_visualize, show_final_layer_attention_maps, etc.) remain the same
-[Previous functions remain unchanged...]
+def classify_and_visualize(img, device="cpu", discard_ratio=0.9, head_fusion="mean"):
+    img = img.convert("RGB")
+    processed_input = processor(images=img, return_tensors="pt").to(device)
+    processed_input = processed_input["pixel_values"].to(device)
+
+    with torch.no_grad():
+        outputs = model(processed_input, output_attentions=True)
+        logits = outputs.logits
+        probabilities = torch.softmax(logits, dim=1)[0].tolist()
+        prediction = torch.argmax(logits, dim=-1).item()
+        predicted_class = class_names[prediction]
+
+    result = {class_name: prob for class_name, prob in zip(class_names, probabilities)}
+
+    # Generate attention heatmap
+    heatmap_img = show_final_layer_attention_maps(
+        outputs, processed_input, device, discard_ratio, head_fusion
+    )
+
+    return {"probabilities": result, "heatmap": heatmap_img}
+
+def show_final_layer_attention_maps(
+    outputs,
+    processed_input,
+    device,
+    discard_ratio=0.6,
+    head_fusion="max",
+    only_last_layer=False,
+):
+    with torch.no_grad():
+        image = processed_input.squeeze(0)
+        image = image - image.min()
+        image = image / image.max()
+
+        result = torch.eye(outputs.attentions[0].size(-1)).to(device)
+        if only_last_layer:
+            attention_list = outputs.attentions[-1].unsqueeze(0).to(device)
+        else:
+            attention_list = outputs.attentions
+
+        for attention in attention_list:
+            if head_fusion == "mean":
+                attention_heads_fused = attention.mean(axis=1)
+            elif head_fusion == "max":
+                attention_heads_fused = attention.max(axis=1)[0]
+            elif head_fusion == "min":
+                attention_heads_fused = attention.min(axis=1)[0]
+
+            flat = attention_heads_fused.view(attention_heads_fused.size(0), -1)
+            _, indices = flat.topk(int(flat.size(-1) * discard_ratio), -1, False)
+            indices = indices[indices != 0]
+            flat[0, indices] = 0
+
+            I = torch.eye(attention_heads_fused.size(-1)).to(device)
+            a = (attention_heads_fused + 1.0 * I) / 2
+            a = a / a.sum(dim=-1)
+            result = torch.matmul(a, result)
+
+        mask = result[0, 0, 1:]
+        width = int(mask.size(-1) ** 0.5)
+        mask = mask.reshape(width, width).cpu().numpy()
+        mask = mask / np.max(mask)
+
+        mask = cv2.resize(mask, (224, 224))
+        mask = (mask - np.min(mask)) / (np.max(mask) - np.min(mask))
+        heatmap = plt.cm.jet(mask)[:, :, :3]
+
+        showed_img = image.permute(1, 2, 0).detach().cpu().numpy()
+        showed_img = (showed_img - np.min(showed_img)) / (
+            np.max(showed_img) - np.min(showed_img)
+        )
+        superimposed_img = heatmap * 0.4 + showed_img * 0.6
+
+        superimposed_img_pil = Image.fromarray(
+            (superimposed_img * 255).astype(np.uint8)
+        )
+        return superimposed_img_pil
+
+def load_examples_from_folder(folder_path):
+    examples = []
+    if os.path.exists(folder_path):
+        for file in os.listdir(folder_path):
+            if file.endswith((".png", ".jpg", ".jpeg")):
+                examples.append(os.path.join(folder_path, file))
+    return examples
 
 def create_interface():
     # Custom CSS