Spaces:
Sleeping
Sleeping
File size: 3,216 Bytes
ee579d0 | 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 cv2
import numpy as np
from ultralytics import YOLO
# ============================================================
# Cargar YOLO
# ============================================================
WEIGHTS_PATH = "yolov8n.pt" # el modelo que ya te funcionaba
print(f"[INFO] Cargando modelo YOLO: {WEIGHTS_PATH}")
model = YOLO(WEIGHTS_PATH)
print("[INFO] Modelo YOLO cargado correctamente.")
# ============================================================
# Preprocesado básico
# ============================================================
def preprocess_frame(frame):
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
blur = cv2.GaussianBlur(gray, (5, 5), 0)
# Detecta cambios bruscos de intensidad (útil para bordes de objetos alargados)
sobel = cv2.Sobel(blur, cv2.CV_64F, 1, 1, ksize=3)
sobel = cv2.convertScaleAbs(sobel)
# Filtro de bordes más duro
edges = cv2.Canny(blur, 80, 150)
# Pequeña expansión de bordes
kernel = np.ones((3, 3), np.uint8)
dilated = cv2.dilate(edges, kernel, iterations=1)
# Cierra huecos pequeños en formas
morph = cv2.morphologyEx(dilated, cv2.MORPH_CLOSE, kernel, iterations=1)
return gray, sobel, morph
# ============================================================
# Detección usando YOLO
# ============================================================
def detect_knives(frame):
results = model(frame, verbose=False)[0]
detections = []
for box in results.boxes:
cls = int(box.cls[0])
label = results.names[cls].lower()
# Verifica si YOLO considera la clase como cuchillo
if "knife" in label or "cuchillo" in label:
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
detections.append((int(x1), int(y1), int(x2), int(y2)))
return detections
# ============================================================
# Lógica principal
# ============================================================
def main():
print("[INFO] Iniciando cámara...")
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("[ERROR] No se pudo abrir la cámara.")
return
while True:
ret, frame = cap.read()
if not ret:
print("[ERROR] No se pudo leer el frame.")
break
gray, sobel, morph = preprocess_frame(frame)
detections = detect_knives(frame)
# Dibuja detecciones
for (x1, y1, x2, y2) in detections:
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
cv2.putText(frame, "KNIFE", (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
# Ventana principal
cv2.imshow("Detección de Cuchillos (YOLO)", frame)
# ------- Ventana extra con 2 filtros juntos -------
filters = np.hstack((gray, sobel))
cv2.imshow("Filtros (Gray | Sobel)", filters)
# Salir con Q
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
print("[INFO] Cámara finalizada.")
if __name__ == "__main__":
main()
|