Spaces:
Running
Running
| import torch | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| import torchvision.transforms as transforms | |
| from model import Yolo_V1 | |
| from utils import cellboxes_to_boxes, non_max_suppression | |
| import cv2 | |
| import os | |
| import glob | |
| import time | |
| from huggingface_hub import hf_hub_download | |
| # Classes PASCAL VOC | |
| CLASSES = [ | |
| "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", | |
| "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", | |
| "pottedplant", "sheep", "sofa", "train", "tvmonitor" | |
| ] | |
| np.random.seed(42) | |
| COLORS = np.random.randint(50, 255, size=(len(CLASSES), 3), dtype=np.uint8) | |
| DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" | |
| MODEL_REPO_ID = "nathbns/yolov1_from_scratch" | |
| MODEL_FILENAME = "checkpoint_epoch_50.pth.tar" | |
| # Charger le modèle depuis Hugging Face Hub | |
| print(f"Chargement du modèle depuis {MODEL_REPO_ID}...") | |
| try: | |
| model_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=MODEL_FILENAME) | |
| print(f"Modèle téléchargé depuis Hugging Face Hub: {model_path}") | |
| except Exception as e: | |
| print(f"Erreur lors du téléchargement: {e}") | |
| print("Tentative de chargement local...") | |
| model_path = MODEL_FILENAME | |
| model = Yolo_V1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE) | |
| checkpoint = torch.load(model_path, map_location=DEVICE) | |
| model.load_state_dict(checkpoint["state_dict"]) | |
| model.eval() | |
| print(f"Modèle chargé avec succès!") | |
| # Info sur le modèle | |
| MODEL_INFO = { | |
| "mAP": checkpoint.get("mAP", "N/A"), | |
| "epoch": checkpoint.get("epoch", "N/A"), | |
| "device": DEVICE, | |
| "classes": len(CLASSES) | |
| } | |
| print(f"entraînement: {MODEL_INFO['mAP']}") | |
| print(f"Device: {DEVICE}") | |
| # Charger des images d'exemple depuis le dossier data | |
| EXAMPLE_IMAGES = [] | |
| if os.path.exists("data/images"): | |
| image_files = glob.glob("data/images/*.jpg")[:20] # Prendre 20 images | |
| EXAMPLE_IMAGES = sorted(image_files) | |
| print(f"{len(EXAMPLE_IMAGES)} images d'exemple chargées") | |
| def draw_boxes(image, boxes): | |
| """Dessine les bounding boxes sur l'image""" | |
| img_array = np.array(image) | |
| height, width = img_array.shape[:2] | |
| for box in boxes: | |
| # box format: [class_pred, prob_score, x, y, width, height] | |
| class_pred = int(box[0]) | |
| confidence = float(box[1]) | |
| x_center, y_center, box_width, box_height = box[2:6] | |
| # Convertir de coordonnées normalisées à pixels | |
| x1 = int((x_center - box_width / 2) * width) | |
| y1 = int((y_center - box_height / 2) * height) | |
| x2 = int((x_center + box_width / 2) * width) | |
| y2 = int((y_center + box_height / 2) * height) | |
| # Couleur de la classe | |
| color = tuple(int(c) for c in COLORS[class_pred]) | |
| # Dessiner le rectangle | |
| cv2.rectangle(img_array, (x1, y1), (x2, y2), color, 2) | |
| # Texte | |
| label = f"{CLASSES[class_pred]}: {confidence:.2f}" | |
| # Fond du texte | |
| (text_width, text_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) | |
| cv2.rectangle(img_array, (x1, y1 - text_height - 5), (x1 + text_width, y1), color, -1) | |
| # Texte blanc | |
| cv2.putText(img_array, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) | |
| return Image.fromarray(img_array) | |
| def detect_objects(image, confidence_threshold, iou_threshold, show_confidence=True): | |
| """Détecte les objets dans une image avec statistiques détaillées""" | |
| if image is None: | |
| return None, None, "**Veuillez uploader ou sélectionner une image**" | |
| start_time = time.time() | |
| # Prétraiter l'image | |
| transform = transforms.Compose([ | |
| transforms.Resize((448, 448)), | |
| transforms.ToTensor(), | |
| ]) | |
| # Garder l'image originale pour l'affichage | |
| original_image = image.copy() | |
| original_size = image.size # (width, height) | |
| # Transformer l'image | |
| img_tensor = transform(image).unsqueeze(0).to(DEVICE) | |
| # Prédiction | |
| with torch.no_grad(): | |
| predictions = model(img_tensor) | |
| # Convertir les prédictions en bounding boxes | |
| bboxes = cellboxes_to_boxes(predictions) | |
| # Non-maximum suppression | |
| nms_boxes = non_max_suppression( | |
| bboxes[0], | |
| iou_threshold=iou_threshold, | |
| threshold=confidence_threshold, | |
| box_format="midpoint" | |
| ) | |
| inference_time = time.time() - start_time | |
| # Dessiner les boxes | |
| result_image = draw_boxes(original_image, nms_boxes) | |
| # Statistiques détaillées | |
| num_detections = len(nms_boxes) | |
| detected_classes = [CLASSES[int(box[0])] for box in nms_boxes] | |
| class_counts = {} | |
| confidence_scores = [] | |
| for box in nms_boxes: | |
| cls = CLASSES[int(box[0])] | |
| conf = float(box[1]) | |
| class_counts[cls] = class_counts.get(cls, 0) + 1 | |
| confidence_scores.append(conf) | |
| # Créer un rapport détaillé | |
| stats = f"##Résultats de détection\n\n" | |
| stats += f"**{num_detections} objet(s) détecté(s)**\n\n" | |
| if num_detections > 0: | |
| stats += f"Temps d'inférence: **{inference_time:.3f}s**\n" | |
| stats += f"Taille image: **{original_size[0]}x{original_size[1]}**\n" | |
| stats += f"Confiance moyenne: **{np.mean(confidence_scores):.2%}**\n\n" | |
| stats += "### Objets détectés:\n\n" | |
| for cls, count in sorted(class_counts.items(), key=lambda x: x[1], reverse=True): | |
| stats += f"- **{cls}**: {count}\n" | |
| if show_confidence: | |
| stats += "\n### Confiances individuelles:\n\n" | |
| for i, box in enumerate(nms_boxes[:10], 1): # Top 10 | |
| cls = CLASSES[int(box[0])] | |
| conf = float(box[1]) | |
| stats += f"{i}. {cls}: {conf:.1%}\n" | |
| if len(nms_boxes) > 10: | |
| stats += f"\n*...et {len(nms_boxes)-10} détection(s) de plus*\n" | |
| else: | |
| stats += "Aucun objet détecté.\n\n" | |
| return original_image, result_image, stats | |
| # Interface Gradio améliorée | |
| with gr.Blocks(title="YOLO v1 - Détection d'objets", theme=gr.themes.Soft(), css=""" | |
| .gradio-container {max-width: 1400px !important} | |
| .example-gallery {height: 400px; overflow-y: auto} | |
| """) as demo: | |
| # En-tête | |
| mAP_display = f"{MODEL_INFO['mAP']:.4f}" if isinstance(MODEL_INFO['mAP'], (int, float)) else MODEL_INFO['mAP'] | |
| gr.Markdown(f""" | |
| # YOLO v1 - Détection d'objets en temps réel | |
| --- | |
| """) | |
| with gr.Tabs(): | |
| # Onglet principal - Détection | |
| with gr.Tab("Détection"): | |
| gr.Markdown(""" | |
| ### Uploadez votre image ou sélectionnez un exemple | |
| **Classes PASCAL VOC :** aeroplane, bicycle, bird, boat, bottle, bus, car, cat, chair, cow, | |
| diningtable, dog, horse, motorbike, person, pottedplant, sheep, sofa, train, tvmonitor | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| input_image = gr.Image(type="pil", label="Image d'entrée") | |
| with gr.Accordion("Paramètres avancés", open=True): | |
| confidence_slider = gr.Slider( | |
| minimum=0.05, | |
| maximum=0.95, | |
| value=0.4, | |
| step=0.05, | |
| label="Seuil de confiance", | |
| info="Plus bas = plus de détections" | |
| ) | |
| iou_slider = gr.Slider( | |
| minimum=0.1, | |
| maximum=0.9, | |
| value=0.5, | |
| step=0.05, | |
| label="Seuil", | |
| info="Plus haut = garde plus de boxes qui se chevauchent" | |
| ) | |
| show_conf_check = gr.Checkbox( | |
| value=True, | |
| label="Afficher les confiances détaillées" | |
| ) | |
| detect_btn = gr.Button("Détecter les objets", variant="primary", size="lg") | |
| with gr.Column(scale=2): | |
| with gr.Row(): | |
| original_display = gr.Image(type="pil", label="Image originale") | |
| output_image = gr.Image(type="pil", label="Résultat avec détections") | |
| output_stats = gr.Markdown("**Uploadez une image et cliquez sur 'Détecter' pour commencer !**") | |
| # Galerie d'exemples | |
| if EXAMPLE_IMAGES: | |
| gr.Markdown("### Exemples (cliquez pour tester)") | |
| examples_list = [[img, 0.4, 0.5, True] for img in EXAMPLE_IMAGES[:12]] | |
| gr.Examples( | |
| examples=examples_list, | |
| inputs=[input_image, confidence_slider, iou_slider, show_conf_check], | |
| outputs=[original_display, output_image, output_stats], | |
| fn=detect_objects, | |
| cache_examples=False, | |
| examples_per_page=6, | |
| ) | |
| # Actions | |
| detect_btn.click( | |
| fn=detect_objects, | |
| inputs=[input_image, confidence_slider, iou_slider, show_conf_check], | |
| outputs=[original_display, output_image, output_stats] | |
| ) | |
| input_image.change( | |
| fn=detect_objects, | |
| inputs=[input_image, confidence_slider, iou_slider, show_conf_check], | |
| outputs=[original_display, output_image, output_stats] | |
| ) | |
| # Onglet Info | |
| with gr.Tab("À propos"): | |
| mAP_info = f"{MODEL_INFO['mAP']:.4f}" if isinstance(MODEL_INFO['mAP'], (int, float)) else 'N/A' | |
| epoch_info = MODEL_INFO['epoch'] if MODEL_INFO['epoch'] != 'N/A' else 'N/A' | |
| # Lancer l'app | |
| if __name__ == "__main__": | |
| print("\n" + "="*60) | |
| print("Lancement de l'application Gradio YOLO v1") | |
| print("="*60) | |
| print(f"Modèle: {MODEL_PATH}") | |
| print(f"Device: {DEVICE}") | |
| print(f"Exemples chargés: {len(EXAMPLE_IMAGES)}") | |
| print("="*60 + "\n") | |
| demo.launch( | |
| share=True, | |
| server_name="0.0.0.0", # Accessible depuis le réseau local | |
| server_port=7860, | |
| show_error=True | |
| ) | |