Blame for app.py

"""
NextJump Broadcast — Hugging Face Space.

Upload a jump-rope video and get back a broadcast-ready overlay video (vertical 9:16 or
horizontal 16:9) with the AI jump count, speed gauge, splits, rhythm and — optionally —
official start/end beep detection. All of the work is done by claude_anim_v2.main();
this file is only the Gradio front end.

Layout expected next to this file (run sync_assets.py from the repo to populate it):
    claude_anim_v2.py, beep.WAV, elitejumps_logo.png, nextjump_logo.png, graphics/art/SalmaproMedium.otf,
    movenet-pytorch/{movenet/, _models/movenet_thunder.pth}
When run from inside the repo without syncing, the repo root is used instead.

The counting checkpoint is private and is NOT part of the Space: it is downloaded at startup from the
private model repo (like NextJump-Tournament/app.py), authenticated with the DATASET_SECRET Space secret.
    NEXTJUMP_MODEL_REPO   model repo id            (default lumos-motion/nextjump)
    NEXTJUMP_MODEL_FILE   ONNX file in that repo   (default nextjump_all_128.onnx)
    DATASET_SECRET / HF_TOKEN   read token for the private repo (locally, a `huggingface-cli login` also works)
    NEXTJUMP_ONNX         optional local path that skips the download
"""
import os
import sys
import time
import uuid
import base64
import tempfile

import cv2
import gradio as gr
from huggingface_hub import hf_hub_download

HERE = os.path.dirname(os.path.abspath(__file__))
for _root in (HERE, os.path.dirname(HERE)):
    if os.path.isfile(os.path.join(_root, 'claude_anim_v2.py')):
        sys.path.insert(0, _root)
        break
import claude_anim_v2 as nj  # noqa: E402

MODEL_REPO = os.environ.get('NEXTJUMP_MODEL_REPO', 'lumos-motion/nextjump')
MODEL_FILE = os.environ.get('NEXTJUMP_MODEL_FILE', 'nextjump_all_128.onnx')
HF_TOKEN = os.environ.get('DATASET_SECRET') or os.environ.get('HF_TOKEN')  # None -> cached CLI login


def resolve_model():
    """Path to the counting checkpoint, downloaded (and cached) from the private model repo."""
    override = os.environ.get('NEXTJUMP_ONNX')
    if override:
        return override
    try:
        path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, repo_type='model', token=HF_TOKEN)
        print(f'Counting model: {MODEL_REPO}/{MODEL_FILE} -> {path}')
        return path
    except Exception as e:
        local = os.path.join(nj.ROOT, 'checkpoint', 'onnx', MODEL_FILE)   # local dev inside the repo
        if os.path.isfile(local):
            print(f'Could not download {MODEL_REPO}/{MODEL_FILE} ({type(e).__name__}); using local {local}')
            return local
        raise RuntimeError(f'Could not download {MODEL_FILE} from the private repo {MODEL_REPO} '
                           f'({type(e).__name__}: {e}). If this is an auth error, add a read token for the repo '
                           'as the DATASET_SECRET secret in the Space settings.') from e


# Render worker processes (multiprocessing 'spawn') re-import this file as __mp_main__: they must not
# download models or build the UI again, so all of that is skipped there.
IS_RENDER_WORKER = __name__ == '__mp_main__'
ONNX_PATH = None if IS_RENDER_WORKER else resolve_model()   # at startup, so the first request doesn't wait
if not IS_RENDER_WORKER:
    try:   # also warm the torchvision person-detector weights used by smart framing (public, ~75 MB)
        nj.load_person_detector('coco', 'cpu')
    except Exception as e:
        print(f'Could not pre-load the person detector ({e}); it will be downloaded on first use')

MAX_SECONDS = float(os.environ.get('NEXTJUMP_MAX_SECONDS', 200))   # longest accepted upload
OUT_DIR = os.path.join(tempfile.gettempdir(), 'nextjump_space')
KEEP_OUTPUTS_S = 2 * 3600
FORMATS = {'Social (TikTok / Reels / Shorts, full screen)': 'social',
           'Broadcast vertical 9:16 (full stats panel)': 'vertical',
           'Broadcast horizontal 16:9 (YouTube / TV)': 'horizontal'}
DEFAULT_FILL = {'social': 0.25, 'vertical': 0.50, 'horizontal': 0.60}
FRAMING = {'Smart: find the most active jumper(s)': 'smart',
           'Pose: follow the first person found': 'pose',
           'Off: fixed centre crop': 'off'}


def probe(path):
    cap = cv2.VideoCapture(path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
    frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    w, h = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    ok = cap.isOpened() and cap.read()[0]
    cap.release()
    return dict(ok=ok, fps=fps, frames=frames, w=w, h=h, duration=frames / fps if fps else 0.0)


def cleanup_old_outputs():
    if not os.path.isdir(OUT_DIR):
        return
    now = time.time()
    for f in os.listdir(OUT_DIR):
        p = os.path.join(OUT_DIR, f)
        try:
            if now - os.path.getmtime(p) > KEEP_OUTPUTS_S:
                os.remove(p)
        except OSError:
            pass


def summary_markdown(s, info):
    lines = [f'## {s["final"]} jumps']
    if s['event'] is not None:
        t0, t1 = s['event']
        lines.append(f'Counted between the beeps: start **{t0:.2f} s**, end **{t1:.2f} s** '
                     f'({t1 - t0:.2f} s event). Nothing before the start beep or after the end beep is counted.')
    elif s['beep_requested']:
        lines.append('⚠️ Beep detection was on, but no start/end beep pair matching the event length was found, '
                     'so the whole video was counted.')
    else:
        lines.append(f'Counted over the detected jumping window ({s["duration"]:.1f} s).')
    if s.get('misses'):
        lines.append(f'\n**{s["misses"]}** miss{"es" if s["misses"] != 1 else ""} detected'
                     + (f'; the first jump after each miss was not counted (rulebook), '
                        f'**{s["skipped"]}** jump{"s" if s["skipped"] != 1 else ""} removed.' if s.get('skipped') else '.'))
    lines.append(f'\nAverage speed **{s["avg_jps"]:.2f}** jumps/s · peak **{s["peak_jps"]:.2f}** jumps/s')
    if s['splits']:
        lines.append(f'\n| {s["split_seconds"]} s split | ' + ' | '.join(str(k + 1) for k in range(len(s['splits']))) + ' |')
        lines.append('|---|' + '---|' * len(s['splits']))
        lines.append('| jumps | ' + ' | '.join(str(v) for v in s['splits']) + ' |')
    lines.append(f'\n<sub>Output {s["size"][0]}×{s["size"][1]} @ {s["fps"]:.2f} fps, {s["frames"]} frames · '
                 f'source {info["w"]}×{info["h"]} @ {info["fps"]:.2f} fps</sub>')
    return '\n'.join(lines)


def generate(video, fmt_label, title, theme_label, turbo, beep, event_length, framing_label, jumpers, skip_after_miss,
             count_on_crop, out_fps, fill, smooth, max_zoom, foot_rings, skeleton, show_pace, mirror,
             progress=gr.Progress()):
    if not video:
        raise gr.Error('Please upload a video first.')
    info = probe(video)
    if not info['ok']:
        raise gr.Error('Could not read that video file. Please upload an MP4 or MOV.')
    if info['duration'] > MAX_SECONDS:
        raise gr.Error(f'That video is {info["duration"]:.0f} s long; the limit here is {MAX_SECONDS:.0f} s. '
                       'Please trim it around the event.')
    if info['fps'] < 45:
        gr.Warning(f'This video is {info["fps"]:.0f} fps. The counting model was trained on 60 fps footage, '
                   'so counts may be less accurate.')
    if beep and not event_length:
        raise gr.Error('Please enter the event length (seconds between the beeps).')

    cleanup_old_outputs()
    os.makedirs(OUT_DIR, exist_ok=True)
    fmt = fmt_label if fmt_label in FORMATS.values() else FORMATS.get(fmt_label, 'social')
    base = os.path.splitext(os.path.basename(video))[0][:40] or 'video'
    output = os.path.join(OUT_DIR, f'{base}_nextjump_{fmt}_{uuid.uuid4().hex[:6]}.mp4')

    argv = ['--video', video, '--output', output, '--onnx', ONNX_PATH, '--format', fmt, '--title', title or '',
            '--theme', theme_label if theme_label in ('elitejumps', 'nextjump') else 'elitejumps',
            '--fill', str(fill), '--smooth', str(smooth), '--max_zoom', str(max_zoom)]
    if out_fps != 'Auto':
        argv += ['--out_fps', str(out_fps)]
    if turbo:
        argv.append('--turbo')
    framing = framing_label if framing_label in FRAMING.values() else FRAMING.get(framing_label, 'smart')
    argv += ['--framing', framing, '--jumpers', str(int(jumpers or 1))]
    if count_on_crop and framing == 'smart':
        argv.append('--count_on_crop')
    if skip_after_miss:
        argv.append('--skip_after_miss')
    if beep:
        argv += ['--beep', '--event_length', str(float(event_length))]
    if foot_rings:
        argv.append('--fx')
    if skeleton:
        argv.append('--skeleton')
    if not show_pace:
        argv.append('--no_pace')
    if mirror:
        argv.append('--mirror')

    progress(0, desc='Starting')
    try:
        s = nj.main(argv, progress=lambda f, desc: progress(f, desc=desc))
    except Exception as e:  # surface a readable message in the UI; full trace goes to the logs
        import traceback
        traceback.print_exc()
        raise gr.Error(f'Processing failed: {e}')
    return s['output'], s['output'], summary_markdown(s, info)


def _logo_html():
    path = os.path.join(nj.ROOT, nj.THEMES[nj.DEFAULT_THEME]['logo'])
    logo = ''
    if os.path.isfile(path):
        with open(path, 'rb') as f:
            logo = f'<img src="data:image/png;base64,{base64.b64encode(f.read()).decode()}" alt="NextJump">'
    return f"""
<div id="nj-header">
  {logo or f'<span class="nj-word">{nj.THEMES[nj.DEFAULT_THEME]["wordmark"]}</span>'}
  <p>Upload a jump-rope video and get a broadcast-ready graphic: AI jump count, speed gauge, splits,
  rhythm and auto-framing on the jumper. Turn on beep detection for official timed events.</p>
</div>"""


CSS = """
#nj-header { background:#111113; border-radius:14px; padding:22px 26px; margin-bottom:6px;
             border-left:6px solid #E52421; }
#nj-header img { height:64px; display:block; margin-bottom:12px; }
#nj-header .nj-word { color:#E52421; font-weight:800; font-size:30px; letter-spacing:1px; }
#nj-header p { color:#D4D4D8; margin:0; font-size:15px; line-height:1.45; }
"""

def build_demo():
    with gr.Blocks(title='NextJump Broadcast', css=CSS, theme=gr.themes.Soft(primary_hue='red'),
                   delete_cache=(3600, KEEP_OUTPUTS_S)) as demo:
        gr.HTML(_logo_html())
        with gr.Row(equal_height=False):
            with gr.Column(scale=1):
                video_in = gr.Video(label='Jump-rope video', sources=['upload'])
                fmt_in = gr.Radio([(k, v) for k, v in FORMATS.items()], value='social', label='Output format')
                title_in = gr.Textbox(label='Title', placeholder='e.g.  Ali  |  30s Speed', max_lines=1)
                theme_in = gr.Radio([('Elite Jumps', 'elitejumps'), ('NextJump', 'nextjump')], value='elitejumps',
                                    label='Graphics theme')
                turbo_in = gr.Checkbox(label='Turbo mode (instant replay)', value=True,
                                       info='About 2x faster: analyses and renders at 30 fps with lighter person detection.')
                with gr.Group():
                    beep_in = gr.Checkbox(label='Detect official start / end beeps', value=False,
                                          info='Counts only between the start beep and the end beep.')
                    event_len_in = gr.Number(label='Event length (seconds between beeps)', value=30,
                                             minimum=5, maximum=600, visible=False)
                with gr.Group():
                    framing_in = gr.Radio([(k, v) for k, v in FRAMING.items()], value='smart', label='Framing',
                                          info='Smart framing ignores judges and bystanders by checking who is actually jumping.')
                    jumpers_in = gr.Slider(1, 6, value=1, step=1, label='People to frame',
                                           info='1 for single rope; 2 for pairs; 3+ for double dutch (turners + jumpers).')
                skip_in = gr.Checkbox(label='Rulebook miss rule: don\'t count the first jump after a miss', value=False)
                with gr.Accordion('Advanced', open=False):
                    crop_count_in = gr.Checkbox(label='Count from a crop around the selected jumper(s)', value=False,
                                                info='Smart framing only. Can help when other people in frame are moving.')
                    fps_in = gr.Radio(['Auto', 30, 60], value='Auto', label='Output frame rate',
                                      info='Auto = 60 fps (30 fps in turbo mode). 30 fps renders about twice as fast.')
                    fill_in = gr.Slider(0.15, 0.85, value=DEFAULT_FILL['social'], step=0.01,
                                        label='Jumper size (fraction of frame height)',
                                        info='Smaller = wider shot. Capped at the widest crop the video allows.')
                    smooth_in = gr.Slider(0.2, 2.0, value=0.8, step=0.1, label='Camera smoothing (s)')
                    zoom_in = gr.Slider(1.0, 4.0, value=2.5, step=0.1, label='Max zoom')
                    fx_in = gr.Checkbox(label='Foot-strike rings', value=True,
                                        info='A ring on the floor at every landing (alternating feet when stepping), during the event.')
                    skeleton_in = gr.Checkbox(label='Show limb tracking (arms & legs)', value=False,
                                          info='Shown only during the event. Slower: runs the pose model on every frame.')
                    pace_in = gr.Checkbox(label='Show start-pace comparison', value=True)
                    mirror_in = gr.Checkbox(label='Mirror video', value=False)
                run_btn = gr.Button('Generate video', variant='primary', size='lg')
                gr.Markdown('<sub>Processing runs on the Space’s hardware and can take several minutes for a '
                            '30 s clip on CPU. Best results: 60 fps, whole body in frame, one jumper.</sub>')
            with gr.Column(scale=1):
                video_out = gr.Video(label='NextJump broadcast video', interactive=False)
                file_out = gr.File(label='Download')
                summary_out = gr.Markdown()

        beep_in.change(lambda on: gr.update(visible=on), beep_in, event_len_in)
        fmt_in.change(lambda fmt: gr.update(value=DEFAULT_FILL.get(fmt, 0.5)), fmt_in, fill_in)
        framing_in.change(lambda fr: gr.update(visible=fr == 'smart'), framing_in, jumpers_in)
        run_btn.click(generate,
                      inputs=[video_in, fmt_in, title_in, theme_in, turbo_in, beep_in, event_len_in, framing_in, jumpers_in,
                              skip_in, crop_count_in, fps_in, fill_in, smooth_in, zoom_in, fx_in, skeleton_in, pace_in, mirror_in],
                      outputs=[video_out, file_out, summary_out],
                      concurrency_limit=1)
    return demo


if not IS_RENDER_WORKER:
    demo = build_demo()
    demo.queue(max_size=16)

if __name__ == '__main__':
    demo.launch()