variants.py

# variants.py

import random
import subprocess

from config import (
    CLIPS_DIR,
    TARGET_WIDTH,
    TARGET_HEIGHT,
    TARGET_FPS,
    MIN_CLIP_DURATION,
    MAX_CLIP_DURATION,
    MIN_SOURCE_DURATION,
    VIDEO_CODEC,
    PRESET,
    CRF,
    VARIANT_RULES,
)

from probe import get_duration


# =====================================================
# VARIANT NAMES
# =====================================================

VARIANT_NAMES = [
    "A", "B", "C", "D", "E",
    "F", "G", "H", "I", "J"
]


# =====================================================
# RANDOM CLIP DURATION
# =====================================================

def random_clip_duration():

    return round(
        random.uniform(
            MIN_CLIP_DURATION,
            MAX_CLIP_DURATION
        ),
        2
    )


# =====================================================
# VARIANT COUNT
# =====================================================

def get_variant_count(duration):

    for min_dur, max_dur, count in VARIANT_RULES:

        if min_dur <= duration < max_dur:
            return count

    return 3


# =====================================================
# BUILD DYNAMIC ZONES
# =====================================================

def build_dynamic_zones(
    duration,
    count
):

    zone_size = duration / count

    zones = []

    current = 0

    for _ in range(count):

        zones.append(
            (
                current,
                current + zone_size
            )
        )

        current += zone_size

    return zones


# =====================================================
# RANDOM START INSIDE ZONE
# =====================================================

def random_start_in_zone(
    zone_start,
    zone_end,
    clip_duration
):

    available = (
        zone_end
        - zone_start
        - clip_duration
    )

    if available <= 0:
        return zone_start

    return round(
        random.uniform(
            zone_start,
            zone_start + available
        ),
        2
    )


# =====================================================
# RENDER SINGLE VARIANT
# =====================================================

def render_variant(
    source_file,
    output_file,
    start_time,
    clip_duration,
    hflip=False
):

    filters = [

        (
            f"scale={TARGET_WIDTH}:{TARGET_HEIGHT}:"
            "force_original_aspect_ratio=increase"
        ),

        (
            f"crop={TARGET_WIDTH}:{TARGET_HEIGHT}"
        ),

        f"fps={TARGET_FPS}",

        "fade=t=in:st=0:d=0.2",

        (
            f"fade=t=out:"
            f"st={max(0, clip_duration - 0.2)}:d=0.2"
        ),
    ]

    if hflip:
        filters.append("hflip")

    vf = ",".join(filters)

    cmd = [

        "ffmpeg",
        "-y",

        "-ss",
        str(start_time),

        "-i",
        str(source_file),

        "-t",
        str(clip_duration),

        "-an",

        "-vf",
        vf,

        "-c:v",
        VIDEO_CODEC,

        "-preset",
        PRESET,

        "-crf",
        str(CRF),

        "-pix_fmt",
        "yuv420p",

        str(output_file),
    ]

    subprocess.run(
        cmd,
        check=True,
        capture_output=True
    )

    return str(output_file)


# =====================================================
# GENERATE VARIANTS
# =====================================================

def generate_variants(
    source_file,
    source_id
):

    duration = get_duration(
        source_file
    )

    if duration < MIN_SOURCE_DURATION:
        return []

    variant_count = get_variant_count(
        duration
    )

    zones = build_dynamic_zones(
        duration,
        variant_count
    )

    scene_pool = []

    for idx, zone in enumerate(zones):

        variant_name = VARIANT_NAMES[idx]

        clip_duration = random_clip_duration()

        start_time = random_start_in_zone(
            zone[0],
            zone[1],
            clip_duration
        )

        # --------------------------------
        # hflip setiap variant ke-3
        # --------------------------------

        hflip = (
            (idx + 1) % 3 == 0
        )

        output_file = (
            CLIPS_DIR
            / f"{source_id}_{variant_name}.mp4"
        )

        render_variant(
            source_file=source_file,
            output_file=output_file,
            start_time=start_time,
            clip_duration=clip_duration,
            hflip=hflip
        )

        scene_pool.append(

            {
                "source_id": source_id,

                "family": source_id,

                "variant": variant_name,

                "path": str(output_file),

                "duration": clip_duration,
            }

        )

    return scene_pool


# =====================================================
# GENERATE MANY
# =====================================================

def generate_many(
    source_files
):

    all_scenes = []

    for idx, source_file in enumerate(
        source_files,
        start=1
    ):

        try:

            scenes = generate_variants(
                source_file,
                idx
            )

            all_scenes.extend(
                scenes
            )

            print(
                f"[Variants] "
                f"Source {idx} "
                f"-> {len(scenes)} variants"
            )

        except Exception as e:

            print(
                f"[Variants] Error:"
                f" {source_file}"
            )

            print(e)

    return all_scenes


# =====================================================
# TEST
# =====================================================

if __name__ == "__main__":

    scenes = generate_variants(
        "downloads/source_0001.mp4",
        1
    )

    print()

    print(
        f"Generated: "
        f"{len(scenes)} scenes"
    )

    print()

    for scene in scenes:

        print(scene)