Update app.py

app.py

@@ -7,15 +7,11 @@ st.set_page_config(page_title="Aide au diagnostic radiologique", layout="wide")
 
 def load_models():
     models = {
-        'Fracture': "kathleen/vit-base-fracture-detection",
-        'Pneumothorax': "nickmuchi/pneumothorax-detection-vit",
-        'Pneumonie': "nickmuchi/chest-xray-pneumonia-detection"
+        'Fracture': pipeline("image-classification", model="stanfordaimi/fracture-detection-vit"),
+        'Pneumothorax': pipeline("image-classification", model="ahmedos/pneumothorax-detection"),
+        'Pneumonie': pipeline("image-classification", model="ahmedos/pneumonia-detection")
     }
-    
-    loaded_models = {}
-    for name, model_id in models.items():
-        loaded_models[name] = pipeline("image-classification", model=model_id)
-    return loaded_models
+    return models
 
 @st.cache_resource
 def get_models():
@@ -24,7 +20,8 @@ def get_models():
 def main():
     st.title("Assistant de diagnostic radiologique")
     
-    models = get_models()
+    with st.spinner("Chargement des modèles..."):
+        models = get_models()
     
     uploaded_file = st.file_uploader("Télécharger une image radiologique", type=["jpg", "jpeg", "png"])
     
@@ -37,23 +34,32 @@ def main():
         with col1:
             if st.button("Détecter Fracture"):
                 with st.spinner("Analyse en cours..."):
-                    result = models['Fracture'](image)
-                    st.write(f"Résultat: {result[0]['label']}")
-                    st.write(f"Confiance: {result[0]['score']:.2%}")
+                    try:
+                        result = models['Fracture'](image)
+                        st.write(f"Résultat: {result[0]['label']}")
+                        st.write(f"Confiance: {result[0]['score']:.2%}")
+                    except Exception as e:
+                        st.error(f"Erreur: {str(e)}")
         
         with col2:
             if st.button("Détecter Pneumothorax"):
                 with st.spinner("Analyse en cours..."):
-                    result = models['Pneumothorax'](image)
-                    st.write(f"Résultat: {result[0]['label']}")
-                    st.write(f"Confiance: {result[0]['score']:.2%}")
+                    try:
+                        result = models['Pneumothorax'](image)
+                        st.write(f"Résultat: {result[0]['label']}")
+                        st.write(f"Confiance: {result[0]['score']:.2%}")
+                    except Exception as e:
+                        st.error(f"Erreur: {str(e)}")
         
         with col3:
             if st.button("Détecter Pneumonie"):
                 with st.spinner("Analyse en cours..."):
-                    result = models['Pneumonie'](image)
-                    st.write(f"Résultat: {result[0]['label']}")
-                    st.write(f"Confiance: {result[0]['score']:.2%}")
+                    try:
+                        result = models['Pneumonie'](image)
+                        st.write(f"Résultat: {result[0]['label']}")
+                        st.write(f"Confiance: {result[0]['score']:.2%}")
+                    except Exception as e:
+                        st.error(f"Erreur: {str(e)}")
 
 if __name__ == "__main__":
     main()