| """Video scene-cut detection, clip extraction, and subject-centered crop/resize. |
| |
| Uses plain (non-relative) imports of sibling modules so this file works whether |
| invoked as ``python dataset_builder/build.py`` or ``python -m dataset_builder.build`` |
| (build.py inserts its own directory onto sys.path before importing). |
| """ |
| import math |
| import statistics |
|
|
| import cv2 |
| import imageio |
| import numpy as np |
| from PIL import Image |
|
|
| import crop_resize |
| from detect import detect_subject_bbox, subdivide_bbox |
|
|
|
|
| def _frame_hist(frame_bgr): |
| hsv = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2HSV) |
| hist = cv2.calcHist([hsv], [0, 1], None, [50, 60], [0, 180, 0, 256]) |
| cv2.normalize(hist, hist) |
| return hist |
|
|
|
|
| def find_scene_cuts(video_path, diff_thresh=0.45, sample_stride=1): |
| """Returns a sorted list of frame indices where a scene cut is detected, |
| via Bhattacharyya distance between consecutive HSV histograms.""" |
| cap = cv2.VideoCapture(video_path) |
| cuts = [] |
| prev_hist = None |
| idx = 0 |
| while True: |
| ok, frame = cap.read() |
| if not ok: |
| break |
| if idx % sample_stride == 0: |
| hist = _frame_hist(frame) |
| if prev_hist is not None: |
| score = cv2.compareHist(prev_hist, hist, cv2.HISTCMP_BHATTACHARYYA) |
| if score > diff_thresh: |
| cuts.append(idx) |
| prev_hist = hist |
| idx += 1 |
| cap.release() |
| return cuts |
|
|
|
|
| def build_clip_ranges(total_frames, fps, cuts, clip_min_sec=2.0, clip_max_sec=6.0): |
| """Turns scene-cut frame indices into (start, end) frame ranges honoring |
| clip_min/max seconds: too-short segments are merged forward, segments |
| longer than clip_max are chopped into roughly-equal sub-segments.""" |
| min_frames = max(1, int(clip_min_sec * fps)) |
| max_frames = max(min_frames, int(clip_max_sec * fps)) |
|
|
| boundaries = [0] + sorted(set(c for c in cuts if 0 < c < total_frames)) + [total_frames] |
| raw_ranges = [] |
| start = boundaries[0] |
| for b in boundaries[1:]: |
| if b - start >= min_frames: |
| raw_ranges.append((start, b)) |
| start = b |
| if start < total_frames: |
| if raw_ranges and (total_frames - start) < min_frames: |
| s, _ = raw_ranges[-1] |
| raw_ranges[-1] = (s, total_frames) |
| else: |
| raw_ranges.append((start, total_frames)) |
|
|
| ranges = [] |
| for s, e in raw_ranges: |
| length = e - s |
| if length <= max_frames: |
| ranges.append((s, e)) |
| continue |
| n_chunks = math.ceil(length / max_frames) |
| chunk_len = length // n_chunks |
| for i in range(n_chunks): |
| cs = s + i * chunk_len |
| ce = e if i == n_chunks - 1 else cs + chunk_len |
| ranges.append((cs, ce)) |
|
|
| return [(s, e) for s, e in ranges if e - s >= max(1, min_frames // 2)] |
|
|
|
|
| def _sample_indices(start, end, n=3): |
| if end - start <= n: |
| return list(range(start, end)) |
| step = (end - start) / (n + 1) |
| return [int(start + step * (i + 1)) for i in range(n)] |
|
|
|
|
| def _median_bbox(bboxes): |
| return ( |
| statistics.median(b[0] for b in bboxes), |
| statistics.median(b[1] for b in bboxes), |
| statistics.median(b[2] for b in bboxes), |
| statistics.median(b[3] for b in bboxes), |
| ) |
|
|
|
|
| def _motion_tag(frames_gray): |
| if len(frames_gray) < 2: |
| return "static, minimal motion" |
| mags = [] |
| for a, b in zip(frames_gray, frames_gray[1:]): |
| flow = cv2.calcOpticalFlowFarneback(a, b, None, 0.5, 3, 15, 3, 5, 1.2, 0) |
| mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1]) |
| mags.append(float(np.mean(mag))) |
| avg = sum(mags) / len(mags) |
| if avg < 0.5: |
| return "static, minimal motion" |
| if avg < 2.0: |
| return "subtle motion" |
| if avg < 5.0: |
| return "moderate motion" |
| return "fast, dynamic motion" |
|
|
|
|
| def extract_clip(video_path, start, end, target_fps, max_tier="1080p", pad_overrides=None, |
| body_region="auto", body_region_frac=0.45): |
| """Reads frames [start, end) from video_path, computes one stable subject |
| crop for the whole clip (median bbox over sampled frames), resamples to |
| target_fps, trims to a 4n+1 frame count, and returns |
| (frames_bgr, (target_w, target_h), representative_pil_image, motion_tag), |
| or None if no frames could be read.""" |
| cap = cv2.VideoCapture(video_path) |
| src_fps = cap.get(cv2.CAP_PROP_FPS) or target_fps |
| img_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| img_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
|
|
| bboxes, kinds = [], [] |
| for idx in _sample_indices(start, end, n=3): |
| cap.set(cv2.CAP_PROP_POS_FRAMES, idx) |
| ok, frame = cap.read() |
| if not ok: |
| continue |
| bbox, kind = detect_subject_bbox(frame) |
| bboxes.append(bbox) |
| kinds.append(kind) |
|
|
| if not bboxes: |
| cap.release() |
| return None |
|
|
| bbox = _median_bbox(bboxes) |
| kind = max(set(kinds), key=kinds.count) |
| if kind == "person" and body_region != "auto": |
| bbox = subdivide_bbox(bbox, body_region, body_region_frac) |
| kind = f"person_{body_region}" |
| crop_box, (tw, th), _bucket_name, _tier = crop_resize.compute_crop( |
| img_w, img_h, bbox, kind, pad_overrides, max_tier |
| ) |
|
|
| step = src_fps / target_fps |
| raw_indices = [] |
| f = float(start) |
| while f < end: |
| raw_indices.append(int(round(f))) |
| f += step |
| if len(raw_indices) >= 5: |
| n = ((len(raw_indices) - 1) // 4) * 4 + 1 |
| raw_indices = raw_indices[:n] |
|
|
| frames_out, frames_gray = [], [] |
| for idx in raw_indices: |
| cap.set(cv2.CAP_PROP_POS_FRAMES, idx) |
| ok, frame = cap.read() |
| if not ok: |
| continue |
| cropped = crop_resize.crop_resize_np(frame, crop_box, tw, th) |
| frames_out.append(cropped) |
| if len(frames_gray) < 6: |
| frames_gray.append(cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY)) |
| cap.release() |
|
|
| if not frames_out: |
| return None |
|
|
| motion_tag = _motion_tag(frames_gray) |
| mid = frames_out[len(frames_out) // 2] |
| representative = Image.fromarray(cv2.cvtColor(mid, cv2.COLOR_BGR2RGB)) |
|
|
| return frames_out, (tw, th), representative, motion_tag |
|
|
|
|
| def write_clip_mp4(frames_bgr, out_path, fps): |
| with imageio.get_writer(str(out_path), fps=fps, codec="libx264", quality=8) as writer: |
| for frame in frames_bgr: |
| writer.append_data(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) |
|
|
