Update app.py

app.py

@@ -4,90 +4,44 @@ from ultralytics import YOLO
 import cv2, os
 from datetime import datetime
 import numpy as np
-import mediapipe as mp
-from dotenv import load_dotenv
-
-# ---------------- Charger config ----------------
-load_dotenv()
-SAVE_LIMIT_FREE = int(os.getenv("SAVE_LIMIT_FREE", 5))
-PREMIUM_KEY = os.getenv("PREMIUM_KEY", "VOTRE_CLE_PREMIUM")
 
 # ---------------- Config générale ----------------
 MODEL_PATH = "best.pt"
 SAVE_DIR = os.path.join("/tmp", "results")
 os.makedirs(SAVE_DIR, exist_ok=True)
 
-# Charger le modèle YOLO une seule fois
+# ---------------- Charger le modèle (1 seule fois) ----------------
 @st.cache_resource
 def load_model():
     return YOLO(MODEL_PATH)
 
 model = load_model()
 
-# ---------------- MediaPipe Face Detection (chargé 1 seule fois) ----------------
-mp_face_detection = mp.solutions.face_detection.FaceDetection(min_detection_confidence=0.6)
-
-def _largest_face_bbox(np_img):
-    h, w = np_img.shape[:2]
-    results = mp_face_detection.process(cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR))
-    if not results.detections:
-        return None
-    boxes = []
-    for det in results.detections:
-        rel = det.location_data.relative_bounding_box
-        x1 = int(max(0, rel.xmin) * w)
-        y1 = int(max(0, rel.ymin) * h)
-        x2 = int(min(1.0, rel.xmin + rel.width) * w)
-        y2 = int(min(1.0, rel.ymin + rel.height) * h)
-        boxes.append((x1, y1, x2, y2))
-    boxes.sort(key=lambda b: (b[2]-b[0])*(b[3]-b[1]), reverse=True)
-    return boxes[0] if boxes else None
-
-# ---------------- Etat utilisateur ----------------
-if "uploads_count" not in st.session_state:
-    st.session_state.uploads_count = 0
-if "premium_access" not in st.session_state:
-    st.session_state.premium_access = False
+# ---------------- Limitation compte gratuit ----------------
+MAX_IMAGES = 3
+MAX_VIDEOS = 1
+
+if "image_count" not in st.session_state:
+    st.session_state.image_count = 0
+if "video_count" not in st.session_state:
+    st.session_state.video_count = 0
 
 # ---------------- Fonctions utilitaires ----------------
-def check_limit():
-    """Vérifie la limite gratuite."""
-    if not st.session_state.premium_access and st.session_state.uploads_count >= SAVE_LIMIT_FREE:
-        st.warning(f"⚠️ Limite gratuite atteinte ({SAVE_LIMIT_FREE} uploads). Passez en mode premium pour continuer.")
-        return False
-    return True
-
-def predict_image(image, conf=0.85, show_labels=True):
-    if not check_limit():
-        return None
-    np_img = np.array(image)
-
-    # Détection visage obligatoire
-    face_bbox = _largest_face_bbox(np_img)
-    if face_bbox is None:
-        st.warning("⚠️ Aucun visage humain détecté. Veuillez centrer le visage.")
-        return None
-
-    if np_img.shape[2] == 4:
-        np_img = cv2.cvtColor(np_img, cv2.COLOR_RGBA2BGR)
+def predict_image(image, conf=0.25, show_labels=True):
+    image = np.array(image)
+    if image.shape[2] == 4:
+        image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)
     else:
-        np_img = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)
-
-    results = model.predict(source=np_img, conf=conf, verbose=False)
-    if len(results[0].boxes) == 0:
-        return None
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
 
+    results = model.predict(source=image, conf=conf, verbose=False)
     annotated_image = results[0].plot(labels=show_labels)
+
     out_path = os.path.join(SAVE_DIR, f"image_result_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png")
     cv2.imwrite(out_path, annotated_image)
-
-    st.session_state.uploads_count += 1
     return out_path
 
-def predict_video(video_path, conf=0.85, show_labels=True):
-    if not check_limit():
-        return None
-
+def predict_video(video_path, conf=0.25, show_labels=True):
     cap = cv2.VideoCapture(video_path)
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out_path = os.path.join(SAVE_DIR, f"video_result_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4")
@@ -95,76 +49,85 @@ def predict_video(video_path, conf=0.85, show_labels=True):
     fps = cap.get(cv2.CAP_PROP_FPS) or 30
     width, height = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
-    detections = 0
 
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
         results = model.predict(frame, conf=conf, verbose=False)
-        if len(results[0].boxes) > 0:
-            detections += 1
         annotated = results[0].plot(labels=show_labels)
         out.write(annotated)
 
     cap.release()
     out.release()
-    if detections == 0:
-        return None
-
-    st.session_state.uploads_count += 1
     return out_path
 
 # ---------------- Interface Streamlit ----------------
 st.title("🧠 Stroke-IA – Détection AVC par IA")
 
-# ---------------- Sidebar ----------------
-st.sidebar.header("⚙️ Paramètres utilisateur")
-conf_threshold = st.sidebar.slider("Seuil de confiance", 0.1, 1.0, 0.85, 0.05)
+# Sidebar (paramètres utilisateur)
+st.sidebar.header("⚙️ Paramètres")
+conf_threshold = st.sidebar.slider("Seuil de confiance", 0.1, 1.0, 0.25, 0.05)
 show_labels = st.sidebar.checkbox("Afficher les labels", value=True)
 
-st.sidebar.header("🔑 Premium / Essai")
-if not st.session_state.premium_access:
-    user_key = st.sidebar.text_input("Entrez votre clé premium :", type="password")
-    if user_key == PREMIUM_KEY:
-        st.session_state.premium_access = True
-        st.sidebar.success("✅ Mode premium activé ! La limitation est levée.")
+# Sidebar quota global
+st.sidebar.header("📊 Utilisation gratuite")
+st.sidebar.write(f"🖼️ Images utilisées : **{st.session_state.image_count}/{MAX_IMAGES}**")
+st.sidebar.write(f"🎥 Vidéos utilisées : **{st.session_state.video_count}/{MAX_VIDEOS}**")
+
+st.sidebar.header("📂 Exemples rapides")
+if st.sidebar.button("Tester une image exemple"):
+    if os.path.exists("example.jpg"):
+        img = Image.open("example.jpg")
+        path = predict_image(img, conf=conf_threshold, show_labels=show_labels)
+        st.image(path, caption="Exemple annoté", use_container_width=True)
+    else:
+        st.warning("⚠️ Aucun fichier example.jpg trouvé.")
 
-# Afficher compteur
-if not st.session_state.premium_access:
-    st.sidebar.info(f"📊 Utilisation gratuite : {st.session_state.uploads_count}/{SAVE_LIMIT_FREE}")
+if st.sidebar.button("Tester une vidéo exemple"):
+    if os.path.exists("example.mp4"):
+        path = predict_video("example.mp4", conf=conf_threshold, show_labels=show_labels)
+        st.video(path)
+    else:
+        st.warning("⚠️ Aucun fichier example.mp4 trouvé.")
 
-# ---------------- Upload vidéo ----------------
+# Section vidéo upload
 st.header("🎥 Détection sur vidéo")
-video_file = st.file_uploader("Uploader une vidéo", type=["mp4", "mov"])
-if video_file and st.button("Analyser la vidéo"):
-    temp_path = os.path.join(SAVE_DIR, "temp_video.mp4")
-    with open(temp_path, "wb") as f:
-        f.write(video_file.read())
-    result_path = predict_video(temp_path, conf=conf_threshold, show_labels=show_labels)
-    if result_path is None:
-        st.success(f"✅ Aucun AVC détecté ou limite gratuite atteinte.")
-    else:
+
+remaining_videos = MAX_VIDEOS - st.session_state.video_count
+st.info(f"🎬 Il vous reste **{remaining_videos} vidéo(s)** gratuite(s).")
+
+if st.session_state.video_count >= MAX_VIDEOS:
+    st.error("⚠️ Limite vidéo atteinte. Passez en premium pour continuer.")
+else:
+    video_file = st.file_uploader("Uploader une vidéo (mp4, mov, etc.)", type=["mp4", "mov"], key="video")
+    if video_file and st.button("Analyser la vidéo"):
+        st.session_state.video_count += 1
+        temp_path = os.path.join(SAVE_DIR, "temp_video.mp4")
+        with open(temp_path, "wb") as f:
+            f.write(video_file.read())
+        result_path = predict_video(temp_path, conf=conf_threshold, show_labels=show_labels)
         st.video(result_path)
 
-# ---------------- Upload image ----------------
+# Section image upload
 st.header("🖼️ Détection sur image")
-image_file = st.file_uploader("Uploader une image", type=["jpg", "jpeg", "png"])
-if image_file and st.button("Analyser l'image"):
-    image = Image.open(image_file)
-    result_path = predict_image(image, conf=conf_threshold, show_labels=show_labels)
-    if result_path is None:
-        st.success(f"✅ Aucun AVC détecté ou limite gratuite atteinte.")
-    else:
+
+remaining_images = MAX_IMAGES - st.session_state.image_count
+st.info(f"🖼️ Il vous reste **{remaining_images} image(s)** gratuite(s).")
+
+if st.session_state.image_count >= MAX_IMAGES:
+    st.error("⚠️ Limite images atteinte. Passez en premium pour continuer.")
+else:
+    image_file = st.file_uploader("Uploader une image", type=["jpg", "jpeg", "png"], key="image")
+    if image_file and st.button("Analyser l'image"):
+        st.session_state.image_count += 1
+        image = Image.open(image_file)
+        result_path = predict_image(image, conf=conf_threshold, show_labels=show_labels)
         st.image(result_path, caption="Image annotée", use_container_width=True)
 
-# ---------------- Disclaimer ----------------
+# Disclaimer
 st.markdown(f"""
 ---
-👨‍💻 **Badsi Djilali** — Ingénieur Deep Learning  
-🚀 Créateur de **Stroke_IA_Detection**  
-🧠 (Détection d'asymétrie faciale & AVC par IA)
-
 ⚠️ **Disclaimer :** Stroke-IA est une démo technique, pas un avis médical.  
 © {datetime.now().year} — Badsi Djilali.
 """)