Update app.py

app.py

@@ -1,57 +1,30 @@
 import streamlit as st
-from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
+from ultralytics import YOLO
 from PIL import Image
-import torch
-import numpy as np
+import io
 
-st.set_page_config(page_title="Segmentation pulmonaire", layout="wide")
+st.set_page_config(layout="centered")
 
 @st.cache_resource
 def load_model():
-    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):
-    mask_rgb = np.zeros((*mask.shape, 3), dtype=np.uint8)
-    mask_rgb[mask == 1] = [255, 0, 0]
-    
-    image_np = np.array(image)
-    if len(image_np.shape) == 2:
-        image_np = np.stack([image_np] * 3, axis=-1)
-    
-    if mask_rgb.shape[:2] != image_np.shape[:2]:
-        mask_rgb = Image.fromarray(mask_rgb).resize(image.size)
-        mask_rgb = np.array(mask_rgb)
-    
-    overlay = Image.fromarray((image_np * (1 - alpha) + mask_rgb * alpha).astype(np.uint8))
-    return overlay
+    return YOLO("keremberke/yolov8m-chest-xray-classification")
 
 def main():
-    st.title("Segmentation des poumons sur radiographie")
+    st.title("Analyse Radiographie Thoracique")
     
-    processor, model = load_model()
+    model = load_model()
     
-    uploaded_file = st.file_uploader("Télécharger une radiographie thoracique", type=["jpg", "jpeg", "png"])
+    uploaded_file = st.file_uploader("Télécharger une radiographie", type=["jpg", "jpeg", "png"])
     
     if 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)
+        image = Image.open(uploaded_file)
+        resized_image = image.resize((640, 640))
+        st.image(resized_image, width=400)
         
         if st.button("Analyser"):
-            with st.spinner("Analyse en cours..."):
-                try:
-                    inputs = processor(images=image, return_tensors="pt")
-                    outputs = model(**inputs)
-                    logits = outputs.logits
-                    predicted_mask = torch.argmax(logits, dim=1).squeeze().numpy()
-                    
-                    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)}")
+            results = model.predict(source=resized_image)
+            st.write(f"Résultat: {results[0].names[results[0].probs.argmax()]}")
+            st.write(f"Confiance: {results[0].probs.max():.2%}")
 
 if __name__ == "__main__":
     main()