Update app.py

app.py

@@ -13,21 +13,17 @@ def load_model():
     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
+    mask_rgb[mask == 1] = [255, 0, 0]
     
-    # 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
+    if len(image_np.shape) == 2:
         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
 
@@ -39,21 +35,18 @@ def main():
     uploaded_file = st.file_uploader("Télécharger une radiographie thoracique", type=["jpg", "jpeg", "png"])
     
     if uploaded_file:
-        image = Image.open(uploaded_file)
+        # Ouvrir et convertir l'image en RGB
+        image = Image.open(uploaded_file).convert('RGB')
         st.image(image, caption="Image originale", use_column_width=True)
         
         if st.button("Analyser"):
             with st.spinner("Analyse en cours..."):
                 try:
-                    # 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)