Update app.py

app.py

@@ -1,38 +1,63 @@
 import gradio as gr
 from ultralytics import YOLO  # Importer YOLOv8
 import numpy as np
+import cv2
 
 # Charger le modèle YOLOv8
 model = YOLO('best.pt')  # Remplace 'best.pt' par le chemin de ton modèle
 
+# Fonction pour dessiner les boîtes de détection sur l'image
+def draw_boxes(image, detections, class_names, confidence_threshold):
+    for detection in detections:
+        if detection.conf[0] >= confidence_threshold:
+            class_index = int(detection.cls[0])  # Obtenir l'indice de la classe
+            class_name = class_names[class_index]  # Nom de la classe
+            confidence = format(detection.conf[0], ".2f")  # Confiance
+            
+            # Obtenir les coordonnées des boîtes
+            xyxy = detection.xyxy[0].cpu().numpy().astype(int)
+            (x1, y1, x2, y2) = xyxy
+            
+            # Dessiner la boîte
+            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            label = f"{class_name} ({confidence})"
+            
+            # Ajouter le texte (classe et confiance) au-dessus de la boîte
+            cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+    
+    return image
+
 # Fonction d'inférence
-def detect(img):
+def detect(img, confidence_threshold):
+    # Convertir l'image PIL en format approprié pour OpenCV
+    img_cv = np.array(img)
+    img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
+    
     # Faire une prédiction avec YOLOv8
-    results = model(img)  # YOLOv8 prend directement l'image en entrée
-
-    # Si des objets sont détectés, renvoyer la classe de l'objet
+    results = model(img_cv)  # YOLOv8 prend directement l'image en entrée
+    
+    # Si des objets sont détectés, dessiner les boîtes
     if results:
         detections = results[0].boxes  # Les résultats des détections
         if len(detections) > 0:
-            # Obtenir la classe de la détection avec la probabilité la plus élevée
-            detection = detections[0]  # On prend la première détection (peut être ajusté)
-            class_index = int(detection.cls[0])  # Obtenir l'indice de la classe
-            class_names = ['AMAZONE', 'BIOGUERRA','PORTE DU NON RETOUR', 'REVENANT', 'ZANGBETO']
-            class_name = class_names[class_index]  # Nom de la classe prédite
-            confidence = format(detection.conf[0], ".2f")  # Confiance de la détection
-            return f"Classe : {class_name}, Confiance : {confidence}"
+            class_names = ['AMAZONE', 'BIOGUERRA', 'PORTE DU NON RETOUR', 'REVENANT', 'ZANGBETO']
+            img_with_boxes = draw_boxes(img_cv, detections, class_names, confidence_threshold)
+            return cv2.cvtColor(img_with_boxes, cv2.COLOR_BGR2RGB)  # Convertir l'image en format RGB pour Gradio
         else:
-            return "Aucune détection"
+            return img  # Si aucune détection, retourner l'image originale
     else:
-        return "Aucune détection"
+        return img  # Si aucune détection, retourner l'image originale
 
 # Interface utilisateur avec Gradio
 title = "Orisha YOLOv8"
 
 iface = gr.Interface(
     fn=detect, 
-    inputs=gr.Image(type="pil", image_mode='RGB'),  # Charger une image
-    outputs=gr.Textbox(label="Résultat", lines=2),  # Afficher le résultat
+    inputs=[
+        gr.Image(type="pil", image_mode='RGB', label="Image à détecter"),  # Charger une image
+        gr.Slider(0.0, 1.0, value=0.25, step=0.01, label="Confiance minimale"),  # Paramétrer la confiance
+    ],  
+    outputs=gr.Image(type="pil", label="Résultat avec détections"),  # Afficher le résultat avec boîtes
     title=title
 )
 
@@ -40,6 +65,48 @@ iface = gr.Interface(
 iface.launch(inline=False)
 
 
+# import gradio as gr
+# from ultralytics import YOLO  # Importer YOLOv8
+# import numpy as np
+
+# # Charger le modèle YOLOv8
+# model = YOLO('best.pt')  # Remplace 'best.pt' par le chemin de ton modèle
+
+# # Fonction d'inférence
+# def detect(img):
+#     # Faire une prédiction avec YOLOv8
+#     results = model(img)  # YOLOv8 prend directement l'image en entrée
+
+#     # Si des objets sont détectés, renvoyer la classe de l'objet
+#     if results:
+#         detections = results[0].boxes  # Les résultats des détections
+#         if len(detections) > 0:
+#             # Obtenir la classe de la détection avec la probabilité la plus élevée
+#             detection = detections[0]  # On prend la première détection (peut être ajusté)
+#             class_index = int(detection.cls[0])  # Obtenir l'indice de la classe
+#             class_names = ['AMAZONE', 'BIOGUERRA','PORTE DU NON RETOUR', 'REVENANT', 'ZANGBETO']
+#             class_name = class_names[class_index]  # Nom de la classe prédite
+#             confidence = format(detection.conf[0], ".2f")  # Confiance de la détection
+#             return f"Classe : {class_name}, Confiance : {confidence}"
+#         else:
+#             return "Aucune détection"
+#     else:
+#         return "Aucune détection"
+
+# # Interface utilisateur avec Gradio
+# title = "Orisha YOLOv8"
+
+# iface = gr.Interface(
+#     fn=detect, 
+#     inputs=gr.Image(type="pil", image_mode='RGB'),  # Charger une image
+#     outputs=gr.Textbox(label="Résultat", lines=2),  # Afficher le résultat
+#     title=title
+# )
+
+# # Lancer l'application
+# iface.launch(inline=False)
+
+
 # import gradio as gr
 # import tensorflow as tf
 # import numpy as np