MAI/dataset_prep/prepare_cvat_dataset.py

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()