scripts/edit_comparison.py

"""Compare four complete recorded painting attempts in one 30-second video.

Each attempt uses its entire action sequence, sped up uniformly, followed by a
brief 1x final view. The last six seconds compare final camera-derived canvases.
No paint textures, manufactured intermediate frames, or evaluator scores are used.
"""
from __future__ import annotations

import argparse
import bisect
import json
import math
from pathlib import Path
import shutil
import subprocess

import cv2
import numpy as np
from PIL import Image, ImageDraw

import edit_video as layout

ROOT = Path(__file__).resolve().parents[1]
FPS = 30
PER_ATTEMPT_FRAMES = 180
COMPARISON_FRAMES = 180
TOTAL_FRAMES = 4 * PER_ATTEMPT_FRAMES + COMPARISON_FRAMES
DEFAULT_LABELS = ['Initial timing', 'Remove endpoint pauses', 'More paint / slower sweep', 'Lighter tower pressure']
FRAME_BYTES = 960 * 720 * 3


def plan_attempt(run_dir, number, change):
    timeline = json.loads((run_dir / 'timeline.json').read_text())
    source_duration = float(timeline['duration_seconds'])
    events = sorted(timeline['events'], key=lambda e: float(e['time_seconds']))
    finished_events = [e for e in events if e.get('type') == 'camera' and e.get('label') == 'finished']
    if not finished_events:
        raise ValueError(f'{run_dir}: no finished camera marker; use completed attempts for this comparison.')
    finished = float(finished_events[-1]['time_seconds'])
    if not 0 < finished < source_duration:
        raise ValueError(f'{run_dir}: invalid finished timestamp.')
    speed = max(1, math.ceil(finished / 5.5))
    moving_frames = min(165, max(1, round(finished / speed * FPS)))
    final_frames = PER_ATTEMPT_FRAMES - moving_frames
    final_duration = final_frames / FPS
    final_source_end = min(source_duration, finished + final_duration)
    messages = [dict(time_seconds=0., title='Preparing the brush',
                     detail='Load paint and locate the paper in the camera.')]

    def first_event(*words):
        return next((float(e['time_seconds']) for e in events if e.get('type') == 'move'
                     and any(word in e.get('label', '').lower() for word in words)), None)

    tower = first_event('tower', 'eiffel')
    water = first_event('bank', 'seine', 'river')
    if tower is not None:
        messages.append(dict(time_seconds=tower / speed,
                             title='Painting the Eiffel Tower',
                             detail='Trace the tower and add the surrounding details.'))
    if water is not None:
        messages.append(dict(time_seconds=water / speed, title='Painting the Seine',
                             detail='Add the riverbanks and flowing water.'))
    messages.append(dict(time_seconds=moving_frames / FPS, title='Attempt finished',
                         detail='The recorded camera view after the brush is lifted.'))
    messages.sort(key=lambda message: message['time_seconds'])
    return dict(number=number, change=change, source_duration_seconds=source_duration,
                playback_speed=speed, moving_frames=moving_frames, frame_count=PER_ATTEMPT_FRAMES,
                final_frame_hold_seconds=max(0., final_duration - (final_source_end - finished)),
                segments=[dict(source_start=0., source_end=finished, speed=speed, output_frames=moving_frames),
                          dict(source_start=finished, source_end=final_source_end, speed=1, output_frames=final_frames)],
                messages=messages)


def attempt_layout(plan, final=False):
    speed_label = '1× SPEED' if final else f'{plan["playback_speed"]}× SPEED'
    if final and plan['final_frame_hold_seconds'] > 1 / FPS:
        speed_label += ' · FINAL HOLD'
    image = layout.render_layout(speed_label)
    draw = ImageDraw.Draw(image)
    draw.rectangle((20, 15, 1250, 77), fill=layout.BG)
    title = f'Attempt {plan["number"]:02d} · {plan["change"]}'
    draw.text((32, 30), title, font=layout.font(28, True), fill=layout.WHITE)
    return image


def final_rgb(ffmpeg, source):
    result = subprocess.run([ffmpeg, '-hide_banner', '-loglevel', 'error', '-sseof', '-0.15',
                             '-i', str(source), '-vf', 'scale=960:720',
                             '-f', 'rawvideo', '-pix_fmt', 'rgb24', 'pipe:1'],
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
    if len(result.stdout) < FRAME_BYTES:
        raise RuntimeError(f'Could not decode the final camera image from {source}')
    return np.frombuffer(result.stdout[-FRAME_BYTES:], dtype=np.uint8).reshape(720, 960, 3).copy()


def comparison_frame(images, plans, selected=None):
    image = Image.new('RGB', (1920, 1080), layout.BG)
    draw = ImageDraw.Draw(image)
    draw.text((32, 24), 'Four attempts · final camera views', font=layout.font(34, True), fill=layout.WHITE)
    label = 'SIMPLIFIED WET-PAINT MODEL'
    draw.text((1888 - draw.textlength(label, font=layout.font(18, True)), 35),
              label, font=layout.font(18, True), fill=layout.SECONDARY)
    for index, (canvas, plan) in enumerate(zip(images, plans)):
        x = 32 + (index % 2) * 944
        y = 100 + (index // 2) * 470
        draw.rounded_rectangle((x, y, x + 911, y + 443), radius=17,
                               fill=layout.PANEL_BG, outline=layout.BORDER, width=2)
        image.paste(Image.fromarray(canvas).resize((416, 416), Image.Resampling.LANCZOS), (x + 16, y + 14))
        draw = ImageDraw.Draw(image)
        draw.text((x + 466, y + 43), f'Attempt {plan["number"]:02d}', font=layout.font(33, True), fill=layout.WHITE)
        layout.text_block(draw, plan['change'], x + 466, y + 115, 408, 31,
                          color=layout.SECONDARY, bold=True, line_height=40, max_lines=4)
        if selected == plan['number']:
            draw.text((x + 466, y + 313), 'SELECTED RESULT', font=layout.font(20, True), fill=layout.MINT)
        draw.text((x + 466, y + 379), 'Final RGB camera', font=layout.font(22), fill=layout.SECONDARY)
    draw.text((32, 1039), 'Same scripted picture. Each canvas is rectified from its run’s final front-camera frame.',
              font=layout.font(21), fill=layout.SECONDARY)
    return image


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--run-dirs', type=Path, nargs=4,
                        default=[ROOT / f'results/attempt-{i:02d}' for i in range(1, 5)])
    parser.add_argument('--labels', nargs=4, default=DEFAULT_LABELS,
                        help='Four concise, factual descriptions of the changes between attempts.')
    parser.add_argument('--output', type=Path, default=ROOT / 'results/comparison.mp4')
    parser.add_argument('--selected', type=int, choices=(1, 2, 3, 4),
                        help='Optional explicitly selected result; no winner is inferred automatically.')
    parser.add_argument('--plan-only', action='store_true')
    args = parser.parse_args()
    plans = [plan_attempt(run_dir, i + 1, label)
             for i, (run_dir, label) in enumerate(zip(args.run_dirs, args.labels))]
    scratch = ROOT / 'work/video/comparison'
    scratch.mkdir(parents=True, exist_ok=True)
    manifest = dict(duration_seconds=30., width=1920, height=1080, fps=FPS,
                    frame_count=TOTAL_FRAMES, attempts=plans, selected_attempt=args.selected,
                    final_comparison_seconds=6.,
                    source='Full action sequences and rectified final images from recorded RGB camera videos.')
    (scratch / 'edit-manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')
    print('Comparison: four 6-second runs + 6-second camera comparison = 30 seconds.', flush=True)
    print('Action speeds:', ', '.join(f'{p["number"]:02d}: {p["playback_speed"]}×' for p in plans), flush=True)
    if args.plan_only:
        return
    ffmpeg = shutil.which('ffmpeg')
    if not ffmpeg:
        raise RuntimeError('ffmpeg is required.')
    homographies = [np.asarray(json.loads((run_dir / 'camera-localization.json').read_text())['camera_to_canvas'], dtype=float)
                    for run_dir in args.run_dirs]
    sources = [run_dir / 'painting-run.mp4' for run_dir in args.run_dirs]
    for source, homography in zip(sources, homographies):
        if not source.is_file():
            raise FileNotFoundError(source)
        if homography.shape != (3, 3) or not np.isfinite(homography).all():
            raise ValueError('Every run must have a finite 3×3 camera homography.')
    canvases = []
    for i, (source, homography) in enumerate(zip(sources, homographies)):
        rgb = final_rgb(ffmpeg, source)
        canvas = cv2.warpPerspective(rgb, homography, (512, 512), flags=cv2.INTER_CUBIC)
        Image.fromarray(rgb).save(scratch / f'attempt-{i+1:02d}-final-camera.png')
        Image.fromarray(canvas).save(scratch / f'attempt-{i+1:02d}-final-canvas.png')
        canvases.append(canvas)
    comparison = comparison_frame(canvases, plans, args.selected)
    comparison.save(scratch / 'comparison-preview.png')
    args.output.parent.mkdir(parents=True, exist_ok=True)
    temporary = args.output.with_name(args.output.stem + '.rendering.mp4')
    encoder = decoder = None
    blank = Image.new('RGB', (layout.CAPTION_RECT[2] - layout.CAPTION_RECT[0],
                              layout.CAPTION_RECT[3] - layout.CAPTION_RECT[1]), layout.PANEL_BG)
    with (scratch / 'encoder.log').open('w') as encode_log:
        try:
            encoder = subprocess.Popen([ffmpeg, '-hide_banner', '-loglevel', 'error', '-y',
                '-f', 'rawvideo', '-pixel_format', 'rgb24', '-video_size', '1920x1080', '-framerate', str(FPS),
                '-i', 'pipe:0', '-an', '-c:v', 'libx264', '-preset', 'fast', '-crf', '14',
                '-pix_fmt', 'yuv420p', '-movflags', '+faststart', '-frames:v', str(TOTAL_FRAMES), str(temporary)],
                stdin=subprocess.PIPE, stderr=encode_log)
            for plan, source, homography in zip(plans, sources, homographies):
                filters = ['[0:v]split=2[s0][s1]']
                for i, segment in enumerate(plan['segments']):
                    filters.append(f'[s{i}]trim=start={segment["source_start"]:.9f}:end={segment["source_end"]:.9f},'
                        f'setpts=(PTS-STARTPTS)/{segment["speed"]},fps={FPS},'
                        f'tpad=stop_mode=clone:stop_duration=6,trim=end_frame={segment["output_frames"]},'
                        f'setpts=PTS-STARTPTS[v{i}]')
                filters.append('[v0][v1]concat=n=2:v=1:a=0,scale=960:720,format=rgb24[out]')
                panels = [layout.render_panel(message) for message in plan['messages']]
                times = [message['time_seconds'] for message in plan['messages']]
                backgrounds = [attempt_layout(plan), attempt_layout(plan, final=True)]
                with (scratch / f'decoder-{plan["number"]:02d}.log').open('w') as decode_log:
                    decoder = subprocess.Popen([ffmpeg, '-hide_banner', '-loglevel', 'error', '-i', str(source),
                        '-filter_complex', ';'.join(filters), '-map', '[out]', '-frames:v', str(PER_ATTEMPT_FRAMES),
                        '-f', 'rawvideo', '-pix_fmt', 'rgb24', 'pipe:1'], stdout=subprocess.PIPE, stderr=decode_log)
                    for index in range(PER_ATTEMPT_FRAMES):
                        raw = decoder.stdout.read(FRAME_BYTES)
                        if len(raw) != FRAME_BYTES:
                            raise RuntimeError(f'Attempt {plan["number"]} decoder stopped at frame {index}.')
                        rgb = np.frombuffer(raw, dtype=np.uint8).reshape(720, 960, 3)
                        t = index / FPS
                        active = max(0, bisect.bisect_right(times, t) - 1)
                        panel = panels[active]
                        elapsed = t - times[active]
                        if active and elapsed < .20:
                            panel = panel.copy()
                            if elapsed < .08:
                                caption = Image.blend(panels[active - 1].crop(layout.CAPTION_RECT), blank, elapsed / .08)
                            else:
                                caption = Image.blend(blank, panels[active].crop(layout.CAPTION_RECT), (elapsed - .08) / .12)
                            panel.paste(caption, layout.CAPTION_RECT[:2])
                        frame = layout.compose_frame(rgb, backgrounds[index >= plan['moving_frames']], panel, homography)
                        encoder.stdin.write(frame.tobytes())
                        if index == 80:
                            frame.save(scratch / f'attempt-{plan["number"]:02d}-action-preview.png')
                    decoder.stdout.close()
                    if decoder.wait(timeout=60):
                        raise RuntimeError(f'Attempt {plan["number"]} decoder failed.')
                    decoder = None
                print(f'Rendered attempt {plan["number"]}/4.', flush=True)
            final_bytes = comparison.tobytes()
            for _ in range(COMPARISON_FRAMES):
                encoder.stdin.write(final_bytes)
            encoder.stdin.close()
            if encoder.wait(timeout=60):
                raise RuntimeError('Comparison encoding failed.')
            temporary.replace(args.output)
            print(f'Saved {args.output}: 30s, 1920×1080, 30 fps.', flush=True)
        except BaseException:
            for process in (decoder, encoder):
                if process is not None and process.poll() is None:
                    process.kill()
                    process.wait()
            temporary.unlink(missing_ok=True)
            raise


if __name__ == '__main__':
    main()