486 lines
17 KiB
Python
486 lines
17 KiB
Python
"""
|
|
Batch Video Cropper — Rekam Video RTSP per Sesi Batch Truk
|
|
|
|
Program ini membaca livestream CCTV (RTSP), menjalankan algoritma penentuan batch
|
|
menggunakan YOLO + State Machine, dan menyimpan potongan video per-batch ke folder
|
|
yang terorganisir berdasarkan tanggal.
|
|
|
|
Struktur Output:
|
|
~/reTraining/data/archive/
|
|
├── 2026-08-05/
|
|
│ ├── batch_1_09-15-30.mp4
|
|
│ ├── batch_2_10-22-45.mp4
|
|
│ └── batch_3_14-08-12.mp4
|
|
└── 2026-08-06/
|
|
└── batch_1_07-30-00.mp4
|
|
|
|
Menjalankan:
|
|
cd ~/reTraining/algoritma-batch
|
|
python3 batch_video_cropper.py
|
|
"""
|
|
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
import time
|
|
import json
|
|
import threading
|
|
from datetime import datetime, timedelta
|
|
from shapely.geometry import Point, Polygon
|
|
from ultralytics import YOLO
|
|
|
|
# Import modules
|
|
from src.tracking import ByteTrackTracker
|
|
from src.stabilizer import BboxStabilizer
|
|
from src.truck_roi import TruckROI
|
|
from src.counting import LineCrossCounter
|
|
from src.batch import BatchLifecycleManager, BatchRecord
|
|
|
|
# =====================================================================
|
|
# 1. KONFIGURASI
|
|
# =====================================================================
|
|
|
|
import platform
|
|
IS_WINDOWS = platform.system() == "Windows"
|
|
|
|
# --- Path Konfigurasi ---
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
ZONES_JSON = os.path.join(BASE_DIR, "zones.json")
|
|
|
|
if IS_WINDOWS:
|
|
MODEL_PATH = os.path.join(BASE_DIR, "v1-best.pt")
|
|
ARCHIVE_BASE = os.path.join(BASE_DIR, "archive_output")
|
|
else:
|
|
MODEL_PATH = os.path.join(BASE_DIR, "v1-best.pt")
|
|
ARCHIVE_BASE = os.path.expanduser("~/reTraining/data/archive")
|
|
|
|
# --- Sumber Video RTSP ---
|
|
RTSP_URL = "rtsp://frigate:zenai@192.168.192.209:8554/camera_stream_640"
|
|
|
|
# --- Batas Pergantian Hari (Cutoff) ---
|
|
DAILY_CUTOFF_TIME = "20:00"
|
|
|
|
# --- Parameter State Machine ---
|
|
SACK_IDLE_TIMEOUT = 5.0 # Jeda aktivitas sebelum masuk WAITING_FOR_ACTIVITY
|
|
MIN_BATCH_DURATION = 2.0 # Durasi minimal batch sebelum boleh masuk WAITING
|
|
TOLERANCE_LOW_COUNT = 0.0 # Instan (0s) — batch langsung berakhir saat area truk kosong
|
|
TOLERANCE_MED_COUNT = 0.0
|
|
TOLERANCE_HIGH_COUNT = 0.0
|
|
|
|
# --- Video Recording ---
|
|
VIDEO_FPS = 10.0 # FPS output video (10 fps sudah cukup untuk rekaman arsip)
|
|
VIDEO_CODEC = "mp4v" # Codec untuk .mp4
|
|
|
|
|
|
# =====================================================================
|
|
# 2. THREADED RTSP READER (Menghindari Lag Buffer)
|
|
# =====================================================================
|
|
class RTSPStreamReader:
|
|
"""Threaded RTSP reader yang selalu mengambil frame terbaru."""
|
|
|
|
def __init__(self, source_url):
|
|
self.source_url = source_url
|
|
self.cap = cv2.VideoCapture(source_url)
|
|
self.frame = None
|
|
self.ret = False
|
|
self.running = True
|
|
self.lock = threading.Lock()
|
|
self.new_frame_event = threading.Event()
|
|
self.thread = threading.Thread(target=self._update, daemon=True)
|
|
self.thread.start()
|
|
|
|
def _update(self):
|
|
while self.running:
|
|
if not self.cap.isOpened():
|
|
print("[RTSP] Stream terputus, mencoba reconnect dalam 5 detik...")
|
|
time.sleep(5)
|
|
self.cap = cv2.VideoCapture(self.source_url)
|
|
continue
|
|
ret, frame = self.cap.read()
|
|
if not ret:
|
|
time.sleep(0.01)
|
|
continue
|
|
with self.lock:
|
|
self.ret = ret
|
|
self.frame = frame
|
|
self.new_frame_event.set()
|
|
time.sleep(0.001)
|
|
|
|
def read(self):
|
|
if self.new_frame_event.wait(timeout=2.0):
|
|
self.new_frame_event.clear()
|
|
with self.lock:
|
|
if self.frame is None:
|
|
return False, None
|
|
return self.ret, self.frame.copy()
|
|
else:
|
|
with self.lock:
|
|
if self.frame is None:
|
|
return False, None
|
|
return self.ret, self.frame.copy()
|
|
|
|
def isOpened(self):
|
|
return self.cap.isOpened()
|
|
|
|
def release(self):
|
|
self.running = False
|
|
if self.cap.isOpened():
|
|
self.cap.release()
|
|
|
|
|
|
# =====================================================================
|
|
# 3. FUNGSI UTILITAS TANGGAL & FOLDER
|
|
# =====================================================================
|
|
def get_counting_date(dt=None):
|
|
"""Menentukan tanggal kerja berdasarkan cutoff harian."""
|
|
if dt is None:
|
|
dt = datetime.now()
|
|
try:
|
|
cutoff = datetime.strptime(DAILY_CUTOFF_TIME, "%H:%M").time()
|
|
except Exception:
|
|
cutoff = datetime.strptime("20:00", "%H:%M").time()
|
|
|
|
if cutoff.hour == 0 and cutoff.minute == 0:
|
|
return dt.date().isoformat()
|
|
|
|
if dt.time() < cutoff:
|
|
return dt.date().isoformat()
|
|
return (dt.date() + timedelta(days=1)).isoformat()
|
|
|
|
|
|
def ensure_date_folder(counting_date):
|
|
"""Membuat folder tanggal di archive jika belum ada. Mengembalikan path folder."""
|
|
folder = os.path.join(ARCHIVE_BASE, counting_date)
|
|
os.makedirs(folder, exist_ok=True)
|
|
return folder
|
|
|
|
|
|
# =====================================================================
|
|
# 4. BATCH VIDEO RECORDER (Mengelola VideoWriter per Batch)
|
|
# =====================================================================
|
|
class BatchVideoRecorder:
|
|
"""Mengelola pembukaan dan penutupan file video per sesi batch."""
|
|
|
|
def __init__(self, archive_base, video_fps=10.0, codec="mp4v"):
|
|
self.archive_base = archive_base
|
|
self.video_fps = video_fps
|
|
self.codec = codec
|
|
self.writer = None
|
|
self.current_path = None
|
|
self.frame_count = 0
|
|
|
|
def start_recording(self, batch_number, counting_date, frame_width=1280, frame_height=720):
|
|
"""Membuka file video baru untuk batch ini."""
|
|
self.stop_recording() # Pastikan writer sebelumnya ditutup
|
|
|
|
folder = ensure_date_folder(counting_date)
|
|
timestamp_str = datetime.now().strftime("%H-%M-%S")
|
|
filename = f"batch_{batch_number}_{timestamp_str}.mp4"
|
|
self.current_path = os.path.join(folder, filename)
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*self.codec)
|
|
self.writer = cv2.VideoWriter(
|
|
self.current_path, fourcc, self.video_fps, (frame_width, frame_height)
|
|
)
|
|
self.frame_count = 0
|
|
|
|
if self.writer.isOpened():
|
|
print(f"[RECORD] Mulai merekam video batch #{batch_number} -> {self.current_path}")
|
|
else:
|
|
print(f"[RECORD ERROR] Gagal membuka VideoWriter untuk: {self.current_path}")
|
|
self.writer = None
|
|
|
|
def write_frame(self, frame):
|
|
"""Menulis satu frame ke video aktif."""
|
|
if self.writer is not None and self.writer.isOpened():
|
|
self.writer.write(frame)
|
|
self.frame_count += 1
|
|
|
|
def stop_recording(self):
|
|
"""Menutup file video yang sedang aktif."""
|
|
if self.writer is not None:
|
|
self.writer.release()
|
|
self.writer = None
|
|
if self.current_path and self.frame_count > 0:
|
|
print(f"[RECORD] Video selesai disimpan: {self.current_path} ({self.frame_count} frames)")
|
|
elif self.current_path and self.frame_count == 0:
|
|
# Hapus file kosong
|
|
try:
|
|
os.remove(self.current_path)
|
|
print(f"[RECORD] File video kosong dihapus: {self.current_path}")
|
|
except Exception:
|
|
pass
|
|
self.current_path = None
|
|
self.frame_count = 0
|
|
|
|
@property
|
|
def is_recording(self):
|
|
return self.writer is not None and self.writer.isOpened()
|
|
|
|
|
|
# =====================================================================
|
|
# 5. FUNGSI UTAMA — LOOP UTAMA DETEKSI & PEREKAMAN
|
|
# =====================================================================
|
|
def run_batch_video_cropper():
|
|
print("=" * 60)
|
|
print(" BATCH VIDEO CROPPER — RTSP → Per-Batch MP4 Recorder")
|
|
print("=" * 60)
|
|
print(f" Model : {MODEL_PATH}")
|
|
print(f" RTSP : {RTSP_URL}")
|
|
print(f" Archive : {ARCHIVE_BASE}")
|
|
print(f" Cutoff : {DAILY_CUTOFF_TIME}")
|
|
print(f" Tolerance : Instan (0s)")
|
|
print("=" * 60)
|
|
|
|
# Pastikan folder archive ada
|
|
os.makedirs(ARCHIVE_BASE, exist_ok=True)
|
|
|
|
# --- Load Model YOLO ---
|
|
print("[INFO] Memuat model YOLO...")
|
|
model = YOLO(MODEL_PATH)
|
|
|
|
# Warm-up model
|
|
print("[INFO] Warm-up model...")
|
|
dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
|
|
device = "cuda" if os.path.exists("/usr/local/cuda") else "cpu"
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
device = "cuda"
|
|
except ImportError:
|
|
pass
|
|
_ = model(dummy, imgsz=640, device=device, verbose=False)
|
|
print(f"[INFO] Model siap. Device: {device}")
|
|
|
|
# --- Setup Components ---
|
|
tracker = ByteTrackTracker(model, conf=0.55)
|
|
stabilizer = BboxStabilizer(
|
|
ema_alpha=0.35,
|
|
max_hold_frames=10,
|
|
max_height_ratio=1.5,
|
|
min_height_ratio=0.70,
|
|
)
|
|
|
|
# Skala koordinat (kalibrasi 1920x1080 -> 1280x720)
|
|
scale_x = 1280.0 / 1920.0
|
|
scale_y = 720.0 / 1080.0
|
|
|
|
# Detection polygon
|
|
detection_poly_pts = [
|
|
[int(574 * scale_x), int(50 * scale_y)],
|
|
[int(586 * scale_x), int(1077 * scale_y)],
|
|
[int(1418 * scale_x), int(1076 * scale_y)],
|
|
[int(1397 * scale_x), int(50 * scale_y)],
|
|
]
|
|
detection_polygon = Polygon(detection_poly_pts)
|
|
|
|
# Truck polygon (untuk monitoring kehadiran karung)
|
|
truck_poly_pts = [
|
|
[int(600 * scale_x), int(385 * scale_y)],
|
|
[int(609 * scale_x), int(1076 * scale_y)],
|
|
[int(1404 * scale_x), int(1078 * scale_y)],
|
|
[int(1381 * scale_x), int(343 * scale_y)],
|
|
]
|
|
truck_polygon = Polygon(truck_poly_pts)
|
|
|
|
# Line crossing
|
|
static_line_y = int(330 * scale_y)
|
|
static_line_x_start = int(577 * scale_x)
|
|
static_line_x_end = int(1401 * scale_x)
|
|
|
|
static_roi = TruckROI(
|
|
x1=int(600 * scale_x),
|
|
y1=int(343 * scale_y),
|
|
x2=int(1404 * scale_x),
|
|
y2=int(1078 * scale_y),
|
|
line_y=static_line_y,
|
|
confidence=1.0,
|
|
)
|
|
|
|
counter = LineCrossCounter(
|
|
line_y=static_line_y,
|
|
line_x_start=static_line_x_start,
|
|
line_x_end=static_line_x_end,
|
|
margin=20,
|
|
dedup_radius=60.0,
|
|
)
|
|
|
|
batch_mgr = BatchLifecycleManager(
|
|
stabilize_seconds=0.0,
|
|
stabilize_threshold_px=9999.0,
|
|
sack_idle_timeout=SACK_IDLE_TIMEOUT,
|
|
min_batch_duration=MIN_BATCH_DURATION,
|
|
truck_gone_tolerance=TOLERANCE_LOW_COUNT,
|
|
)
|
|
|
|
recorder = BatchVideoRecorder(
|
|
archive_base=ARCHIVE_BASE,
|
|
video_fps=VIDEO_FPS,
|
|
codec=VIDEO_CODEC,
|
|
)
|
|
|
|
# --- Buka RTSP Stream ---
|
|
print(f"[INFO] Membuka RTSP stream: {RTSP_URL}")
|
|
cap = RTSPStreamReader(RTSP_URL)
|
|
if not cap.isOpened():
|
|
print(f"[ERROR] Gagal membuka RTSP stream: {RTSP_URL}")
|
|
return
|
|
|
|
# Tracking state
|
|
batch_counter = 0
|
|
frame_idx = 0
|
|
last_fps_time = time.time()
|
|
fps_counter = 0
|
|
|
|
print("\n[INFO] Memulai loop utama... Tekan Ctrl+C untuk berhenti.\n")
|
|
|
|
try:
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret or frame is None:
|
|
time.sleep(0.01)
|
|
continue
|
|
|
|
# Resize ke 1280x720 (sesuai kalibrasi koordinat zona)
|
|
frame = cv2.resize(frame, (1280, 720))
|
|
timestamp = time.time()
|
|
frame_idx += 1
|
|
|
|
# FPS counter
|
|
fps_counter += 1
|
|
if fps_counter % 100 == 0:
|
|
elapsed = time.time() - last_fps_time
|
|
fps = 100.0 / elapsed if elapsed > 0 else 0
|
|
last_fps_time = time.time()
|
|
state_str = batch_mgr.state
|
|
rec_str = "REC" if recorder.is_recording else "---"
|
|
print(f"[FPS] {fps:.1f} fps | State: {state_str} | {rec_str} | Frames: {frame_idx}")
|
|
|
|
# Simpan state sebelum update
|
|
prev_active = batch_mgr.is_active
|
|
prev_state = batch_mgr.state
|
|
|
|
# --- 1. YOLO Tracking ---
|
|
raw_tracked_all = tracker.update(frame, [])
|
|
raw_tracked_sacks = [d for d in raw_tracked_all if d.class_name == "sack"]
|
|
|
|
# --- 2. Stabilizer ---
|
|
stable = stabilizer.update(raw_tracked_sacks)
|
|
|
|
# --- 3. Filter Detection Polygon ---
|
|
stable = [
|
|
d for d in stable
|
|
if detection_polygon.contains(
|
|
Point((d.bbox[0] + d.bbox[2]) / 2.0, (d.bbox[1] + d.bbox[3]) / 2.0)
|
|
)
|
|
]
|
|
|
|
# --- 4. Hitung karung di 70% area bawah truk ---
|
|
min_ty, max_ty = truck_polygon.bounds[1], truck_polygon.bounds[3]
|
|
truck_height = max_ty - min_ty
|
|
truck_cutoff_y = min_ty + 0.30 * truck_height
|
|
|
|
sacks_in_truck_area = 0
|
|
for d in stable:
|
|
cx = (d.bbox[0] + d.bbox[2]) / 2.0
|
|
cy = (d.bbox[1] + d.bbox[3]) / 2.0
|
|
if truck_polygon.contains(Point(cx, cy)) and cy >= truck_cutoff_y:
|
|
sacks_in_truck_area += 1
|
|
|
|
# --- 5. Line Crossing ---
|
|
tracked_in_roi = [
|
|
d for d in stable
|
|
if static_roi.contains_x((d.bbox[0] + d.bbox[2]) / 2.0)
|
|
]
|
|
events = counter.update(tracked_in_roi)
|
|
has_crossing = len(events) > 0
|
|
|
|
# =============================================================
|
|
# LOGIKA ALGORITMA PENENTUAN BATCH (STATE MACHINE)
|
|
# =============================================================
|
|
|
|
# A. Mulai Batch
|
|
if batch_mgr.state in ("IDLE", "TRUCK_STABILIZING"):
|
|
batch_mgr.update_truck(has_crossing, (0.0, 0.0), timestamp)
|
|
|
|
# B. Monitoring Batch Aktif
|
|
if batch_mgr.state in ("COUNTING_SACKS", "WAITING_FOR_ACTIVITY"):
|
|
current_count = counter.loading_count
|
|
if current_count < 20:
|
|
batch_mgr._truck_gone_tolerance = TOLERANCE_LOW_COUNT
|
|
elif current_count >= 40:
|
|
batch_mgr._truck_gone_tolerance = TOLERANCE_HIGH_COUNT
|
|
else:
|
|
batch_mgr._truck_gone_tolerance = TOLERANCE_MED_COUNT
|
|
|
|
batch_mgr.update_sacks(
|
|
has_crossing_event=has_crossing,
|
|
sacks_in_area_count=sacks_in_truck_area,
|
|
timestamp=timestamp,
|
|
loading_count=counter.loading_count,
|
|
unloading_count=counter.unloading_count,
|
|
)
|
|
|
|
if batch_mgr.state == "WAITING_FOR_ACTIVITY":
|
|
batch_mgr.update_truck(sacks_in_truck_area > 0, None, timestamp)
|
|
|
|
for ev in events:
|
|
now_str = datetime.now().strftime("%H:%M:%S")
|
|
print(f"[{now_str}] [KARUNG] #{ev['track_id']} melintasi garis. Total: {counter.loading_count}")
|
|
|
|
# =============================================================
|
|
# TRANSISI BATCH — MULAI/SELESAI REKAMAN VIDEO
|
|
# =============================================================
|
|
|
|
# C. Batch baru saja dimulai
|
|
if batch_mgr.is_active and not prev_active:
|
|
counting_date = get_counting_date()
|
|
batch_counter += 1
|
|
now_str = datetime.now().strftime("%H:%M:%S")
|
|
print(f"\n>>> [{now_str}] BATCH #{batch_counter} DIMULAI (tanggal: {counting_date}) <<<")
|
|
recorder.start_recording(batch_counter, counting_date, 1280, 720)
|
|
|
|
# D. Batch baru saja selesai
|
|
elif not batch_mgr.is_active and prev_active:
|
|
final_count = counter.loading_count
|
|
now_str = datetime.now().strftime("%H:%M:%S")
|
|
print(f"\n>>> [{now_str}] BATCH #{batch_counter} SELESAI. Total karung: {final_count} <<<")
|
|
recorder.stop_recording()
|
|
|
|
# Reset counter dan stabilizer untuk batch berikutnya
|
|
counter.reset()
|
|
stabilizer.reset()
|
|
|
|
# E. Log transisi status
|
|
if batch_mgr.state != prev_state:
|
|
now_str = datetime.now().strftime("%H:%M:%S")
|
|
print(f"[{now_str}] [STATE] {prev_state} -> {batch_mgr.state}")
|
|
|
|
# =============================================================
|
|
# TULIS FRAME KE VIDEO (jika batch aktif)
|
|
# =============================================================
|
|
if batch_mgr.is_active and recorder.is_recording:
|
|
recorder.write_frame(frame)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n\n[INFO] Program dihentikan oleh pengguna (Ctrl+C).")
|
|
|
|
finally:
|
|
# Tutup rekaman yang masih terbuka
|
|
if recorder.is_recording:
|
|
print("[INFO] Menyimpan rekaman batch terakhir...")
|
|
recorder.stop_recording()
|
|
|
|
cap.release()
|
|
|
|
print("\n" + "=" * 60)
|
|
print(" BATCH VIDEO CROPPER SELESAI")
|
|
print(f" Total batch terekam: {batch_counter}")
|
|
print(f" Total frame diproses: {frame_idx}")
|
|
print(f" Folder output: {ARCHIVE_BASE}")
|
|
print("=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_batch_video_cropper()
|