""" EagleVision CV Microservice – main.py ======================================== Entry point. Reads video → runs detection + tracking + optical flow → classifies activity → publishes to Kafka → writes to TimescaleDB. """ import os import sys import time import logging import argparse import cv2 import numpy as np from detector import EquipmentDetector from tracker_state import MachineStateTracker from flow_analyzer import OpticalFlowAnalyzer from classifier import ActivityClassifier from kafka_pub import KafkaPublisher from db_writer import DBWriter # ─────────────────── Logging ──────────────────────────────────────────────── logging.basicConfig( level = logging.INFO, format = "%(asctime)s [%(levelname)s] %(name)s – %(message)s", handlers=[logging.StreamHandler(sys.stdout)] ) log = logging.getLogger("cv_service") # ─────────────────── Config from env ──────────────────────────────────────── CONFIG = { "video_source" : os.getenv("VIDEO_SOURCE", "data/input.mp4"), "kafka_servers" : os.getenv("KAFKA_BOOTSTRAP_SERVERS", "localhost:9092"), "kafka_topic" : os.getenv("KAFKA_TOPIC", "equipment_events"), "yolo_model" : os.getenv("YOLO_MODEL", "yolov8m.pt"), "flow_threshold" : float(os.getenv("FLOW_THRESHOLD", "1.5")), "inactive_threshold" : float(os.getenv("INACTIVE_THRESHOLD", "0.5")), "db_host" : os.getenv("DB_HOST", "localhost"), "db_port" : int(os.getenv("DB_PORT", "5432")), "db_name" : os.getenv("DB_NAME", "eaglevision"), "db_user" : os.getenv("DB_USER", "eagle"), "db_password" : os.getenv("DB_PASSWORD", "eagle_secret"), "output_video" : os.getenv("OUTPUT_VIDEO", "data/output.mp4"), "target_fps" : int(os.getenv("TARGET_FPS", "15")), # process every N fps "display_local" : os.getenv("DISPLAY_LOCAL", "false").lower() == "true", } def build_equipment_id(cls_name: str, track_id: int) -> str: prefix = {"excavator": "EX", "dump_truck": "DT", "bulldozer": "BZ"}.get(cls_name, "EQ") return f"{prefix}-{track_id:03d}" def run_pipeline(config: dict) -> None: log.info("🦅 EagleVision CV Service starting …") # ── Init components ─────────────────────────────────────────────────── detector = EquipmentDetector(config["yolo_model"]) flow_anal = OpticalFlowAnalyzer( flow_threshold = config["flow_threshold"], inactive_threshold = config["inactive_threshold"], ) classifier = ActivityClassifier() state_mgr = MachineStateTracker() publisher = KafkaPublisher(config["kafka_servers"], config["kafka_topic"]) db_writer = DBWriter(config) # ── Open video ──────────────────────────────────────────────────────── cap = cv2.VideoCapture(config["video_source"]) if not cap.isOpened(): log.error(f"Cannot open video: {config['video_source']}") sys.exit(1) orig_fps = cap.get(cv2.CAP_PROP_FPS) or 30 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) total_f = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) log.info(f"Video: {width}×{height} @ {orig_fps:.1f}fps total_frames={total_f}") # ── Output writer ───────────────────────────────────────────────────── out_writer = None if config["output_video"]: fourcc = cv2.VideoWriter_fourcc(*"mp4v") out_writer = cv2.VideoWriter(config["output_video"], fourcc, orig_fps, (width, height)) # ── Frame-skip factor ───────────────────────────────────────────────── skip = max(1, int(orig_fps / config["target_fps"])) prev_gray = None frame_id = 0 t_start = time.time() try: while True: ret, frame = cap.read() if not ret: log.info("End of video stream.") break frame_id += 1 if frame_id % skip != 0: # still write original frame to output for smooth playback if out_writer: out_writer.write(frame) continue # ── Gray for optical flow ───────────────────────────────────── gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # ── Detect + Track ──────────────────────────────────────────── detections = detector.detect(frame) # list of Detection objects # ── Compute dense optical flow (full frame) ─────────────────── flow = None if prev_gray is not None: flow = flow_anal.compute_flow(prev_gray, gray) annotated = frame.copy() for det in detections: # ── Regional flow analysis ──────────────────────────────── flow_upper, flow_lower, motion_source = flow_anal.analyze_region( flow, det.bbox, height ) if flow is not None else (0.0, 0.0, "none") # ── State: ACTIVE / INACTIVE ────────────────────────────── state = "ACTIVE" if motion_source != "none" else "INACTIVE" # ── Activity classification ─────────────────────────────── activity = classifier.classify( track_id = det.track_id, bbox = det.bbox, flow_upper = flow_upper, flow_lower = flow_lower, motion_source = motion_source, ) # ── Time accumulation ───────────────────────────────────── eq_id = build_equipment_id(det.cls_name, det.track_id) dt_secs = skip / orig_fps stats = state_mgr.update(eq_id, state, dt_secs) # ── Build Kafka payload ─────────────────────────────────── video_ts = frame_id / orig_fps h = int(video_ts // 3600) m = int((video_ts % 3600) // 60) s = video_ts % 60 ts = f"{h:02d}:{m:02d}:{s:06.3f}" payload = { "frame_id" : frame_id, "equipment_id" : eq_id, "equipment_class": det.cls_name, "track_id" : det.track_id, "timestamp" : ts, "utilization" : { "current_state" : state, "current_activity": activity, "motion_source" : motion_source, }, "flow_metrics" : { "flow_upper": round(flow_upper, 3), "flow_lower": round(flow_lower, 3), }, "bbox" : { "x1": det.bbox[0], "y1": det.bbox[1], "x2": det.bbox[2], "y2": det.bbox[3], }, "time_analytics" : { "total_tracked_seconds" : round(stats["total_tracked"], 2), "total_active_seconds" : round(stats["total_active"], 2), "total_idle_seconds" : round(stats["total_idle"], 2), "utilization_percent" : round(stats["utilization_pct"], 1), }, } # ── Publish to Kafka ────────────────────────────────────── publisher.publish(payload) # ── Write to DB ─────────────────────────────────────────── db_writer.insert(payload) # ── Annotate frame ──────────────────────────────────────── annotated = draw_annotations(annotated, det, state, activity, motion_source, stats) # ── Write output frame ──────────────────────────────────────── if out_writer: out_writer.write(annotated) if config["display_local"]: cv2.imshow("EagleVision", annotated) if cv2.waitKey(1) & 0xFF == ord("q"): break prev_gray = gray if frame_id % 100 == 0: elapsed = time.time() - t_start log.info(f"Frame {frame_id}/{total_f} – elapsed {elapsed:.1f}s – " f"active machines: {state_mgr.active_count()}") finally: cap.release() if out_writer: out_writer.release() cv2.destroyAllWindows() publisher.close() db_writer.close() log.info("✅ Pipeline finished.") # ─────────────────── Annotation helpers ───────────────────────────────────── COLORS = { "ACTIVE" : (0, 200, 100), "INACTIVE": (0, 100, 220), } ACTIVITY_ICONS = { "DIGGING" : "⛏", "SWINGING/LOADING": "🔄", "DUMPING" : "📤", "WAITING" : "⏸", } def draw_annotations(frame, det, state, activity, motion_source, stats): x1, y1, x2, y2 = det.bbox color = COLORS.get(state, (200, 200, 200)) # BBox cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2) # Header label label = f"{det.cls_name.upper()} #{det.track_id} | {state}" lw, lh = cv2.getTextSize(label, cv2.FONT_HERSHEY_DUPLEX, 0.55, 1)[0] cv2.rectangle(frame, (x1, y1 - lh - 10), (x1 + lw + 8, y1), color, -1) cv2.putText(frame, label, (x1 + 4, y1 - 4), cv2.FONT_HERSHEY_DUPLEX, 0.55, (255, 255, 255), 1) # Activity + motion source sub = f"{activity} [{motion_source}]" cv2.putText(frame, sub, (x1 + 4, y1 + 18), cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 1) # Utilization strip at bottom of box util = stats["utilization_pct"] bar_w = x2 - x1 filled = int(bar_w * util / 100) cv2.rectangle(frame, (x1, y2 + 2), (x2, y2 + 10), (50, 50, 50), -1) cv2.rectangle(frame, (x1, y2 + 2), (x1 + filled, y2 + 10), color, -1) util_txt = f"{util:.1f}% util | A:{stats['total_active']:.0f}s I:{stats['total_idle']:.0f}s" cv2.putText(frame, util_txt, (x1, y2 + 22), cv2.FONT_HERSHEY_SIMPLEX, 0.38, color, 1) return frame # ─────────────────── CLI ──────────────────────────────────────────────────── if __name__ == "__main__": parser = argparse.ArgumentParser(description="EagleVision CV Service") parser.add_argument("--video", default=CONFIG["video_source"]) parser.add_argument("--display", action="store_true") args = parser.parse_args() CONFIG["video_source"] = args.video CONFIG["display_local"] = args.display run_pipeline(CONFIG)