Spaces:
Sleeping
Sleeping
| import os | |
| from datasets import load_dataset | |
| from ultralytics import YOLO | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score | |
| # ------------------------------------------------------------------ | |
| # CONFIG | |
| # ------------------------------------------------------------------ | |
| MODEL_PATH = "weapons_yolo.pt" | |
| METRICS_DIR = "Metrics" | |
| os.makedirs(METRICS_DIR, exist_ok=True) | |
| print("[INFO] Cargando modelo...") | |
| model = YOLO(MODEL_PATH) | |
| # ------------------------------------------------------------------ | |
| # EVALUACIÓN | |
| # ------------------------------------------------------------------ | |
| def evaluate_model(): | |
| print("[INFO] Descargando dataset...") | |
| ds = load_dataset("cerealkiller2527/pick_knife") | |
| images = ds["train"] | |
| y_true = [] | |
| y_pred = [] | |
| total = len(images) | |
| print(f"[INFO] Total de imágenes: {total}") | |
| for i, item in enumerate(images): | |
| img = item["image"] | |
| label = int(item["label"]) # 0 = no knife, 1 = knife | |
| results = model.predict(img, verbose=False)[0] | |
| detected_knife = False | |
| for box in results.boxes: | |
| cls = int(box.cls[0]) | |
| name = results.names[cls].lower() | |
| if "knife" in name: | |
| detected_knife = True | |
| break | |
| y_true.append(label) | |
| y_pred.append(1 if detected_knife else 0) | |
| if i % 50 == 0: | |
| print(f"[INFO] Procesadas {i}/{total} imágenes...") | |
| # ------------------------------------------------------------------ | |
| # Métricas | |
| # ------------------------------------------------------------------ | |
| precision = precision_score(y_true, y_pred) | |
| recall = recall_score(y_true, y_pred) | |
| f1 = f1_score(y_true, y_pred) | |
| cm = confusion_matrix(y_true, y_pred) | |
| print("\n===== RESULTADOS =====") | |
| print(f"Precisión: {precision:.4f}") | |
| print(f"Sensibilidad (Recall): {recall:.4f}") | |
| print(f"F1 Score: {f1:.4f}") | |
| print("Confusion Matrix:") | |
| print(cm) | |
| # Guardamos métricas en TXT | |
| with open(os.path.join(METRICS_DIR, "metrics.txt"), "w") as f: | |
| f.write(f"Precisión: {precision:.4f}\n") | |
| f.write(f"Recall: {recall:.4f}\n") | |
| f.write(f"F1 Score: {f1:.4f}\n") | |
| f.write(f"Matriz de confusión:\n{cm}\n") | |
| # ------------------------------------------------------------------ | |
| # GRÁFICAS | |
| # ------------------------------------------------------------------ | |
| # Confusion matrix plot | |
| plt.figure(figsize=(5, 4)) | |
| plt.imshow(cm, cmap="Blues") | |
| plt.title("Matriz de Confusión") | |
| plt.xlabel("Predicción") | |
| plt.ylabel("Valor Real") | |
| plt.colorbar() | |
| plt.savefig(os.path.join(METRICS_DIR, "confusion_matrix.png")) | |
| plt.close() | |
| # Precision-Recall (simple) | |
| plt.figure() | |
| plt.bar(["Precisión", "Recall", "F1"], [precision, recall, f1]) | |
| plt.ylim(0, 1) | |
| plt.title("Reporte de Métricas") | |
| plt.savefig(os.path.join(METRICS_DIR, "metrics_barplot.png")) | |
| plt.close() | |
| print(f"\n[INFO] Gráficas guardadas en: {METRICS_DIR}/") | |
| print("[INFO] Evaluación completa.") | |
| # ------------------------------------------------------------------ | |
| if __name__ == "__main__": | |
| evaluate_model() | |