Update app.py

app.py

@@ -1,58 +1,33 @@
 import streamlit as st
-from transformers import AutoImageProcessor, AutoModelForImageClassification
-import torch
+from transformers import pipeline
 from PIL import Image
-import numpy as np
-
-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():
-    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
+    return pipeline("image-classification", model="mrm8488/vit-base-patch16-224_finetuned-pneumothorax", top_k=2)
 
 def main():
     st.title("Détection de Pneumothorax")
     
-    processor, model = load_model()
+    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).convert('RGB')
         resized_image = image.resize((224, 224))
-        st.image(resized_image, width=400)
+        
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(resized_image, width=300, caption="Image originale")
         
         if st.button("Analyser"):
             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%}")
-                
+                results = model(resized_image)
                 with col2:
-                    st.image(overlay, caption="Zones suspectes", width=400)
+                    st.write("Résultats:")
+                    for result in results:
+                        st.write(f"{result['label']}: {result['score']:.2%}")
 
 if __name__ == "__main__":
     main()