CoolFace
Apppublic

naveenkm13/occupancyos

sourceHugging Faceupdated 2mo agoView on Hugging Face
0likes
debug_raw.py105 linesDownload Raw Back to scripts
1"""2Diagnostic — draws RAW YOLOv8l detections (person + chair) on a video,3bypassing all tracking, filtering, and occupancy logic. Use this to see4what the underlying model is actually detecting before tuning.5 6Usage:7    python -m backend.scripts.debug_raw --source samples\office.mp4 --show8    python -m backend.scripts.debug_raw --source 0 --show9"""10from __future__ import annotations11 12import argparse13from pathlib import Path14 15import cv216import torch17from ultralytics import YOLO18 19from backend import config20from backend.streams.video_source import VideoSource21 22 23def run(source: str, output: str, show: bool, max_frames: int | None,24        conf: float, imgsz: int) -> None:25    device = config.DEVICE if torch.cuda.is_available() else "cpu"26    model = YOLO(config.YOLO_WEIGHTS)27    model.to(device)28    print(f"[debug_raw] weights={config.YOLO_WEIGHTS}  device={device}  "29          f"imgsz={imgsz}  conf={conf}")30    print(f"[debug_raw] person id={config.PERSON_CLASS_ID}  "31          f"chair id={config.CHAIR_CLASS_ID}")32 33    src = VideoSource(source, loop=False, reconnect=False)34    out_path = Path(output)35    out_path.parent.mkdir(parents=True, exist_ok=True)36    writer = cv2.VideoWriter(37        str(out_path), cv2.VideoWriter_fourcc(*"mp4v"),38        src.fps or 25.0, (src.width, src.height),39    )40 41    frame_idx = 042    try:43        while True:44            ok, frame = src.read()45            if not ok or frame is None:46                break47            frame_idx += 148 49            res = model.predict(50                source=frame, conf=conf, iou=config.IOU_THRESHOLD,51                classes=[config.PERSON_CLASS_ID, config.CHAIR_CLASS_ID],52                imgsz=imgsz, device=device, verbose=False,53                augment=getattr(config, "TTA", False),54            )55 56            n_p = n_c = 057            if res and res[0].boxes is not None:58                b = res[0].boxes59                xyxy  = b.xyxy.cpu().numpy().astype(int)60                confs = b.conf.cpu().numpy().astype(float)61                clss  = b.cls.cpu().numpy().astype(int)62                for (x1, y1, x2, y2), cf, k in zip(xyxy, confs, clss):63                    if k == config.PERSON_CLASS_ID:64                        color = (0, 220, 255); label = f"person {cf:.2f}"; n_p += 165                    elif k == config.CHAIR_CLASS_ID:66                        color = (60, 220, 120); label = f"chair {cf:.2f}";  n_c += 167                    else:68                        continue69                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)70                    cv2.putText(frame, label, (x1, max(15, y1 - 6)),71                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)72 73            hud = f"RAW YOLO | people={n_p}  chairs={n_c}  conf>={conf}  imgsz={imgsz}"74            cv2.rectangle(frame, (0, 0), (frame.shape[1], 32), (15, 22, 35), -1)75            cv2.putText(frame, hud, (10, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.6,76                        (180, 240, 255), 2)77 78            writer.write(frame)79            if show:80                cv2.imshow("DEBUG raw YOLO", frame)81                if cv2.waitKey(1) & 0xFF == ord("q"):82                    break83            if max_frames is not None and frame_idx >= max_frames:84                break85            if frame_idx % 50 == 0:86                print(f"  frame {frame_idx}  people={n_p}  chairs={n_c}")87    finally:88        src.release()89        writer.release()90        if show:91            cv2.destroyAllWindows()92    print(f"[debug_raw] done. wrote {frame_idx} frames -> {out_path}")93 94 95if __name__ == "__main__":96    p = argparse.ArgumentParser()97    p.add_argument("--source", default=config.DEFAULT_VIDEO_SOURCE)98    p.add_argument("--output", default=str(config.OUTPUT_DIR / "debug_raw.mp4"))99    p.add_argument("--show", action="store_true")100    p.add_argument("--max-frames", type=int, default=None)101    p.add_argument("--conf", type=float, default=0.25)102    p.add_argument("--imgsz", type=int, default=config.IMG_SIZE)103    args = p.parse_args()104    run(args.source, args.output, args.show, args.max_frames, args.conf, args.imgsz)105