Update app.py

app.py

@@ -1,32 +1,58 @@
 import streamlit as st
-from ultralytics import YOLO
-from PIL import Image
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 import torch
+from PIL import Image
+import numpy as np
 
-st.set_page_config(layout="centered")
+def create_overlay(image, attention_map, alpha=0.5):
+    attention_map = (attention_map - attention_map.min()) / (attention_map.max() - attention_map.min())
+    heatmap = np.uint8(255 * attention_map)
+    heatmap = Image.fromarray(heatmap).resize(image.size)
+    heatmap = np.array(heatmap)
+    heatmap = np.stack([heatmap, np.zeros_like(heatmap), np.zeros_like(heatmap)], axis=-1)
+    
+    image_array = np.array(image)
+    overlay = Image.fromarray(np.uint8(image_array * (1 - alpha) + heatmap * alpha))
+    return overlay
 
 @st.cache_resource
 def load_model():
-    model = YOLO('yolov8m.pt')  # Load base YOLOv8 model
-    model.load('keremberke/yolov8m-chest-xray-classification.pt')  # Load weights
-    return model
+    processor = AutoImageProcessor.from_pretrained("mrm8488/vit-base-patch16-224_finetuned-pneumothorax")
+    model = AutoModelForImageClassification.from_pretrained("mrm8488/vit-base-patch16-224_finetuned-pneumothorax")
+    return processor, model
 
 def main():
-    st.title("Analyse Radiographie Thoracique")
+    st.title("Détection de Pneumothorax")
     
-    model = load_model()
+    processor, model = load_model()
     
     uploaded_file = st.file_uploader("Télécharger une radiographie", type=["jpg", "jpeg", "png"])
     
     if uploaded_file:
-        image = Image.open(uploaded_file)
-        resized_image = image.resize((640, 640))
+        image = Image.open(uploaded_file).convert('RGB')
+        resized_image = image.resize((224, 224))
         st.image(resized_image, width=400)
         
         if st.button("Analyser"):
-            results = model.predict(source=resized_image)
-            st.write(f"Résultat: {results[0].names[results[0].probs.argmax()]}")
-            st.write(f"Confiance: {results[0].probs.max():.2%}")
+            with st.spinner("Analyse en cours..."):
+                inputs = processor(images=resized_image, return_tensors="pt")
+                outputs = model(**inputs)
+                probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+                
+                # Obtenir attentions des dernières couches
+                attention = outputs.hidden_states[-1].mean(1)[0].detach().numpy()
+                attention_map = attention.reshape(14, 14)  # ViT patch size
+                
+                # Créer overlay
+                overlay = create_overlay(resized_image, attention_map)
+                
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.write("Résultat:", model.config.id2label[outputs.logits.argmax(-1).item()])
+                    st.write(f"Confiance: {probs.max().item():.2%}")
+                
+                with col2:
+                    st.image(overlay, caption="Zones suspectes", width=400)
 
 if __name__ == "__main__":
     main()