1534 lines
59 KiB
Python
1534 lines
59 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from dataclasses import dataclass, replace
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
DEFAULT_MODEL = ROOT / "DATA" / "models" / "karung-dimuat-seg-200e.pt"
|
|
DEFAULT_VIDEO = ROOT / "DATA" / "Camera2_segments" / "clip_005.mp4"
|
|
DEFAULT_CONFIG = ROOT / "configs" / "area_truk.json"
|
|
SACK_CLASS_ID = 1
|
|
DEFAULT_ORIENTATION = "horizontal"
|
|
DEFAULT_DIRECTION = "bottom_to_top"
|
|
DEFAULT_LINE_FRACTION = 0.55
|
|
DEFAULT_MIN_APPROACH_DEPTH = 15.0
|
|
COUNTING_LOGIC_VERSION = "geometric_v15"
|
|
DEFAULT_APPROACH_MARGIN = 0
|
|
DEFAULT_COUNT_COOLDOWN_DIST = 85.0
|
|
DEFAULT_COUNT_COOLDOWN_FRAMES = 40
|
|
DEFAULT_SAME_SACK_RADIUS = 70.0
|
|
DEFAULT_STACK_SACK_RADIUS = 58.0
|
|
DEFAULT_OUTSIDE_CONFIRM_FRAMES = 2
|
|
DEFAULT_CROSSING_POINT_RATIO = 0.72
|
|
DEFAULT_STAGING_COOLDOWN_FRAMES = 15
|
|
DEFAULT_MIN_STAGING_DEPTH = 40.0
|
|
DEFAULT_MIN_TRACK_FRAMES = 8
|
|
DEFAULT_GHOST_TRACK_FRAMES = 0
|
|
DEFAULT_MIN_POST_CROSS_INSIDE_DEPTH = 0.0
|
|
DEFAULT_BURST_COOLDOWN_FRAMES = 0
|
|
DEFAULT_BURST_COOLDOWN_DIST = 40.0
|
|
DEFAULT_CLIP_WARMUP_FRAMES = 25
|
|
DEFAULT_INFERENCE_CONF = 0.2
|
|
DEFAULT_SEGMENT_FRAMES = 250
|
|
DEFAULT_MASK_ALPHA = 0.35
|
|
FLASH_DURATION_SEC = 0.35
|
|
DEFAULT_ZONE_X_MIN_FRAC = 0.52
|
|
DEFAULT_ZONE_X_MAX_FRAC = 0.68
|
|
DEFAULT_ZONE_Y_MIN_FRAC = 0.16
|
|
DEFAULT_ZONE_Y_MAX_FRAC = 0.58
|
|
COUNTING_PARAMS_FILE = ROOT / "configs" / "counting_params.json"
|
|
|
|
|
|
def load_counting_params() -> dict:
|
|
if not COUNTING_PARAMS_FILE.exists():
|
|
return {}
|
|
data = json.loads(COUNTING_PARAMS_FILE.read_text(encoding="utf-8"))
|
|
keys = (
|
|
"min_approach_depth",
|
|
"same_sack_radius",
|
|
"stack_sack_radius",
|
|
"outside_confirm_frames",
|
|
"count_cooldown_dist",
|
|
"count_cooldown_frames",
|
|
"crossing_point_ratio",
|
|
"staging_cooldown_frames",
|
|
"min_staging_depth",
|
|
"min_track_frames",
|
|
"ghost_track_frames",
|
|
"clip_warmup_frames",
|
|
"min_post_cross_inside_depth",
|
|
"burst_cooldown_frames",
|
|
"burst_cooldown_dist",
|
|
"conf",
|
|
)
|
|
return {k: data[k] for k in keys if k in data}
|
|
|
|
|
|
def _apply_counting_params_file() -> None:
|
|
global DEFAULT_MIN_APPROACH_DEPTH, DEFAULT_SAME_SACK_RADIUS, DEFAULT_STACK_SACK_RADIUS
|
|
global DEFAULT_OUTSIDE_CONFIRM_FRAMES, DEFAULT_COUNT_COOLDOWN_DIST, DEFAULT_COUNT_COOLDOWN_FRAMES
|
|
global DEFAULT_CROSSING_POINT_RATIO, DEFAULT_STAGING_COOLDOWN_FRAMES
|
|
global DEFAULT_MIN_STAGING_DEPTH, DEFAULT_MIN_TRACK_FRAMES, DEFAULT_GHOST_TRACK_FRAMES
|
|
global DEFAULT_MIN_POST_CROSS_INSIDE_DEPTH, DEFAULT_BURST_COOLDOWN_FRAMES
|
|
global DEFAULT_BURST_COOLDOWN_DIST, DEFAULT_CLIP_WARMUP_FRAMES, DEFAULT_INFERENCE_CONF
|
|
params = load_counting_params()
|
|
if "min_approach_depth" in params:
|
|
DEFAULT_MIN_APPROACH_DEPTH = float(params["min_approach_depth"])
|
|
if "same_sack_radius" in params:
|
|
DEFAULT_SAME_SACK_RADIUS = float(params["same_sack_radius"])
|
|
if "stack_sack_radius" in params:
|
|
DEFAULT_STACK_SACK_RADIUS = float(params["stack_sack_radius"])
|
|
if "outside_confirm_frames" in params:
|
|
DEFAULT_OUTSIDE_CONFIRM_FRAMES = int(params["outside_confirm_frames"])
|
|
if "count_cooldown_dist" in params:
|
|
DEFAULT_COUNT_COOLDOWN_DIST = float(params["count_cooldown_dist"])
|
|
if "count_cooldown_frames" in params:
|
|
DEFAULT_COUNT_COOLDOWN_FRAMES = int(params["count_cooldown_frames"])
|
|
if "crossing_point_ratio" in params:
|
|
DEFAULT_CROSSING_POINT_RATIO = float(params["crossing_point_ratio"])
|
|
if "staging_cooldown_frames" in params:
|
|
DEFAULT_STAGING_COOLDOWN_FRAMES = int(params["staging_cooldown_frames"])
|
|
if "min_staging_depth" in params:
|
|
DEFAULT_MIN_STAGING_DEPTH = float(params["min_staging_depth"])
|
|
if "min_track_frames" in params:
|
|
DEFAULT_MIN_TRACK_FRAMES = int(params["min_track_frames"])
|
|
if "ghost_track_frames" in params:
|
|
DEFAULT_GHOST_TRACK_FRAMES = int(params["ghost_track_frames"])
|
|
if "min_post_cross_inside_depth" in params:
|
|
DEFAULT_MIN_POST_CROSS_INSIDE_DEPTH = float(params["min_post_cross_inside_depth"])
|
|
if "burst_cooldown_frames" in params:
|
|
DEFAULT_BURST_COOLDOWN_FRAMES = int(params["burst_cooldown_frames"])
|
|
if "burst_cooldown_dist" in params:
|
|
DEFAULT_BURST_COOLDOWN_DIST = float(params["burst_cooldown_dist"])
|
|
if "clip_warmup_frames" in params:
|
|
DEFAULT_CLIP_WARMUP_FRAMES = int(params["clip_warmup_frames"])
|
|
if "conf" in params:
|
|
DEFAULT_INFERENCE_CONF = float(params["conf"])
|
|
|
|
|
|
def reload_counting_params(path: Path | None = None) -> None:
|
|
"""Muat ulang parameter counting dari file JSON (mis. override per klip panjang)."""
|
|
global COUNTING_PARAMS_FILE
|
|
if path is not None:
|
|
COUNTING_PARAMS_FILE = path
|
|
_apply_counting_params_file()
|
|
|
|
|
|
_apply_counting_params_file()
|
|
|
|
|
|
def counting_params_snapshot() -> dict:
|
|
"""Parameter counting aktif (setelah load configs/counting_params.json)."""
|
|
return {
|
|
"counting_logic": COUNTING_LOGIC_VERSION,
|
|
"min_approach_depth": DEFAULT_MIN_APPROACH_DEPTH,
|
|
"same_sack_radius": DEFAULT_SAME_SACK_RADIUS,
|
|
"stack_sack_radius": DEFAULT_STACK_SACK_RADIUS,
|
|
"outside_confirm_frames": DEFAULT_OUTSIDE_CONFIRM_FRAMES,
|
|
"count_cooldown_dist": DEFAULT_COUNT_COOLDOWN_DIST,
|
|
"count_cooldown_frames": DEFAULT_COUNT_COOLDOWN_FRAMES,
|
|
"crossing_point_ratio": DEFAULT_CROSSING_POINT_RATIO,
|
|
"staging_cooldown_frames": DEFAULT_STAGING_COOLDOWN_FRAMES,
|
|
"min_staging_depth": DEFAULT_MIN_STAGING_DEPTH,
|
|
"min_track_frames": DEFAULT_MIN_TRACK_FRAMES,
|
|
"ghost_track_frames": DEFAULT_GHOST_TRACK_FRAMES,
|
|
"clip_warmup_frames": DEFAULT_CLIP_WARMUP_FRAMES,
|
|
}
|
|
|
|
|
|
def counting_params_match(stored: dict | None) -> bool:
|
|
if not stored:
|
|
return False
|
|
current = counting_params_snapshot()
|
|
if stored.get("counting_logic") != current["counting_logic"]:
|
|
return False
|
|
for key in (
|
|
"min_approach_depth",
|
|
"same_sack_radius",
|
|
"stack_sack_radius",
|
|
"outside_confirm_frames",
|
|
"count_cooldown_dist",
|
|
"count_cooldown_frames",
|
|
"crossing_point_ratio",
|
|
):
|
|
if key not in stored:
|
|
return False
|
|
if key == "outside_confirm_frames":
|
|
if int(stored[key]) != int(current[key]):
|
|
return False
|
|
elif float(stored[key]) != float(current[key]):
|
|
return False
|
|
return True
|
|
|
|
|
|
def counting_inference_kwargs(**overrides) -> dict:
|
|
"""Kwargs untuk count_sacks_in_video dari snapshot + override CLI."""
|
|
snap = counting_params_snapshot()
|
|
return {
|
|
"min_approach_depth": overrides.get("min_approach_depth", snap["min_approach_depth"]),
|
|
"same_sack_radius": overrides.get("same_sack_radius", snap["same_sack_radius"]),
|
|
"stack_sack_radius": overrides.get("stack_sack_radius", snap["stack_sack_radius"]),
|
|
"outside_confirm_frames": overrides.get(
|
|
"outside_confirm_frames", snap["outside_confirm_frames"]
|
|
),
|
|
"count_cooldown_dist": overrides.get("count_cooldown_dist", snap["count_cooldown_dist"]),
|
|
"count_cooldown_frames": overrides.get(
|
|
"count_cooldown_frames", snap["count_cooldown_frames"]
|
|
),
|
|
"crossing_point_ratio": overrides.get("crossing_point_ratio", snap["crossing_point_ratio"]),
|
|
"staging_cooldown_frames": overrides.get(
|
|
"staging_cooldown_frames", snap["staging_cooldown_frames"]
|
|
),
|
|
"min_staging_depth": overrides.get("min_staging_depth", snap["min_staging_depth"]),
|
|
"min_track_frames": overrides.get("min_track_frames", snap["min_track_frames"]),
|
|
"ghost_track_frames": overrides.get("ghost_track_frames", snap["ghost_track_frames"]),
|
|
"clip_warmup_frames": overrides.get("clip_warmup_frames", snap["clip_warmup_frames"]),
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BoundarySettings:
|
|
"""Garis dan zona truk tetap — ubah hanya lewat calibrate.py / configs/*.json."""
|
|
|
|
line_pos: int
|
|
orientation: str
|
|
direction: str
|
|
zone_x_min: int
|
|
zone_x_max: int
|
|
zone_y_min: int
|
|
zone_y_max: int
|
|
|
|
|
|
@dataclass
|
|
class CountFlash:
|
|
number: int
|
|
x: int
|
|
y: int
|
|
frames_left: int
|
|
|
|
|
|
@dataclass
|
|
class CountedSackLabel:
|
|
"""Karung terhitung — tetap ditampilkan sepanjang video demo."""
|
|
|
|
number: int
|
|
box: np.ndarray
|
|
|
|
|
|
def _box_center(box: np.ndarray) -> tuple[float, float]:
|
|
x1, y1, x2, y2 = box
|
|
return (float(x1 + x2) / 2, float(y1 + y2) / 2)
|
|
|
|
|
|
def update_session_label_boxes(
|
|
session_labels: list[CountedSackLabel],
|
|
boxes,
|
|
counter: LineCounter,
|
|
match_dist: float = 130.0,
|
|
) -> None:
|
|
if not session_labels or boxes is None or boxes.id is None:
|
|
return
|
|
|
|
truck_boxes: list[np.ndarray] = []
|
|
for box in boxes.xyxy.cpu().numpy():
|
|
px, py = tracking_point(box)
|
|
if counter._in_truck_zone(px, py):
|
|
truck_boxes.append(box)
|
|
|
|
used: set[int] = set()
|
|
for label in session_labels:
|
|
lx, ly = _box_center(label.box)
|
|
best_idx = None
|
|
best_dist = match_dist
|
|
for i, box in enumerate(truck_boxes):
|
|
if i in used:
|
|
continue
|
|
bx, by = _box_center(box)
|
|
dist = ((lx - bx) ** 2 + (ly - by) ** 2) ** 0.5
|
|
if dist < best_dist:
|
|
best_dist = dist
|
|
best_idx = i
|
|
if best_idx is not None:
|
|
label.box = truck_boxes[best_idx].copy()
|
|
used.add(best_idx)
|
|
|
|
|
|
def draw_persisted_sacks(
|
|
frame: np.ndarray,
|
|
session_labels: list[CountedSackLabel],
|
|
show_mask: bool = False,
|
|
masks=None,
|
|
mask_alpha: float = DEFAULT_MASK_ALPHA,
|
|
) -> None:
|
|
for label in session_labels:
|
|
if show_mask and masks is not None and getattr(masks, "xy", None) is not None:
|
|
lx, ly = _box_center(label.box)
|
|
best_poly = None
|
|
best_dist = 150.0
|
|
for poly in masks.xy:
|
|
if poly is None or len(poly) < 3:
|
|
continue
|
|
px = float(np.mean(poly[:, 0]))
|
|
py = float(np.mean(poly[:, 1]))
|
|
dist = ((lx - px) ** 2 + (ly - py) ** 2) ** 0.5
|
|
if dist < best_dist:
|
|
best_dist = dist
|
|
best_poly = poly
|
|
if best_poly is not None:
|
|
draw_sack_mask(frame, best_poly, alpha=mask_alpha)
|
|
draw_sack_box(frame, label.box)
|
|
draw_sack_number_badge(frame, label.box, label.number)
|
|
|
|
|
|
class LineCounter:
|
|
"""Hitung karung yang melewati garis batas masuk ke dalam kotak truk (aturan geometri)."""
|
|
|
|
def __init__(
|
|
self,
|
|
line_pos: int,
|
|
orientation: str = "horizontal",
|
|
direction: str = "bottom_to_top",
|
|
min_approach_depth: float = DEFAULT_MIN_APPROACH_DEPTH,
|
|
approach_margin: int = DEFAULT_APPROACH_MARGIN,
|
|
count_cooldown_dist: float = DEFAULT_COUNT_COOLDOWN_DIST,
|
|
count_cooldown_frames: int = DEFAULT_COUNT_COOLDOWN_FRAMES,
|
|
same_sack_radius: float = DEFAULT_SAME_SACK_RADIUS,
|
|
stack_sack_radius: float = DEFAULT_STACK_SACK_RADIUS,
|
|
outside_confirm_frames: int = DEFAULT_OUTSIDE_CONFIRM_FRAMES,
|
|
staging_cooldown_frames: int = DEFAULT_STAGING_COOLDOWN_FRAMES,
|
|
min_staging_depth: float = DEFAULT_MIN_STAGING_DEPTH,
|
|
min_track_frames: int = DEFAULT_MIN_TRACK_FRAMES,
|
|
ghost_track_frames: int = DEFAULT_GHOST_TRACK_FRAMES,
|
|
clip_warmup_frames: int = DEFAULT_CLIP_WARMUP_FRAMES,
|
|
min_post_cross_inside_depth: float = DEFAULT_MIN_POST_CROSS_INSIDE_DEPTH,
|
|
burst_cooldown_frames: int = DEFAULT_BURST_COOLDOWN_FRAMES,
|
|
burst_cooldown_dist: float = DEFAULT_BURST_COOLDOWN_DIST,
|
|
zone_x_min: int | None = None,
|
|
zone_x_max: int | None = None,
|
|
zone_y_min: int | None = None,
|
|
zone_y_max: int | None = None,
|
|
enable_diag: bool = False,
|
|
) -> None:
|
|
self._line_pos = line_pos
|
|
self.orientation = orientation
|
|
self.direction = direction
|
|
self.min_approach_depth = min_approach_depth
|
|
self.approach_margin = approach_margin
|
|
self.count_cooldown_dist = count_cooldown_dist
|
|
self.count_cooldown_frames = count_cooldown_frames
|
|
self.same_sack_radius = same_sack_radius
|
|
self.stack_sack_radius = stack_sack_radius
|
|
self.outside_confirm_frames = outside_confirm_frames
|
|
self.staging_cooldown_frames = staging_cooldown_frames
|
|
self.min_staging_depth = min_staging_depth
|
|
self.min_track_frames = min_track_frames
|
|
self.ghost_track_frames = ghost_track_frames
|
|
self.clip_warmup_frames = clip_warmup_frames
|
|
self.min_post_cross_inside_depth = min_post_cross_inside_depth
|
|
self.burst_cooldown_frames = burst_cooldown_frames
|
|
self.burst_cooldown_dist = burst_cooldown_dist
|
|
self.zone_x_min = zone_x_min
|
|
self.zone_x_max = zone_x_max
|
|
self.zone_y_min = zone_y_min
|
|
self.zone_y_max = zone_y_max
|
|
self.prev_cross: dict[int, tuple[float, float]] = {}
|
|
self.prev_foot: dict[int, tuple[float, float]] = {}
|
|
self.seen_approaching: set[int] = set()
|
|
self.seen_outside_zone: set[int] = set()
|
|
self.outside_streak: dict[int, int] = {}
|
|
self.born_in_truck: set[int] = set()
|
|
self.max_below_line: dict[int, float] = {}
|
|
self.recent_counts: list[tuple[float, float, int, int]] = []
|
|
self.burst_counts: list[tuple[float, float, int]] = []
|
|
self.counted_positions: list[tuple[float, float, int]] = []
|
|
self.counted_ids: set[int] = set()
|
|
self.last_count_frame: dict[int, int] = {}
|
|
self.track_count_hits: dict[int, int] = {}
|
|
self.count_numbers: dict[int, int] = {}
|
|
self.count = 0
|
|
self.track_frames: dict[int, int] = {}
|
|
self._diag: dict[int, dict] | None = {} if enable_diag else None
|
|
|
|
def reset_state(self) -> None:
|
|
"""Reset state antar segmen (mis. klip 10 detik digabung)."""
|
|
self.prev_cross.clear()
|
|
self.prev_foot.clear()
|
|
self.seen_approaching.clear()
|
|
self.seen_outside_zone.clear()
|
|
self.outside_streak.clear()
|
|
self.born_in_truck.clear()
|
|
self.max_below_line.clear()
|
|
self.recent_counts.clear()
|
|
self.burst_counts.clear()
|
|
self.counted_positions.clear()
|
|
self.counted_ids.clear()
|
|
self.last_count_frame.clear()
|
|
self.track_count_hits.clear()
|
|
self.count_numbers.clear()
|
|
self.track_frames.clear()
|
|
self.count = 0
|
|
|
|
def _inside_depth(self, axis: float) -> float:
|
|
"""Kedalaman masuk ke dalam truk melewati garis (px)."""
|
|
if self.orientation == "horizontal":
|
|
if self.direction == "bottom_to_top":
|
|
return max(0.0, self.line_pos - axis)
|
|
return max(0.0, axis - self.line_pos)
|
|
if self.direction == "left_to_right":
|
|
return max(0.0, self.line_pos - axis)
|
|
return max(0.0, axis - self.line_pos)
|
|
|
|
def _burst_blocked(self, x: float, y: float, frame_idx: int) -> bool:
|
|
"""Blokir multi-track menghitung satu kejadian (burst temporal + spasial)."""
|
|
self.burst_counts = [
|
|
(bx, by, bf)
|
|
for bx, by, bf in self.burst_counts
|
|
if frame_idx - bf < self.burst_cooldown_frames
|
|
]
|
|
temporal_guard = max(8, self.burst_cooldown_frames // 2)
|
|
for bx, by, bf in self.burst_counts:
|
|
dt = frame_idx - bf
|
|
if dt < temporal_guard:
|
|
return True
|
|
if ((x - bx) ** 2 + (y - by) ** 2) ** 0.5 < self.burst_cooldown_dist:
|
|
return True
|
|
return False
|
|
|
|
def _inside_truck_body(self, x: float, y: float) -> bool:
|
|
"""Area truk di atas garis — zona buta (karung sudah masuk, tidak dilabeli)."""
|
|
axis = y if self.orientation == "horizontal" else x
|
|
return self._in_truck_zone(x, y) and not self._below_line(axis)
|
|
|
|
def _is_approach_visible(self, x: float, y: float) -> bool:
|
|
"""Karung belum masuk truk — boleh ditampilkan di koridor/staging."""
|
|
return self._in_entry_corridor(x, y) or self._outside_truck_zone(x, y)
|
|
|
|
def _diag_entry(self, track_id: int) -> dict | None:
|
|
if self._diag is None:
|
|
return None
|
|
if track_id not in self._diag:
|
|
self._diag[track_id] = {
|
|
"frames": 0,
|
|
"seen_approaching": False,
|
|
"seen_outside_zone": False,
|
|
"born_in_truck": False,
|
|
"max_below_line": 0.0,
|
|
"first_point": None,
|
|
"last_point": None,
|
|
"cross_attempts": [],
|
|
"counted": False,
|
|
}
|
|
return self._diag[track_id]
|
|
|
|
def _diag_cross(self, track_id: int, frame_idx: int, reason: str, counted: bool = False) -> None:
|
|
entry = self._diag_entry(track_id)
|
|
if entry is None:
|
|
return
|
|
entry["cross_attempts"].append({"frame": frame_idx, "reason": reason, "counted": counted})
|
|
|
|
@property
|
|
def line_pos(self) -> int:
|
|
return self._line_pos
|
|
|
|
def _in_x_corridor(self, x: float) -> bool:
|
|
if self.zone_x_min is not None and x < self.zone_x_min:
|
|
return False
|
|
if self.zone_x_max is not None and x > self.zone_x_max:
|
|
return False
|
|
return True
|
|
|
|
def _in_truck_zone(self, x: float, y: float) -> bool:
|
|
if not self._in_x_corridor(x):
|
|
return False
|
|
if self.zone_y_min is not None and y < self.zone_y_min:
|
|
return False
|
|
if self.zone_y_max is not None and y > self.zone_y_max:
|
|
return False
|
|
return True
|
|
|
|
def _outside_truck_zone(self, x: float, y: float) -> bool:
|
|
if self.zone_x_min is not None and x < self.zone_x_min:
|
|
return True
|
|
if self.zone_x_max is not None and x > self.zone_x_max:
|
|
return True
|
|
if self.zone_y_min is not None and y < self.zone_y_min:
|
|
return True
|
|
if self.zone_y_max is not None and y > self.zone_y_max:
|
|
return True
|
|
return False
|
|
|
|
def _below_line(self, axis: float) -> bool:
|
|
if self.orientation == "vertical":
|
|
if self.direction == "left_to_right":
|
|
return axis < self.line_pos
|
|
return axis > self.line_pos
|
|
if self.direction == "top_to_bottom":
|
|
return axis < self.line_pos
|
|
return axis > self.line_pos
|
|
|
|
def _in_entry_corridor(self, x: float, y: float) -> bool:
|
|
"""Koridor masuk: lebar bak truk, di bawah garis hitung."""
|
|
if self.orientation == "horizontal":
|
|
return self._in_x_corridor(x) and self._below_line(y)
|
|
if self.direction == "left_to_right":
|
|
return self._in_truck_zone(x, y) and x < self.line_pos
|
|
return self._in_truck_zone(x, y) and x > self.line_pos
|
|
|
|
def _crossed(self, prev_axis: float, curr_axis: float) -> bool:
|
|
margin = self.approach_margin
|
|
if self.orientation == "vertical":
|
|
if self.direction == "left_to_right":
|
|
return prev_axis + margin < self.line_pos <= curr_axis
|
|
return prev_axis - margin > self.line_pos >= curr_axis
|
|
if self.direction == "top_to_bottom":
|
|
return prev_axis + margin < self.line_pos <= curr_axis
|
|
return prev_axis - margin > self.line_pos >= curr_axis
|
|
|
|
def _approach_depth(self, track_id: int) -> float:
|
|
deepest = self.max_below_line.get(track_id)
|
|
if deepest is None:
|
|
return 0.0
|
|
if self.orientation == "horizontal":
|
|
if self.direction == "bottom_to_top":
|
|
return deepest - self.line_pos
|
|
return self.line_pos - deepest
|
|
if self.direction == "left_to_right":
|
|
return self.line_pos - deepest
|
|
return deepest - self.line_pos
|
|
|
|
def _record_geometry(self, track_id: int, x: float, y: float, entry: dict | None) -> None:
|
|
if self._outside_truck_zone(x, y):
|
|
streak = self.outside_streak.get(track_id, 0) + 1
|
|
self.outside_streak[track_id] = streak
|
|
if streak >= self.outside_confirm_frames:
|
|
self.seen_outside_zone.add(track_id)
|
|
self.seen_approaching.add(track_id)
|
|
if entry is not None:
|
|
entry["seen_outside_zone"] = True
|
|
entry["seen_approaching"] = True
|
|
else:
|
|
self.outside_streak[track_id] = 0
|
|
|
|
axis = y if self.orientation == "horizontal" else x
|
|
if self._in_entry_corridor(x, y):
|
|
self.seen_approaching.add(track_id)
|
|
if entry is not None:
|
|
entry["seen_approaching"] = True
|
|
|
|
if self._below_line(axis) and (self._in_x_corridor(x) or self._outside_truck_zone(x, y)):
|
|
prev_deepest = self.max_below_line.get(track_id, axis)
|
|
if self.direction == "bottom_to_top" or self.direction == "left_to_right":
|
|
deepest = max(prev_deepest, axis)
|
|
else:
|
|
deepest = min(prev_deepest, axis)
|
|
self.max_below_line[track_id] = deepest
|
|
if entry is not None:
|
|
entry["max_below_line"] = round(self._approach_depth(track_id), 1)
|
|
|
|
def _staging_depth_now(self, foot_axis: float) -> float:
|
|
if not self._below_line(foot_axis):
|
|
return 0.0
|
|
if self.orientation == "horizontal":
|
|
if self.direction == "bottom_to_top":
|
|
return foot_axis - self.line_pos
|
|
return self.line_pos - foot_axis
|
|
if self.direction == "left_to_right":
|
|
return self.line_pos - foot_axis
|
|
return foot_axis - self.line_pos
|
|
|
|
def _can_recount_staged(self, track_id: int, foot_axis: float, frame_idx: int) -> bool:
|
|
last = self.last_count_frame.get(track_id)
|
|
if last is None:
|
|
return False
|
|
if self.track_count_hits.get(track_id, 0) >= 2:
|
|
return False
|
|
if frame_idx - last <= self.staging_cooldown_frames:
|
|
return False
|
|
staging = self._staging_depth_now(foot_axis)
|
|
return staging >= self.min_staging_depth * 1.15
|
|
|
|
def _too_close_to_recent(self, x: float, y: float, frame_idx: int, track_id: int | None = None) -> bool:
|
|
if track_id is not None and track_id in self.seen_outside_zone:
|
|
frame_window = self.staging_cooldown_frames
|
|
dist_limit = self.stack_sack_radius
|
|
else:
|
|
frame_window = self.count_cooldown_frames
|
|
dist_limit = self.count_cooldown_dist
|
|
|
|
self.recent_counts = [
|
|
(rx, ry, rf, rtid)
|
|
for rx, ry, rf, rtid in self.recent_counts
|
|
if frame_idx - rf < frame_window
|
|
]
|
|
for rx, ry, _, rtid in self.recent_counts:
|
|
if track_id is not None and rtid == track_id:
|
|
continue
|
|
if ((x - rx) ** 2 + (y - ry) ** 2) ** 0.5 < dist_limit:
|
|
return True
|
|
return False
|
|
|
|
def _dedup_radius(self, track_id: int) -> float:
|
|
if track_id in self.seen_outside_zone:
|
|
return self.stack_sack_radius
|
|
return self.same_sack_radius
|
|
|
|
def _near_counted_position(self, x: float, y: float, track_id: int) -> bool:
|
|
radius = self._dedup_radius(track_id)
|
|
for px, py, _ in self.counted_positions:
|
|
if ((x - px) ** 2 + (y - py) ** 2) ** 0.5 < radius:
|
|
return True
|
|
return False
|
|
|
|
def _reject(self, track_id: int, frame_idx: int, reason: str) -> None:
|
|
self._diag_cross(track_id, frame_idx, reason)
|
|
|
|
def update(
|
|
self,
|
|
track_id: int,
|
|
cross: tuple[float, float],
|
|
frame_idx: int = 0,
|
|
foot: tuple[float, float] | None = None,
|
|
) -> int | None:
|
|
"""Update track: cross=titik crossing garis, foot=bottom-center untuk zona/dedup."""
|
|
if foot is None:
|
|
foot = cross
|
|
cx, cy = cross
|
|
fx, fy = foot
|
|
foot_axis = fy if self.orientation == "horizontal" else fx
|
|
|
|
if track_id in self.born_in_truck:
|
|
return None
|
|
|
|
if track_id in self.counted_ids:
|
|
if self._can_recount_staged(track_id, foot_axis, frame_idx):
|
|
self.counted_ids.discard(track_id)
|
|
else:
|
|
return None
|
|
|
|
self.track_frames[track_id] = self.track_frames.get(track_id, 0) + 1
|
|
track_frames = self.track_frames[track_id]
|
|
|
|
entry = self._diag_entry(track_id)
|
|
if entry is not None:
|
|
entry["frames"] = track_frames
|
|
entry["last_point"] = (round(fx, 1), round(fy, 1))
|
|
if entry["first_point"] is None:
|
|
entry["first_point"] = entry["last_point"]
|
|
|
|
prev_cross = self.prev_cross.get(track_id)
|
|
prev_foot = self.prev_foot.get(track_id)
|
|
self.prev_cross[track_id] = cross
|
|
self.prev_foot[track_id] = foot
|
|
|
|
if track_id not in self.seen_outside_zone and self._near_counted_position(fx, fy, track_id):
|
|
if not self._below_line(foot_axis):
|
|
if frame_idx >= self.clip_warmup_frames:
|
|
self.born_in_truck.add(track_id)
|
|
if entry is not None:
|
|
entry["born_in_truck"] = True
|
|
return None
|
|
|
|
if prev_cross is None:
|
|
self._record_geometry(track_id, fx, fy, entry)
|
|
if self._near_counted_position(fx, fy, track_id) and not self._below_line(foot_axis):
|
|
if frame_idx >= self.clip_warmup_frames:
|
|
self.born_in_truck.add(track_id)
|
|
if entry is not None:
|
|
entry["born_in_truck"] = True
|
|
elif self._in_truck_zone(fx, fy) and not self._below_line(foot_axis):
|
|
if frame_idx >= self.clip_warmup_frames:
|
|
self.born_in_truck.add(track_id)
|
|
if entry is not None:
|
|
entry["born_in_truck"] = True
|
|
return None
|
|
|
|
self._record_geometry(track_id, fx, fy, entry)
|
|
|
|
prev_cx, prev_cy = prev_cross
|
|
crossed_foot = False
|
|
if self.orientation == "horizontal":
|
|
prev_foot_axis = prev_foot[1] if prev_foot is not None else prev_cy
|
|
crossed_point = self._crossed(prev_cy, cy)
|
|
crossed_foot = self._crossed(prev_foot_axis, foot_axis)
|
|
foot_only = crossed_foot and not crossed_point and not self._below_line(cy)
|
|
crossed = crossed_point or foot_only
|
|
else:
|
|
prev_foot_axis = prev_foot[0] if prev_foot is not None else prev_cx
|
|
crossed_point = self._crossed(prev_cx, cx)
|
|
crossed_foot = self._crossed(prev_foot_axis, foot_axis)
|
|
foot_only = crossed_foot and not crossed_point and not self._below_line(cx)
|
|
crossed = crossed_point or foot_only
|
|
|
|
in_zone = self._in_truck_zone(cx, cy) or (
|
|
crossed_foot and not crossed_point and self._in_truck_zone(fx, fy)
|
|
)
|
|
in_x = self._in_x_corridor(cx) or (
|
|
crossed_foot and not crossed_point and self._in_x_corridor(fx)
|
|
)
|
|
if not crossed or not in_zone or not in_x:
|
|
return None
|
|
|
|
if track_id not in self.seen_approaching:
|
|
self._reject(track_id, frame_idx, "belum_di_koridor_masuk")
|
|
return None
|
|
if not self._below_line(prev_foot_axis):
|
|
self._reject(track_id, frame_idx, "prev_tidak_di_bawah_garis")
|
|
return None
|
|
if prev_foot is not None:
|
|
prev_fx, prev_fy = prev_foot
|
|
if not (self._in_entry_corridor(prev_fx, prev_fy) or self._outside_truck_zone(prev_fx, prev_fy)):
|
|
self._reject(track_id, frame_idx, "prev_tidak_di_koridor_masuk")
|
|
return None
|
|
depth = self._approach_depth(track_id)
|
|
shallow_foot_cross = (
|
|
foot_only
|
|
and self.min_approach_depth <= depth < min(self.min_staging_depth, self.min_approach_depth * 1.5)
|
|
)
|
|
if track_id in self.seen_outside_zone and not shallow_foot_cross:
|
|
if (
|
|
self.ghost_track_frames > 0
|
|
and track_frames <= self.ghost_track_frames
|
|
and depth < self.min_staging_depth
|
|
):
|
|
self._reject(track_id, frame_idx, f"ghost_staging<{self.min_staging_depth}")
|
|
return None
|
|
if depth < self.min_approach_depth:
|
|
self._reject(track_id, frame_idx, f"kedalaman_approach<{self.min_approach_depth}")
|
|
return None
|
|
elif depth < self.min_approach_depth:
|
|
self._reject(track_id, frame_idx, f"kedalaman_approach<{self.min_approach_depth}")
|
|
return None
|
|
inside_depth = self._inside_depth(foot_axis)
|
|
min_inside = self.min_post_cross_inside_depth
|
|
if min_inside > 0:
|
|
if foot_only:
|
|
min_inside = max(min_inside, 10.0)
|
|
elif shallow_foot_cross:
|
|
min_inside = max(min_inside, 8.0)
|
|
if inside_depth < min_inside:
|
|
self._reject(track_id, frame_idx, f"belum_masuk_dalam<{min_inside:.0f}")
|
|
return None
|
|
if (
|
|
self.burst_cooldown_frames > 0
|
|
and foot_only
|
|
and inside_depth < 12
|
|
and depth < self.min_staging_depth * 0.55
|
|
):
|
|
self._reject(track_id, frame_idx, "palet_sentuh_garis")
|
|
return None
|
|
if self.burst_cooldown_frames > 0 and self._burst_blocked(cx, cy, frame_idx):
|
|
self._reject(track_id, frame_idx, "burst_spasial")
|
|
return None
|
|
deep_carry = depth >= self.min_approach_depth * 2.5
|
|
if self.min_track_frames > 0 and track_frames < self.min_track_frames and not deep_carry:
|
|
self._reject(track_id, frame_idx, f"track_frames<{self.min_track_frames}")
|
|
return None
|
|
if self._too_close_to_recent(cx, cy, frame_idx, track_id):
|
|
self._reject(track_id, frame_idx, "cooldown_spasial_di_garis")
|
|
return None
|
|
if self._near_counted_position(fx, fy, track_id) and not self._below_line(foot_axis):
|
|
self._reject(track_id, frame_idx, "dekat_posisi_sudah_dihitung")
|
|
return None
|
|
self.counted_ids.add(track_id)
|
|
self.last_count_frame[track_id] = frame_idx
|
|
self.track_count_hits[track_id] = self.track_count_hits.get(track_id, 0) + 1
|
|
self.recent_counts.append((cx, cy, frame_idx, track_id))
|
|
self.burst_counts.append((cx, cy, frame_idx))
|
|
self.counted_positions.append((fx, fy, frame_idx))
|
|
self.count += 1
|
|
self.count_numbers[track_id] = self.count
|
|
if entry is not None:
|
|
entry["counted"] = True
|
|
self._diag_cross(track_id, frame_idx, "OK", counted=True)
|
|
return self.count
|
|
|
|
|
|
def tracking_point_foot(box: np.ndarray) -> tuple[float, float]:
|
|
"""Bottom-center — posisi fisik karung untuk zona & dedup."""
|
|
x1, y1, x2, y2 = box
|
|
return (float(x1 + x2) / 2, float(y2))
|
|
|
|
|
|
def tracking_point(box: np.ndarray, ratio: float | None = None) -> tuple[float, float]:
|
|
"""Titik crossing garis — antara tengah dan bawah bbox."""
|
|
x1, y1, x2, y2 = box
|
|
r = DEFAULT_CROSSING_POINT_RATIO if ratio is None else ratio
|
|
r = max(0.35, min(0.95, r))
|
|
return (float(x1 + x2) / 2, float(y1 + r * (y2 - y1)))
|
|
|
|
|
|
def track_points(box: np.ndarray) -> tuple[tuple[float, float], tuple[float, float]]:
|
|
return tracking_point(box), tracking_point_foot(box)
|
|
|
|
|
|
def draw_truck_counter_box(
|
|
frame: np.ndarray,
|
|
boundary: BoundarySettings,
|
|
count: int,
|
|
) -> None:
|
|
"""Panel hitung di dalam kotak truk."""
|
|
x1, x2 = boundary.zone_x_min, boundary.zone_x_max
|
|
y1 = boundary.zone_y_min
|
|
panel_w, panel_h = 140, 90
|
|
px1 = x1 + 12
|
|
py1 = y1 + 12
|
|
px2 = px1 + panel_w
|
|
py2 = py1 + panel_h
|
|
overlay = frame.copy()
|
|
cv2.rectangle(overlay, (px1, py1), (px2, py2), (0, 0, 0), -1)
|
|
cv2.addWeighted(overlay, 0.65, frame, 0.35, 0, frame)
|
|
cv2.rectangle(frame, (px1, py1), (px2, py2), (0, 255, 0), 2)
|
|
cv2.putText(frame, "KARUNG", (px1 + 14, py1 + 28), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (200, 200, 200), 2)
|
|
cv2.putText(
|
|
frame,
|
|
str(count),
|
|
(px1 + 28, py1 + 78),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
2.0,
|
|
(0, 255, 0),
|
|
4,
|
|
)
|
|
|
|
|
|
def draw_blind_truck_overlay(
|
|
frame: np.ndarray,
|
|
boundary: BoundarySettings,
|
|
line_pos: int,
|
|
) -> None:
|
|
"""Overlay tipis di area truk atas garis — zona buta."""
|
|
x1, x2 = boundary.zone_x_min, boundary.zone_x_max
|
|
y1, y2 = boundary.zone_y_min, boundary.zone_y_max
|
|
if boundary.orientation == "horizontal":
|
|
blind_y2 = min(line_pos, y2)
|
|
if blind_y2 <= y1:
|
|
return
|
|
overlay = frame.copy()
|
|
cv2.rectangle(overlay, (x1, y1), (x2, blind_y2), (40, 40, 40), -1)
|
|
cv2.addWeighted(overlay, 0.12, frame, 0.88, 0, frame)
|
|
|
|
|
|
def draw_boundary(
|
|
frame: np.ndarray,
|
|
line_pos: int,
|
|
orientation: str,
|
|
zone_x_min: int | None = None,
|
|
zone_x_max: int | None = None,
|
|
zone_y_min: int | None = None,
|
|
zone_y_max: int | None = None,
|
|
color: tuple[int, int, int] = (0, 255, 0),
|
|
) -> None:
|
|
height, width = frame.shape[:2]
|
|
if orientation == "vertical":
|
|
cv2.line(frame, (line_pos, 0), (line_pos, height), color, 3)
|
|
return
|
|
|
|
if zone_x_min is not None and zone_x_max is not None:
|
|
y1 = zone_y_min if zone_y_min is not None else 0
|
|
y2 = zone_y_max if zone_y_max is not None else height
|
|
cv2.rectangle(frame, (zone_x_min, y1), (zone_x_max, y2), (0, 180, 255), 3)
|
|
cv2.line(frame, (zone_x_min, line_pos), (zone_x_max, line_pos), color, 4)
|
|
cv2.putText(
|
|
frame,
|
|
"AREA TRUK",
|
|
(zone_x_min + 8, max(y1 - 10, 24)),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.8,
|
|
(0, 180, 255),
|
|
2,
|
|
)
|
|
cv2.putText(
|
|
frame,
|
|
"GARIS BATAS MASUK",
|
|
(zone_x_min + 8, line_pos - 12),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.7,
|
|
color,
|
|
2,
|
|
)
|
|
cv2.putText(
|
|
frame,
|
|
"koridor masuk",
|
|
(zone_x_min + 8, min(line_pos + 28, y2 - 4)),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.55,
|
|
(0, 220, 255),
|
|
1,
|
|
)
|
|
else:
|
|
cv2.line(frame, (0, line_pos), (width, line_pos), color, 3)
|
|
|
|
|
|
def draw_count_hud(frame: np.ndarray, count: int) -> None:
|
|
w = frame.shape[1]
|
|
scale = w / 1280.0
|
|
panel_w = min(int(460 * scale), w - 20)
|
|
panel_h = int(100 * scale)
|
|
overlay = frame.copy()
|
|
cv2.rectangle(overlay, (10, 10), (panel_w, panel_h), (0, 0, 0), -1)
|
|
cv2.addWeighted(overlay, 0.55, frame, 0.45, 0, frame)
|
|
cv2.putText(
|
|
frame,
|
|
f"Muat ke truk: {count}",
|
|
(int(24 * scale), int(78 * scale)),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
1.9 * scale,
|
|
(0, 255, 0),
|
|
max(1, int(4 * scale)),
|
|
)
|
|
|
|
|
|
def draw_sack_number_badge(frame: np.ndarray, box: np.ndarray, number: int) -> None:
|
|
x1, y1, x2, y2 = box
|
|
cx, cy = int((x1 + x2) / 2), int((y1 + y2) / 2)
|
|
label = str(number)
|
|
scale = 1.4
|
|
thickness = 3
|
|
(tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, scale, thickness)
|
|
pad = 8
|
|
cv2.rectangle(
|
|
frame,
|
|
(cx - tw // 2 - pad, cy - th // 2 - pad),
|
|
(cx + tw // 2 + pad, cy + th // 2 + pad),
|
|
(0, 180, 0),
|
|
-1,
|
|
)
|
|
cv2.putText(
|
|
frame,
|
|
label,
|
|
(cx - tw // 2, cy + th // 2),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
scale,
|
|
(255, 255, 255),
|
|
thickness,
|
|
)
|
|
|
|
|
|
def tick_flashes(flashes: list[CountFlash]) -> None:
|
|
for flash in flashes:
|
|
flash.frames_left -= 1
|
|
flashes[:] = [flash for flash in flashes if flash.frames_left > 0]
|
|
|
|
|
|
def draw_count_flashes(frame: np.ndarray, flashes: list[CountFlash], fps: float) -> None:
|
|
hold_frames = max(int(fps * 0.35), 1)
|
|
for flash in flashes:
|
|
pulse = min(1.0, flash.frames_left / hold_frames)
|
|
scale = 1.6 + pulse * 1.4
|
|
thickness = 4 + int(pulse * 2)
|
|
cv2.circle(frame, (flash.x, flash.y), 22, (0, 255, 0), 3)
|
|
cv2.putText(
|
|
frame,
|
|
str(flash.number),
|
|
(flash.x - 18, flash.y - 28),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
scale,
|
|
(0, 255, 0),
|
|
thickness,
|
|
)
|
|
|
|
|
|
def draw_sack_mask(
|
|
frame: np.ndarray,
|
|
polygon: np.ndarray,
|
|
color: tuple[int, int, int] = (0, 255, 0),
|
|
alpha: float = DEFAULT_MASK_ALPHA,
|
|
) -> None:
|
|
if polygon is None or len(polygon) < 3:
|
|
return
|
|
pts = np.asarray(polygon, dtype=np.int32).reshape(-1, 1, 2)
|
|
overlay = frame.copy()
|
|
cv2.fillPoly(overlay, [pts], color)
|
|
cv2.addWeighted(overlay, alpha, frame, 1.0 - alpha, 0, frame)
|
|
|
|
|
|
def draw_sack_box(frame: np.ndarray, box: np.ndarray, color: tuple[int, int, int] = (0, 255, 0)) -> None:
|
|
x1, y1, x2, y2 = map(int, box)
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
|
|
|
|
|
|
def annotate_tracks(
|
|
frame: np.ndarray,
|
|
boxes,
|
|
counter: LineCounter,
|
|
flashes: list[CountFlash],
|
|
fps: float,
|
|
frame_idx: int,
|
|
session_offset: int = 0,
|
|
blind_truck: bool = True,
|
|
session_labels: list[CountedSackLabel] | None = None,
|
|
masks=None,
|
|
show_mask: bool = False,
|
|
mask_alpha: float = DEFAULT_MASK_ALPHA,
|
|
on_count_cb=None,
|
|
) -> None:
|
|
"""Update counter. blind: koridor saja. labeled truck: koridor + session_labels di truk."""
|
|
if boxes is None or boxes.id is None:
|
|
return
|
|
|
|
mask_polygons: list | None = None
|
|
if show_mask and masks is not None and getattr(masks, "xy", None) is not None:
|
|
mask_polygons = masks.xy
|
|
|
|
for idx, (box, track_id) in enumerate(
|
|
zip(boxes.xyxy.cpu().numpy(), boxes.id.int().cpu().tolist())
|
|
):
|
|
cross, foot = track_points(box)
|
|
new_count = counter.update(track_id, cross, frame_idx, foot)
|
|
|
|
if new_count is not None:
|
|
global_num = session_offset + new_count
|
|
flashes.append(
|
|
CountFlash(
|
|
number=global_num,
|
|
x=int(cross[0]),
|
|
y=int(cross[1]),
|
|
frames_left=int(fps * FLASH_DURATION_SEC),
|
|
)
|
|
)
|
|
if session_labels is not None:
|
|
session_labels.append(CountedSackLabel(global_num, box.copy()))
|
|
if on_count_cb is not None:
|
|
on_count_cb(global_num, box.copy())
|
|
|
|
if counter._inside_truck_body(foot[0], foot[1]):
|
|
if not blind_truck:
|
|
continue
|
|
continue
|
|
|
|
if not counter._is_approach_visible(foot[0], foot[1]):
|
|
continue
|
|
|
|
if show_mask and mask_polygons is not None and idx < len(mask_polygons):
|
|
draw_sack_mask(frame, mask_polygons[idx], color=(0, 200, 255), alpha=mask_alpha)
|
|
draw_sack_box(frame, box, color=(0, 200, 255))
|
|
|
|
|
|
def resolve_boundary(args: argparse.Namespace, width: int, height: int) -> BoundarySettings:
|
|
if args.config and args.config.exists():
|
|
try:
|
|
with open(args.config, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
video_width = data.get("video_width")
|
|
video_height = data.get("video_height")
|
|
if video_width and video_height and (video_width != width or video_height != height):
|
|
scale_x = width / video_width
|
|
scale_y = height / video_height
|
|
print(f"[Scale Coordinates] Scaling args from {video_width}x{video_height} to {width}x{height} (x_scale={scale_x:.3f}, y_scale={scale_y:.3f})")
|
|
if args.line_x is not None:
|
|
args.line_x = int(args.line_x * scale_x)
|
|
if args.line_y is not None:
|
|
args.line_y = int(args.line_y * scale_y)
|
|
if args.zone_x_min is not None:
|
|
args.zone_x_min = int(args.zone_x_min * scale_x)
|
|
if args.zone_x_max is not None:
|
|
args.zone_x_max = int(args.zone_x_max * scale_x)
|
|
if args.zone_y_min is not None:
|
|
args.zone_y_min = int(args.zone_y_min * scale_y)
|
|
if args.zone_y_max is not None:
|
|
args.zone_y_max = int(args.zone_y_max * scale_y)
|
|
except Exception as e:
|
|
print(f"[Warning] resolve_boundary scaling error: {e}")
|
|
|
|
line_pos = resolve_line_pos(args, width, height)
|
|
zone_x_min, zone_x_max, zone_y_min, zone_y_max = resolve_truck_zone(args, width, height)
|
|
return BoundarySettings(
|
|
line_pos=line_pos,
|
|
orientation=args.orientation,
|
|
direction=args.direction,
|
|
zone_x_min=zone_x_min,
|
|
zone_x_max=zone_x_max,
|
|
zone_y_min=zone_y_min,
|
|
zone_y_max=zone_y_max,
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Hitung karung yang masuk truk (melewati garis di area truk saja).",
|
|
)
|
|
parser.add_argument("--video", type=Path, default=DEFAULT_VIDEO, help="Input video path.")
|
|
parser.add_argument("--model", type=Path, default=DEFAULT_MODEL, help="YOLO segmentation model.")
|
|
parser.add_argument("--line-x", type=int, default=None, help="Garis vertikal (pixel).")
|
|
parser.add_argument("--line-y", type=int, default=None, help="Garis horizontal (pixel).")
|
|
parser.add_argument("--zone-x-min", type=int, default=None, help="Batas kiri area truk (pixel).")
|
|
parser.add_argument("--zone-x-max", type=int, default=None, help="Batas kanan bak truk (pixel).")
|
|
parser.add_argument("--zone-y-min", type=int, default=None, help="Batas atas bak truk (pixel).")
|
|
parser.add_argument("--zone-y-max", type=int, default=None, help="Batas bawah bak truk (pixel).")
|
|
parser.add_argument("--orientation", choices=("vertical", "horizontal"), default=DEFAULT_ORIENTATION)
|
|
parser.add_argument(
|
|
"--direction",
|
|
choices=("left_to_right", "right_to_left", "top_to_bottom", "bottom_to_top"),
|
|
default=DEFAULT_DIRECTION,
|
|
)
|
|
parser.add_argument("--min-approach-depth", type=float, default=DEFAULT_MIN_APPROACH_DEPTH,
|
|
help="Min kedalaman di bawah garis (px) sebelum dihitung masuk truk.")
|
|
parser.add_argument("--conf", type=float, default=0.25)
|
|
parser.add_argument("--device", default="", help="Perangkat inferensi: cuda, cpu, atau kosong=auto.")
|
|
parser.add_argument("--show-mask", action="store_true", help="Overlay mask segmentasi karung terhitung.")
|
|
parser.add_argument("--mask-alpha", type=float, default=DEFAULT_MASK_ALPHA, help="Opasitas mask (0-1).")
|
|
parser.add_argument("--output", type=Path, default=None)
|
|
parser.add_argument(
|
|
"--segment-frames",
|
|
type=int,
|
|
default=0,
|
|
help="Reset counter tiap N frame (untuk video gabungan beberapa klip). 0=nonaktif.",
|
|
)
|
|
parser.add_argument(
|
|
"--config",
|
|
type=Path,
|
|
default=DEFAULT_CONFIG if DEFAULT_CONFIG.exists() else None,
|
|
help="File JSON dari calibrate.py (garis/zona tetap).",
|
|
)
|
|
parser.add_argument(
|
|
"--dynamic-truk",
|
|
action="store_true",
|
|
help="Gunakan model detektor truk untuk kalibrasi otomatis dinamis.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def apply_config(args: argparse.Namespace) -> None:
|
|
if args.config is None or not args.config.exists():
|
|
return
|
|
data = json.loads(args.config.read_text(encoding="utf-8"))
|
|
for attr in ("zone_x_min", "zone_x_max", "zone_y_min", "zone_y_max", "line_x", "line_y"):
|
|
if getattr(args, attr) is None and data.get(attr) is not None:
|
|
setattr(args, attr, data[attr])
|
|
if data.get("orientation"):
|
|
args.orientation = data["orientation"]
|
|
if data.get("direction"):
|
|
args.direction = data["direction"]
|
|
|
|
|
|
def resolve_line_pos(args: argparse.Namespace, width: int, height: int) -> int:
|
|
if args.orientation == "vertical":
|
|
if args.line_x is not None:
|
|
return args.line_x
|
|
return int(width * DEFAULT_LINE_FRACTION)
|
|
if args.line_y is not None:
|
|
return args.line_y
|
|
return int(height * DEFAULT_LINE_FRACTION)
|
|
|
|
|
|
def resolve_truck_zone(args: argparse.Namespace, width: int, height: int) -> tuple[int, int, int, int]:
|
|
x_min = args.zone_x_min if args.zone_x_min is not None else int(width * DEFAULT_ZONE_X_MIN_FRAC)
|
|
x_max = args.zone_x_max if args.zone_x_max is not None else int(width * DEFAULT_ZONE_X_MAX_FRAC)
|
|
y_min = args.zone_y_min if args.zone_y_min is not None else int(height * DEFAULT_ZONE_Y_MIN_FRAC)
|
|
y_max = args.zone_y_max if args.zone_y_max is not None else int(height * DEFAULT_ZONE_Y_MAX_FRAC)
|
|
return x_min, x_max, y_min, y_max
|
|
|
|
|
|
def calibrate_truck_dynamically(
|
|
video_path: Path,
|
|
model_path: Path,
|
|
device: str | None = None,
|
|
sample_frames: int = 15,
|
|
truck_class_id: int = 2,
|
|
conf_thresh: float = 0.4,
|
|
) -> tuple[int, int, int, int] | None:
|
|
"""Scan first N frames of the video to detect the parked truck bed and return [x1, y1, x2, y2]."""
|
|
dedicated_model = ROOT / "DATA" / "models" / "truck-detector.pt"
|
|
if dedicated_model.exists():
|
|
print(f"[Dynamic Truk] Menggunakan model detektor truk dedikasi: {dedicated_model.name}")
|
|
model = YOLO(str(dedicated_model))
|
|
target_cls = 0
|
|
else:
|
|
model = YOLO(str(model_path))
|
|
target_cls = truck_class_id
|
|
if not model.names or target_cls not in model.names:
|
|
print(f"[Dynamic Truk] Kelas ID {target_cls} tidak ditemukan pada model {model_path.name} dan tidak ada model dedikasi. Kalibrasi dinamis dilewati.")
|
|
return None
|
|
|
|
cap = cv2.VideoCapture(str(video_path))
|
|
if not cap.isOpened():
|
|
return None
|
|
|
|
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
|
|
frames_to_check = min(total, sample_frames)
|
|
truck_boxes = []
|
|
|
|
print(f"[Dynamic Truk] Memulai kalibrasi otomatis...")
|
|
for f in range(frames_to_check):
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, f)
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
results = model.predict(frame, classes=[target_cls], conf=conf_thresh, verbose=False, device=device)
|
|
for box in results[0].boxes:
|
|
cls = int(box.cls[0].item())
|
|
if cls == target_cls:
|
|
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().tolist()
|
|
truck_boxes.append((x1, y1, x2, y2))
|
|
|
|
cap.release()
|
|
|
|
if not truck_boxes:
|
|
print("[Dynamic Truk] Truk tidak terdeteksi pada frame awal. Menggunakan kalibrasi statis.")
|
|
return None
|
|
|
|
# We take the median of the coordinates to ignore outliers / noise
|
|
xs1 = sorted([b[0] for b in truck_boxes])
|
|
ys1 = sorted([b[1] for b in truck_boxes])
|
|
xs2 = sorted([b[2] for b in truck_boxes])
|
|
ys2 = sorted([b[3] for b in truck_boxes])
|
|
|
|
mid_idx = len(truck_boxes) // 2
|
|
tx1 = int(xs1[mid_idx])
|
|
ty1 = int(ys1[mid_idx])
|
|
tx2 = int(xs2[mid_idx])
|
|
ty2 = int(ys2[mid_idx])
|
|
|
|
if tx2 - tx1 < 100 or ty2 - ty1 < 100:
|
|
print(f"[Dynamic Truk] Dimensi truk terlalu kecil ({tx2-tx1}x{ty2-ty1}). Menggunakan kalibrasi statis.")
|
|
return None
|
|
|
|
print(f"[Dynamic Truk] Sukses mendeteksi bak truk: x={tx1}..{tx2}, y={ty1}..{ty2}")
|
|
return tx1, ty1, tx2, ty2
|
|
|
|
|
|
def calculate_truck_coverage(boxes, counter, grid_size: int = 15) -> float:
|
|
"""Hitung persentase luas area bak truk yang tertutup oleh deteksi karung (grid-based)."""
|
|
if boxes is None or len(boxes) == 0:
|
|
return 0.0
|
|
|
|
x_min, x_max = counter.zone_x_min, counter.zone_x_max
|
|
y_min, y_max = counter.zone_y_min, counter.zone_y_max
|
|
|
|
if x_min is None or x_max is None or y_min is None or y_max is None:
|
|
return 0.0
|
|
|
|
# Area bak truk di atas garis batas masuk (untuk horizontal bottom_to_top)
|
|
if counter.orientation == "horizontal":
|
|
if counter.direction == "bottom_to_top":
|
|
truck_x_min, truck_x_max = x_min, x_max
|
|
truck_y_min, truck_y_max = y_min, counter.line_pos
|
|
else:
|
|
truck_x_min, truck_x_max = x_min, x_max
|
|
truck_y_min, truck_y_max = counter.line_pos, y_max
|
|
else:
|
|
# vertical
|
|
truck_y_min, truck_y_max = y_min, y_max
|
|
if counter.direction == "left_to_right":
|
|
truck_x_min, truck_x_max = counter.line_pos, x_max
|
|
else:
|
|
truck_x_min, truck_x_max = x_min, counter.line_pos
|
|
|
|
w_truck = truck_x_max - truck_x_min
|
|
h_truck = truck_y_max - truck_y_min
|
|
if w_truck <= 0 or h_truck <= 0:
|
|
return 0.0
|
|
|
|
# Build active boxes inside the truck bed
|
|
active_boxes = []
|
|
for box in boxes.xyxy.cpu().numpy():
|
|
bx1, by1, bx2, by2 = box
|
|
# Clip box to truck bed boundary
|
|
ix1 = max(bx1, truck_x_min)
|
|
iy1 = max(by1, truck_y_min)
|
|
ix2 = min(bx2, truck_x_max)
|
|
iy2 = min(by2, truck_y_max)
|
|
if ix2 > ix1 and iy2 > iy1:
|
|
active_boxes.append((ix1, iy1, ix2, iy2))
|
|
|
|
if not active_boxes:
|
|
return 0.0
|
|
|
|
# Check cells
|
|
covered_cells = 0
|
|
total_cells = grid_size * grid_size
|
|
|
|
for i in range(grid_size):
|
|
cx = truck_x_min + (i + 0.5) * (w_truck / grid_size)
|
|
for j in range(grid_size):
|
|
cy = truck_y_min + (j + 0.5) * (h_truck / grid_size)
|
|
|
|
# Check if this center is inside any active box
|
|
for bx1, by1, bx2, by2 in active_boxes:
|
|
if bx1 <= cx <= bx2 and by1 <= cy <= by2:
|
|
covered_cells += 1
|
|
break
|
|
|
|
return covered_cells / total_cells
|
|
|
|
|
|
def classify_filling_status(sack_count: int, coverage: float) -> str:
|
|
"""Klasifikasikan status pengisian truk (kosong, sebagian, penuh) berdasarkan jumlah dan sebaran karung."""
|
|
if sack_count == 0:
|
|
return "kosong"
|
|
|
|
# Jika sebaran/luas cakupan deteksi di bak >= 35% atau total karung >= 100
|
|
if coverage >= 0.35 or sack_count >= 100:
|
|
return "penuh"
|
|
|
|
return "sebagian"
|
|
|
|
|
|
|
|
def count_sacks_in_video(
|
|
video_path: Path,
|
|
model_path: Path,
|
|
boundary: BoundarySettings,
|
|
min_approach_depth: float | None = None,
|
|
conf: float | None = None,
|
|
output_path: Path | None = None,
|
|
save_video: bool = True,
|
|
model: YOLO | None = None,
|
|
device: str | None = None,
|
|
enable_diag: bool = False,
|
|
show_mask: bool = False,
|
|
mask_alpha: float = DEFAULT_MASK_ALPHA,
|
|
segment_reset_frames: int | None = None,
|
|
same_sack_radius: float | None = None,
|
|
stack_sack_radius: float | None = None,
|
|
outside_confirm_frames: int | None = None,
|
|
count_cooldown_dist: float | None = None,
|
|
count_cooldown_frames: int | None = None,
|
|
staging_cooldown_frames: int | None = None,
|
|
min_staging_depth: float | None = None,
|
|
min_track_frames: int | None = None,
|
|
ghost_track_frames: int | None = None,
|
|
clip_warmup_frames: int | None = None,
|
|
crossing_point_ratio: float | None = None,
|
|
blind_truck: bool = True,
|
|
dynamic_truk: bool = False,
|
|
) -> tuple[int, Path | None, LineCounter]:
|
|
tuned = counting_inference_kwargs()
|
|
if conf is None:
|
|
conf = 0.25
|
|
min_approach_depth = tuned["min_approach_depth"] if min_approach_depth is None else min_approach_depth
|
|
same_sack_radius = tuned["same_sack_radius"] if same_sack_radius is None else same_sack_radius
|
|
stack_sack_radius = tuned["stack_sack_radius"] if stack_sack_radius is None else stack_sack_radius
|
|
outside_confirm_frames = (
|
|
tuned["outside_confirm_frames"] if outside_confirm_frames is None else outside_confirm_frames
|
|
)
|
|
count_cooldown_dist = (
|
|
tuned["count_cooldown_dist"] if count_cooldown_dist is None else count_cooldown_dist
|
|
)
|
|
count_cooldown_frames = (
|
|
tuned["count_cooldown_frames"] if count_cooldown_frames is None else count_cooldown_frames
|
|
)
|
|
staging_cooldown_frames = (
|
|
tuned["staging_cooldown_frames"] if staging_cooldown_frames is None else staging_cooldown_frames
|
|
)
|
|
min_staging_depth = tuned["min_staging_depth"] if min_staging_depth is None else min_staging_depth
|
|
min_track_frames = tuned["min_track_frames"] if min_track_frames is None else min_track_frames
|
|
ghost_track_frames = tuned["ghost_track_frames"] if ghost_track_frames is None else ghost_track_frames
|
|
clip_warmup_frames = tuned["clip_warmup_frames"] if clip_warmup_frames is None else clip_warmup_frames
|
|
if crossing_point_ratio is None:
|
|
crossing_point_ratio = tuned["crossing_point_ratio"]
|
|
if crossing_point_ratio is not None:
|
|
global DEFAULT_CROSSING_POINT_RATIO
|
|
DEFAULT_CROSSING_POINT_RATIO = crossing_point_ratio
|
|
cap = cv2.VideoCapture(str(video_path))
|
|
if not cap.isOpened():
|
|
raise RuntimeError(f"Unable to open video: {video_path}")
|
|
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
|
|
|
|
if save_video:
|
|
if output_path is None:
|
|
output_dir = ROOT / "runs" / "count"
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
output_path = output_dir / f"{video_path.stem}_counted.mp4"
|
|
writer = cv2.VideoWriter(
|
|
str(output_path),
|
|
cv2.VideoWriter_fourcc(*"mp4v"),
|
|
fps,
|
|
(width, height),
|
|
)
|
|
else:
|
|
output_path = None
|
|
writer = None
|
|
|
|
if model is None:
|
|
model = YOLO(str(model_path))
|
|
|
|
if dynamic_truk:
|
|
truck_box = calibrate_truck_dynamically(
|
|
video_path=video_path,
|
|
model_path=model_path,
|
|
device=device,
|
|
truck_class_id=2,
|
|
conf_thresh=0.4,
|
|
)
|
|
if truck_box is not None:
|
|
tx1, ty1, tx2, ty2 = truck_box
|
|
boundary = replace(
|
|
boundary,
|
|
zone_x_min=tx1,
|
|
zone_x_max=tx2,
|
|
zone_y_min=ty1,
|
|
zone_y_max=ty2,
|
|
line_pos=int(ty2 - 5),
|
|
)
|
|
print(f"[Dynamic Truk] Override garis batas: line_y={boundary.line_pos}, zone_x={boundary.zone_x_min}..{boundary.zone_x_max}")
|
|
|
|
counter = LineCounter(
|
|
line_pos=boundary.line_pos,
|
|
orientation=boundary.orientation,
|
|
direction=boundary.direction,
|
|
min_approach_depth=min_approach_depth,
|
|
same_sack_radius=same_sack_radius,
|
|
stack_sack_radius=stack_sack_radius,
|
|
outside_confirm_frames=outside_confirm_frames,
|
|
count_cooldown_dist=count_cooldown_dist,
|
|
count_cooldown_frames=count_cooldown_frames,
|
|
staging_cooldown_frames=staging_cooldown_frames,
|
|
min_staging_depth=min_staging_depth,
|
|
min_track_frames=min_track_frames,
|
|
ghost_track_frames=ghost_track_frames,
|
|
clip_warmup_frames=clip_warmup_frames,
|
|
min_post_cross_inside_depth=DEFAULT_MIN_POST_CROSS_INSIDE_DEPTH,
|
|
burst_cooldown_frames=DEFAULT_BURST_COOLDOWN_FRAMES,
|
|
burst_cooldown_dist=DEFAULT_BURST_COOLDOWN_DIST,
|
|
zone_x_min=boundary.zone_x_min,
|
|
zone_x_max=boundary.zone_x_max,
|
|
zone_y_min=boundary.zone_y_min,
|
|
zone_y_max=boundary.zone_y_max,
|
|
enable_diag=enable_diag,
|
|
)
|
|
flashes: list[CountFlash] = []
|
|
frame_idx = 0
|
|
segment_total = 0
|
|
reset_tracker = False
|
|
|
|
while cap.isOpened():
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
|
|
if segment_reset_frames and frame_idx > 0 and frame_idx % segment_reset_frames == 0:
|
|
segment_total += counter.count
|
|
counter.reset_state()
|
|
flashes.clear()
|
|
reset_tracker = True
|
|
|
|
try:
|
|
results = model.track(
|
|
frame,
|
|
persist=not reset_tracker,
|
|
classes=[SACK_CLASS_ID],
|
|
conf=conf,
|
|
tracker="bytetrack.yaml",
|
|
device=device or None,
|
|
verbose=False,
|
|
)
|
|
except (ValueError, Exception) as exc:
|
|
if "ingular" not in str(exc).lower():
|
|
raise
|
|
try:
|
|
results = model.track(
|
|
frame,
|
|
persist=False,
|
|
classes=[SACK_CLASS_ID],
|
|
conf=conf,
|
|
tracker="bytetrack.yaml",
|
|
device=device or None,
|
|
verbose=False,
|
|
)
|
|
reset_tracker = True
|
|
except (ValueError, Exception):
|
|
frame_idx += 1
|
|
reset_tracker = True
|
|
continue
|
|
reset_tracker = False
|
|
|
|
if save_video:
|
|
annotated = frame.copy()
|
|
display_count = segment_total + counter.count
|
|
if blind_truck:
|
|
draw_blind_truck_overlay(annotated, boundary, boundary.line_pos)
|
|
annotate_tracks(
|
|
annotated,
|
|
results[0].boxes,
|
|
counter,
|
|
flashes,
|
|
fps,
|
|
frame_idx,
|
|
session_offset=segment_total,
|
|
blind_truck=True,
|
|
session_labels=None,
|
|
masks=results[0].masks,
|
|
show_mask=show_mask,
|
|
mask_alpha=mask_alpha,
|
|
)
|
|
tick_flashes(flashes)
|
|
draw_count_flashes(annotated, flashes, fps)
|
|
draw_boundary(
|
|
annotated,
|
|
boundary.line_pos,
|
|
boundary.orientation,
|
|
boundary.zone_x_min,
|
|
boundary.zone_x_max,
|
|
boundary.zone_y_min,
|
|
boundary.zone_y_max,
|
|
)
|
|
draw_truck_counter_box(annotated, boundary, display_count)
|
|
draw_count_hud(annotated, display_count)
|
|
writer.write(annotated)
|
|
else:
|
|
boxes = results[0].boxes
|
|
if boxes is not None and boxes.id is not None:
|
|
for box, track_id in zip(boxes.xyxy.cpu().numpy(), boxes.id.int().cpu().tolist()):
|
|
cross, foot = track_points(box)
|
|
counter.update(track_id, cross, frame_idx, foot)
|
|
|
|
frame_idx += 1
|
|
|
|
cap.release()
|
|
if writer is not None:
|
|
writer.release()
|
|
total = segment_total + counter.count
|
|
counter.count = total
|
|
return total, output_path, counter
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
apply_config(args)
|
|
if not args.video.exists():
|
|
raise FileNotFoundError(f"Video not found: {args.video}")
|
|
if not args.model.exists():
|
|
raise FileNotFoundError(f"Model not found: {args.model}")
|
|
|
|
cap = cv2.VideoCapture(str(args.video))
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
cap.release()
|
|
|
|
boundary = resolve_boundary(args, width, height)
|
|
total, output_path, _counter = count_sacks_in_video(
|
|
video_path=args.video,
|
|
model_path=args.model,
|
|
boundary=boundary,
|
|
min_approach_depth=args.min_approach_depth,
|
|
conf=args.conf,
|
|
output_path=args.output,
|
|
device=args.device or None,
|
|
show_mask=args.show_mask,
|
|
mask_alpha=args.mask_alpha,
|
|
segment_reset_frames=args.segment_frames or None,
|
|
dynamic_truk=args.dynamic_truk,
|
|
)
|
|
print(f"Garis: {boundary.orientation} @ {_counter.line_pos} ({'DINAMIS' if args.dynamic_truk else 'TETAP'}), arah={boundary.direction}")
|
|
if args.config:
|
|
print(f"Config : {args.config}")
|
|
print(f"Area truk: x={_counter.zone_x_min}..{_counter.zone_x_max}, y={_counter.zone_y_min}..{_counter.zone_y_max}")
|
|
print(f"Karung masuk truk: {total}")
|
|
print(f"Saved to: {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|