Spaces:
Sleeping
Sleeping
| 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() | |