Update app.py

app.py

@@ -1,43 +1,59 @@
 import streamlit as st
-from transformers import AutoProcessor, AutoModelForImageClassification
+from transformers import pipeline
 from PIL import Image
-import requests
-
-# Titre de l'application
-st.title("RADPID: Assistant de diagnostic radiologique")
-st.markdown("**Chargez une radiographie et sélectionnez la tâche souhaitée :**")
-
-# Sélection des tâches
-task = st.radio("Sélectionnez une tâche", ["Fracture Detection", "Pneumothorax Detection", "Pneumonia Detection"])
-
-# Modèles
-models = {
-    "Fracture Detection": "facebook/detr-resnet-50",
-    "Pneumothorax Detection": "RGDancer/Pneumothorax_detection",
-    "Pneumonia Detection": "wanghaoy/Chest_XRay_Pneumonia",
-}
-
-# Charger le modèle et le processeur correspondant
-model_name = models[task]
-processor = AutoProcessor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
-
-# Upload de l'image
-uploaded_file = st.file_uploader("Upload your Chest X-Ray image", type=["png", "jpg", "jpeg"])
-
-if uploaded_file is not None:
-    # Charger l'image
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Image Uploadée", use_column_width=True)
-
-    # Effectuer la prédiction
-    st.markdown("### Résultat de la prédiction :")
-    with st.spinner("Analyse en cours..."):
-        inputs = processor(images=image, return_tensors="pt")
-        outputs = model(**inputs)
-        predictions = outputs.logits.softmax(dim=-1).tolist()
-
-        # Afficher les scores
-        st.write(f"Scores pour la tâche '{task}':")
-        st.json(predictions)
-
+import torch
+
+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"
+    }
+    
+    loaded_models = {}
+    for name, model_id in models.items():
+        loaded_models[name] = pipeline("image-classification", model=model_id)
+    return loaded_models
+
+@st.cache_resource
+def get_models():
+    return load_models()
+
+def main():
+    st.title("Assistant de diagnostic radiologique")
+    
+    models = get_models()
+    
+    uploaded_file = st.file_uploader("Télécharger une image radiologique", type=["jpg", "jpeg", "png"])
+    
+    if uploaded_file:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Image téléchargée", use_column_width=True)
+        
+        col1, col2, col3 = st.columns(3)
+        
+        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%}")
+        
+        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%}")
+        
+        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%}")
+
+if __name__ == "__main__":
+    main()