Spaces:
Sleeping
Sleeping
File size: 3,273 Bytes
273ca83 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | 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()
|