import argparse import csv import json import math import shutil import subprocess from dataclasses import dataclass from pathlib import Path VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".mts", ".m2ts"} @dataclass class VideoMeta: abs_path: Path rel_path: Path group_name: str file_name: str size_bytes: int width: int height: int fps: float duration_sec: float nb_frames: int osd_present: bool def parse_args(): parser = argparse.ArgumentParser( description="Prepare grouped frame datasets from many videos for CVAT and YOLO." ) subparsers = parser.add_subparsers(dest="command", required=True) catalog = subparsers.add_parser("catalog", help="Scan videos and write catalog CSV files.") catalog.add_argument("--input-root", required=True, help="Root folder with grouped videos.") catalog.add_argument("--output-root", required=True, help="Workspace folder for CSV outputs.") preview = subparsers.add_parser( "preview", help="Extract low-rate preview frames for interval selection.", ) preview.add_argument("--input-root", required=True) preview.add_argument("--output-root", required=True) preview.add_argument("--every-sec", type=float, default=2.0) preview.add_argument("--image-ext", choices=("jpg", "png"), default="jpg") preview.add_argument("--jpg-quality", type=int, default=95) preview.add_argument("--overwrite", action="store_true") extract = subparsers.add_parser( "extract", help="Extract frames for selected intervals from intervals CSV.", ) extract.add_argument("--input-root", required=True) extract.add_argument("--intervals-csv", required=True) extract.add_argument("--output-root", required=True) extract.add_argument("--overwrite", action="store_true") extract_all = subparsers.add_parser( "extract-all", help="Extract frames from every video while preserving the input group structure.", ) extract_all.add_argument("--input-root", required=True) extract_all.add_argument("--output-root", required=True) extract_all.add_argument("--frame-step", type=int, default=2) extract_all.add_argument("--image-ext", choices=("jpg", "png"), default="png") extract_all.add_argument("--jpg-quality", type=int, default=95) extract_all.add_argument("--split", default="train") extract_all.add_argument("--role", default="bulk_raw") extract_all.add_argument( "--group-filter", action="append", default=[], help="Optional substring filter for group names. Can be repeated.", ) extract_all.add_argument("--overwrite", action="store_true") extract_flat_all = subparsers.add_parser( "extract-flat-all", help="Extract frames from every video into one flat output folder.", ) extract_flat_all.add_argument("--input-root", required=True) extract_flat_all.add_argument( "--output-root", required=True, help="Flat folder where all images will be written.", ) extract_flat_all.add_argument("--frame-step", type=int, default=1) extract_flat_all.add_argument("--image-ext", choices=("jpg", "png"), default="jpg") extract_flat_all.add_argument("--jpg-quality", type=int, default=95) extract_flat_all.add_argument( "--group-filter", action="append", default=[], help="Optional substring filter for group names. Can be repeated.", ) extract_flat_all.add_argument("--overwrite", action="store_true") return parser.parse_args() def main(): args = parse_args() input_root = Path(args.input_root).expanduser().resolve() output_root = Path(args.output_root).expanduser().resolve() if not input_root.exists(): raise SystemExit(f"Input root not found: {input_root}") ensure_tool("ffprobe") ensure_tool("ffmpeg") if args.command == "catalog": videos = scan_videos(input_root) write_catalog_outputs(videos, output_root) print(f"Wrote catalog for {len(videos)} videos into {output_root}") return if args.command == "preview": videos = scan_videos(input_root) preview_root = output_root / "previews" preview_root.mkdir(parents=True, exist_ok=True) manifest_path = output_root / "preview_manifest.csv" rows = [] for meta in videos: out_dir = preview_root / sanitize_path_component(meta.group_name) / sanitize_path_component(meta.rel_path.stem) if args.overwrite and out_dir.exists(): shutil.rmtree(out_dir) out_dir.mkdir(parents=True, exist_ok=True) pattern = out_dir / f"preview_%06d.{args.image_ext}" cmd = [ "ffmpeg", "-hide_banner", "-loglevel", "warning", "-y", "-i", str(meta.abs_path), "-vf", f"fps=1/{float(args.every_sec)}", ] if args.image_ext == "jpg": cmd.extend(["-q:v", str(jpeg_qscale(args.jpg_quality))]) cmd.append(str(pattern)) run_cmd(cmd) files = sorted(out_dir.glob(f"*.{args.image_ext}")) for idx, file_path in enumerate(files): rows.append( { "group_name": meta.group_name, "video_relpath": str(meta.rel_path).replace("\\", "/"), "preview_file": str(file_path.relative_to(output_root)).replace("\\", "/"), "preview_index": idx + 1, "approx_time_sec": round(idx * float(args.every_sec), 3), } ) write_csv( manifest_path, rows, ["group_name", "video_relpath", "preview_file", "preview_index", "approx_time_sec"], ) print(f"Wrote previews into {preview_root}") print(f"Wrote preview manifest: {manifest_path}") return if args.command == "extract": video_map = {str(v.rel_path).replace("\\", "/"): v for v in scan_videos(input_root)} intervals_csv = Path(args.intervals_csv).expanduser().resolve() if not intervals_csv.exists(): raise SystemExit(f"Intervals CSV not found: {intervals_csv}") rows = load_intervals(intervals_csv) manifest_rows = [] extracted_root = output_root / "frames" extracted_root.mkdir(parents=True, exist_ok=True) for row in rows: if not is_enabled(row.get("enabled", "")): continue rel_key = normalize_relpath(row["video_relpath"]) if rel_key not in video_map: raise SystemExit(f"Video not found in catalog: {rel_key}") meta = video_map[rel_key] split = row.get("split", "train").strip() or "train" role = row.get("role", "positive").strip() or "positive" frame_step = max(1, int(float(row.get("frame_step", "1") or "1"))) start_sec = clamp_float(row.get("start_sec", "0"), 0.0, meta.duration_sec) end_sec_raw = row.get("end_sec", "") end_sec = meta.duration_sec if end_sec_raw == "" else clamp_float(end_sec_raw, 0.0, meta.duration_sec) if end_sec <= start_sec: raise SystemExit(f"Invalid interval for {rel_key}: end_sec <= start_sec") image_ext = (row.get("image_ext", "png").strip().lower() or "png") if image_ext not in {"png", "jpg"}: raise SystemExit(f"Unsupported image_ext '{image_ext}' in intervals CSV") jpg_quality = max(70, min(100, int(float(row.get("jpg_quality", "95") or "95")))) extract_rows = extract_range_to_dir( meta=meta, output_root=output_root, extracted_root=extracted_root, split=split, role=role, start_sec=start_sec, end_sec=end_sec, frame_step=frame_step, image_ext=image_ext, jpg_quality=jpg_quality, overwrite=args.overwrite, interval_name=format_interval_name(start_sec, end_sec, frame_step), ) manifest_rows.extend(extract_rows) manifest_path = output_root / "extracted_frames_manifest.csv" write_csv( manifest_path, manifest_rows, ["split", "role", "group_name", "video_relpath", "output_file", "frame_idx", "time_sec", "frame_step", "image_ext"], ) print(f"Wrote extracted frames into {extracted_root}") print(f"Wrote extraction manifest: {manifest_path}") return if args.command == "extract-all": videos = scan_videos(input_root) if args.group_filter: filters = [v.strip().lower() for v in args.group_filter if v.strip()] videos = [meta for meta in videos if any(token in meta.group_name.lower() for token in filters)] frame_step = max(1, int(args.frame_step)) image_ext = args.image_ext jpg_quality = max(70, min(100, int(args.jpg_quality))) split = args.split.strip() or "train" role = args.role.strip() or "bulk_raw" extracted_root = output_root / "frames_all" extracted_root.mkdir(parents=True, exist_ok=True) manifest_rows = [] for meta in videos: extract_rows = extract_range_to_dir( meta=meta, output_root=output_root, extracted_root=extracted_root, split=split, role=role, start_sec=0.0, end_sec=meta.duration_sec, frame_step=frame_step, image_ext=image_ext, jpg_quality=jpg_quality, overwrite=args.overwrite, interval_name=f"all_step{frame_step}", ) manifest_rows.extend(extract_rows) manifest_path = output_root / "extract_all_manifest.csv" write_csv( manifest_path, manifest_rows, ["split", "role", "group_name", "video_relpath", "output_file", "frame_idx", "time_sec", "frame_step", "image_ext"], ) print(f"Wrote bulk extracted frames into {extracted_root}") print(f"Wrote bulk extraction manifest: {manifest_path}") return if args.command == "extract-flat-all": videos = scan_videos(input_root) if args.group_filter: filters = [v.strip().lower() for v in args.group_filter if v.strip()] videos = [meta for meta in videos if any(token in meta.group_name.lower() for token in filters)] frame_step = max(1, int(args.frame_step)) image_ext = args.image_ext jpg_quality = max(70, min(100, int(args.jpg_quality))) flat_root = output_root if args.overwrite and flat_root.exists(): shutil.rmtree(flat_root) flat_root.mkdir(parents=True, exist_ok=True) manifest_rows = [] for meta in videos: manifest_rows.extend( extract_range_flat( meta=meta, flat_root=flat_root, frame_step=frame_step, image_ext=image_ext, jpg_quality=jpg_quality, ) ) manifest_path = flat_root / "_flat_manifest.csv" write_csv( manifest_path, manifest_rows, ["group_name", "video_relpath", "output_file", "frame_idx", "time_sec", "frame_step", "image_ext"], ) print(f"Wrote flat extracted frames into {flat_root}") print(f"Wrote flat extraction manifest: {manifest_path}") return def ensure_tool(name: str): if shutil.which(name) is None: raise SystemExit(f"Required tool not found in PATH: {name}") def scan_videos(input_root: Path): videos = [] for path in sorted(input_root.rglob("*")): if not path.is_file() or path.suffix.lower() not in VIDEO_EXTS: continue rel_path = path.relative_to(input_root) group_name = rel_path.parts[0] if len(rel_path.parts) > 1 else "UNGROUPED" info = probe_video(path) videos.append( VideoMeta( abs_path=path, rel_path=rel_path, group_name=group_name, file_name=path.name, size_bytes=path.stat().st_size, width=info["width"], height=info["height"], fps=info["fps"], duration_sec=info["duration_sec"], nb_frames=info["nb_frames"], osd_present=("без osd" not in group_name.lower()), ) ) return videos def probe_video(path: Path): cmd = [ "ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width,height,r_frame_rate,avg_frame_rate,nb_frames,duration", "-of", "json", str(path), ] data = json.loads(run_cmd(cmd)) streams = data.get("streams", []) if not streams: raise SystemExit(f"No video stream found: {path}") stream = streams[0] fps_str = stream.get("avg_frame_rate") or stream.get("r_frame_rate") or "0/1" fps = parse_fraction(fps_str) duration = float(stream.get("duration") or 0.0) nb_frames_raw = stream.get("nb_frames") nb_frames = int(nb_frames_raw) if str(nb_frames_raw).isdigit() else int(round(duration * fps)) return { "width": int(stream.get("width") or 0), "height": int(stream.get("height") or 0), "fps": fps, "duration_sec": duration, "nb_frames": nb_frames, } def extract_range_to_dir( meta: VideoMeta, output_root: Path, extracted_root: Path, split: str, role: str, start_sec: float, end_sec: float, frame_step: int, image_ext: str, jpg_quality: int, overwrite: bool, interval_name: str, ): target_dir = ( extracted_root / sanitize_path_component(split) / sanitize_path_component(role) / sanitize_path_component(meta.group_name) / sanitize_path_component(meta.rel_path.stem) / sanitize_path_component(interval_name) ) if overwrite and target_dir.exists(): shutil.rmtree(target_dir) existing_files = sorted(target_dir.glob(f"*.{image_ext}")) if target_dir.exists() else [] if existing_files and (not overwrite): return build_manifest_rows( output_root=output_root, meta=meta, split=split, role=role, files=existing_files, frame_step=frame_step, ) target_dir.mkdir(parents=True, exist_ok=True) tmp_dir = target_dir / "_tmp" if tmp_dir.exists(): shutil.rmtree(tmp_dir) tmp_dir.mkdir(parents=True, exist_ok=True) tmp_pattern = tmp_dir / f"tmp_%06d.{image_ext}" cmd = [ "ffmpeg", "-hide_banner", "-loglevel", "warning", "-y", "-ss", f"{start_sec:.3f}", "-to", f"{end_sec:.3f}", "-i", str(meta.abs_path), "-vf", f"select=not(mod(n\\,{frame_step}))", "-vsync", "vfr", ] if image_ext == "jpg": cmd.extend(["-q:v", str(jpeg_qscale(jpg_quality))]) cmd.append(str(tmp_pattern)) run_cmd(cmd) files = sorted(tmp_dir.glob(f"*.{image_ext}")) start_frame_idx = max(0, int(round(start_sec * meta.fps))) final_files = [] for idx, tmp_file in enumerate(files): frame_idx = start_frame_idx + idx * frame_step time_sec = frame_idx / max(1e-6, meta.fps) final_name = ( f"{sanitize_path_component(meta.rel_path.stem)}__" f"f{frame_idx:06d}__t{time_sec:010.3f}.{image_ext}" ) final_path = target_dir / final_name tmp_file.replace(final_path) final_files.append(final_path) shutil.rmtree(tmp_dir) return build_manifest_rows( output_root=output_root, meta=meta, split=split, role=role, files=final_files, frame_step=frame_step, ) def build_manifest_rows(output_root: Path, meta: VideoMeta, split: str, role: str, files, frame_step: int): rows = [] for file_path in files: frame_idx, time_sec = parse_frame_info_from_name(file_path.stem, meta.fps) rows.append( { "split": split, "role": role, "group_name": meta.group_name, "video_relpath": normalize_relpath(meta.rel_path), "output_file": normalize_relpath(file_path.relative_to(output_root)), "frame_idx": frame_idx, "time_sec": round(time_sec, 3), "frame_step": frame_step, "image_ext": file_path.suffix.lstrip(".").lower(), } ) return rows def parse_frame_info_from_name(stem: str, fps: float): frame_idx = 0 time_sec = 0.0 for part in stem.split("__"): if part.startswith("f") and part[1:].isdigit(): frame_idx = int(part[1:]) elif part.startswith("t"): try: time_sec = float(part[1:]) except ValueError: time_sec = frame_idx / max(1e-6, fps) if time_sec <= 0.0 and frame_idx > 0: time_sec = frame_idx / max(1e-6, fps) return frame_idx, time_sec def extract_range_flat( meta: VideoMeta, flat_root: Path, frame_step: int, image_ext: str, jpg_quality: int, ): tmp_dir = flat_root / f"_tmp_{sanitize_path_component(meta.rel_path.stem)}" if tmp_dir.exists(): shutil.rmtree(tmp_dir) tmp_dir.mkdir(parents=True, exist_ok=True) tmp_pattern = tmp_dir / f"tmp_%06d.{image_ext}" cmd = [ "ffmpeg", "-hide_banner", "-loglevel", "warning", "-y", "-i", str(meta.abs_path), "-vf", f"select=not(mod(n\\,{frame_step}))", "-vsync", "vfr", ] if image_ext == "jpg": cmd.extend(["-q:v", str(jpeg_qscale(jpg_quality))]) cmd.append(str(tmp_pattern)) run_cmd(cmd) rows = [] files = sorted(tmp_dir.glob(f"*.{image_ext}")) safe_group = sanitize_path_component(meta.group_name) safe_video = sanitize_path_component(meta.rel_path.stem) for idx, tmp_file in enumerate(files): frame_idx = idx * frame_step time_sec = frame_idx / max(1e-6, meta.fps) final_name = f"{safe_group}__{safe_video}__f{frame_idx:06d}__t{time_sec:010.3f}.{image_ext}" final_path = flat_root / final_name tmp_file.replace(final_path) rows.append( { "group_name": meta.group_name, "video_relpath": normalize_relpath(meta.rel_path), "output_file": final_path.name, "frame_idx": frame_idx, "time_sec": round(time_sec, 3), "frame_step": frame_step, "image_ext": image_ext, } ) shutil.rmtree(tmp_dir) return rows def parse_fraction(value: str): num, den = value.split("/") den_v = float(den) if abs(den_v) < 1e-9: return 0.0 return float(num) / den_v def write_catalog_outputs(videos, output_root: Path): output_root.mkdir(parents=True, exist_ok=True) catalog_rows = [] summary_map = {} template_rows = [] for meta in videos: rel_path_str = normalize_relpath(meta.rel_path) catalog_rows.append( { "group_name": meta.group_name, "osd_present": int(meta.osd_present), "video_relpath": rel_path_str, "file_name": meta.file_name, "size_bytes": meta.size_bytes, "size_gb": round(meta.size_bytes / (1024 ** 3), 4), "duration_sec": round(meta.duration_sec, 3), "duration_min": round(meta.duration_sec / 60.0, 2), "fps": round(meta.fps, 6), "width": meta.width, "height": meta.height, "nb_frames": meta.nb_frames, } ) key = meta.group_name info = summary_map.setdefault(key, {"videos": 0, "bytes": 0, "duration": 0.0}) info["videos"] += 1 info["bytes"] += meta.size_bytes info["duration"] += meta.duration_sec template_rows.append( { "enabled": "no", "split": "train", "role": "positive", "video_relpath": rel_path_str, "group_name": meta.group_name, "start_sec": "", "end_sec": "", "frame_step": "2", "image_ext": "png", "jpg_quality": "95", "notes": "Duplicate row for more intervals or use role=hard_negative for sparse negatives.", } ) summary_rows = [] for group_name in sorted(summary_map): info = summary_map[group_name] summary_rows.append( { "group_name": group_name, "video_count": info["videos"], "total_size_gb": round(info["bytes"] / (1024 ** 3), 3), "total_duration_min": round(info["duration"] / 60.0, 2), "osd_present": int("без osd" not in group_name.lower()), } ) write_csv( output_root / "video_catalog.csv", catalog_rows, [ "group_name", "osd_present", "video_relpath", "file_name", "size_bytes", "size_gb", "duration_sec", "duration_min", "fps", "width", "height", "nb_frames", ], ) write_csv( output_root / "group_summary.csv", summary_rows, ["group_name", "video_count", "total_size_gb", "total_duration_min", "osd_present"], ) write_csv( output_root / "intervals_template.csv", template_rows, [ "enabled", "split", "role", "video_relpath", "group_name", "start_sec", "end_sec", "frame_step", "image_ext", "jpg_quality", "notes", ], ) def write_csv(path: Path, rows, fieldnames): path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", newline="", encoding="utf-8-sig") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() for row in rows: writer.writerow(row) def load_intervals(path: Path): with path.open("r", newline="", encoding="utf-8-sig") as f: return list(csv.DictReader(f)) def is_enabled(value: str): return str(value).strip().lower() in {"1", "true", "yes", "y", "on"} def clamp_float(value, min_value: float, max_value: float): value_f = float(value) return max(min_value, min(max_value, value_f)) def format_interval_name(start_sec: float, end_sec: float, frame_step: int): return f"s{int(math.floor(start_sec)):06d}_e{int(math.ceil(end_sec)):06d}_step{frame_step}" def sanitize_path_component(value: str): keep = [] for ch in value: if ch in '<>:"/\\|?*': keep.append("_") else: keep.append(ch) return "".join(keep).strip().rstrip(".") def normalize_relpath(path): return str(path).replace("\\", "/") def jpeg_qscale(quality_pct: int): quality_pct = max(70, min(100, int(quality_pct))) if quality_pct >= 98: return 1 if quality_pct >= 92: return 2 if quality_pct >= 88: return 3 if quality_pct >= 82: return 4 return 5 def run_cmd(cmd): result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: raise SystemExit( f"Command failed ({result.returncode}): {' '.join(cmd)}\n" f"stdout:\n{result.stdout}\n" f"stderr:\n{result.stderr}" ) return result.stdout if __name__ == "__main__": main()