naveenkm13/occupancyos

1"""2Run YOLOv8l + desk-zone occupancy inference on a video file.3Writes an annotated mp4 and a CSV log of per-frame occupancy.4 5Usage:6    python -m backend.scripts.video_inference --source path/to/video.mp4 \7        --output backend/output/annotated.mp48"""9from __future__ import annotations10 11import argparse12import csv13from pathlib import Path14 15import cv216 17from backend import config18from backend.detection.occupancy_detector import OccupancyDetector19from backend.detection.chair_occupancy_detector import ChairOccupancyDetector20from backend.detection.desk_state_detector import DeskStateDetector21from backend.streams.video_source import VideoSource22 23 24def run(source: str, output: str, csv_path: str | None = None,25        show: bool = False, max_frames: int | None = None) -> None:26    mode = config.DETECTOR_MODE.lower()27    if mode == "desk_state":28        detector = DeskStateDetector()29    elif mode == "chair":30        detector = ChairOccupancyDetector()31    else:32        detector = OccupancyDetector()33    print(f"[video_inference] detector mode = {mode}")34    src = VideoSource(source, loop=False, reconnect=False)35 36    fourcc = cv2.VideoWriter_fourcc(*"mp4v")37    out_path = Path(output)38    out_path.parent.mkdir(parents=True, exist_ok=True)39    writer = cv2.VideoWriter(str(out_path), fourcc, src.fps or 25.0,40                             (src.width, src.height))41 42    csv_file = None43    csv_writer = None44    if csv_path:45        csv_file = open(csv_path, "w", newline="", encoding="utf-8")46        csv_writer = csv.writer(csv_file)47        csv_writer.writerow(["frame", "people", "occupied", "total", "occupancy_pct", "fps"])48 49    print(f"[video_inference] {source} -> {out_path}")50    frame_idx = 051    try:52        while True:53            ok, frame = src.read()54            if not ok or frame is None:55                break56            frame_idx += 157            result = detector.detect(frame)58            annotated = detector.annotate(frame, result)59            writer.write(annotated)60            occupied = getattr(result, "occupied_chairs", None)61            total    = getattr(result, "total_chairs", None)62            if occupied is None:63                occupied = result.occupied_desks64                total    = result.total_desks65            if csv_writer:66                csv_writer.writerow([67                    frame_idx, result.total_people, occupied, total,68                    f"{result.occupancy_pct:.2f}", f"{result.fps:.2f}",69                ])70            if frame_idx % 50 == 0:71                print(f"  frame {frame_idx} | people={result.total_people} "72                      f"| occ={occupied}/{total} | util={result.occupancy_pct:.1f}% "73                      f"| fps={result.fps:.1f}")74            if show:75                cv2.imshow("Chair Occupancy", annotated)76                if cv2.waitKey(1) & 0xFF == ord("q"):77                    break78            if max_frames is not None and frame_idx >= max_frames:79                break80    finally:81        src.release()82        writer.release()83        if csv_file:84            csv_file.close()85        if show:86            cv2.destroyAllWindows()87    print(f"[video_inference] done. wrote {frame_idx} frames.")88 89 90if __name__ == "__main__":91    p = argparse.ArgumentParser()92    p.add_argument("--source", default=config.DEFAULT_VIDEO_SOURCE)93    p.add_argument("--output", default=str(config.OUTPUT_DIR / "chair_demo.mp4"))94    p.add_argument("--csv", default=str(config.OUTPUT_DIR / "occupancy_log.csv"))95    p.add_argument("--show", action="store_true", help="display annotated frames in a window")96    p.add_argument("--max-frames", type=int, default=None, help="stop after N frames")97    args = p.parse_args()98    run(args.source, args.output, args.csv, show=args.show, max_frames=args.max_frames)99