Update app.py

app.py

@@ -1,31 +1,62 @@
 import streamlit as st
-from transformers import AutoImageProcessor, AutoModelForImageClassification, pipeline
+from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
 from PIL import Image
 import torch
+import numpy as np
 
-st.set_page_config(page_title="Détection maladies pulmonaires", layout="wide")
+st.set_page_config(page_title="Segmentation pulmonaire", layout="wide")
 
 @st.cache_resource
 def load_model():
-    return pipeline("image-classification", model="gianlab/swin-tiny-patch4-window7-224-finetuned-lungs-disease")
+    processor = AutoImageProcessor.from_pretrained("Tianmu28/segformer-b0-segments-lungs-xray")
+    model = SegformerForSemanticSegmentation.from_pretrained("Tianmu28/segformer-b0-segments-lungs-xray")
+    return processor, model
+
+def create_overlay(image, mask, alpha=0.5):
+    # Convertir le masque en RGB
+    mask_rgb = np.zeros((*mask.shape, 3), dtype=np.uint8)
+    mask_rgb[mask == 1] = [255, 0, 0]  # Rouge pour les zones détectées
+    
+    # Convertir l'image en array numpy
+    image_np = np.array(image)
+    if len(image_np.shape) == 2:  # Si l'image est en niveaux de gris
+        image_np = np.stack([image_np] * 3, axis=-1)
+    
+    # Redimensionner le masque si nécessaire
+    if mask_rgb.shape[:2] != image_np.shape[:2]:
+        mask_rgb = Image.fromarray(mask_rgb).resize(image.size)
+        mask_rgb = np.array(mask_rgb)
+    
+    # Créer l'overlay
+    overlay = Image.fromarray((image_np * (1 - alpha) + mask_rgb * alpha).astype(np.uint8))
+    return overlay
 
 def main():
-    st.title("Détection de maladies pulmonaires")
+    st.title("Segmentation des poumons sur radiographie")
     
-    model = load_model()
+    processor, model = load_model()
     
     uploaded_file = st.file_uploader("Télécharger une radiographie thoracique", type=["jpg", "jpeg", "png"])
     
     if uploaded_file:
         image = Image.open(uploaded_file)
-        st.image(image, caption="Radiographie", use_column_width=True)
+        st.image(image, caption="Image originale", use_column_width=True)
         
         if st.button("Analyser"):
             with st.spinner("Analyse en cours..."):
                 try:
-                    result = model(image)
-                    st.write(f"Diagnostic: {result[0]['label']}")
-                    st.write(f"Confiance: {result[0]['score']:.2%}")
+                    # Préparer l'image
+                    inputs = processor(images=image, return_tensors="pt")
+                    
+                    # Prédiction
+                    outputs = model(**inputs)
+                    logits = outputs.logits
+                    predicted_mask = torch.argmax(logits, dim=1).squeeze().numpy()
+                    
+                    # Créer et afficher l'overlay
+                    overlay = create_overlay(image, predicted_mask)
+                    st.image(overlay, caption="Zones détectées", use_column_width=True)
+                    
                 except Exception as e:
                     st.error(f"Erreur: {str(e)}")