| """Paired graybox / styled team-battle clips from matrix-shooter-js. |
| |
| usage: matrix_capture.py plan --game CAPTURE_DIR --out DIR [--clips 300] [--block 10] |
| matrix_capture.py work --game CAPTURE_DIR --out DIR --map paris [--shard 0 --shards 1] [--port 8131] |
| |
| CAPTURE_DIR is the output of matrix_capture_game.py. One Chromium (GPU, ANGLE) |
| opens the game four times, once per look (original graybox, cartoon, war, |
| punk), and keeps those pages for every clip of its map. Every frame the |
| scripted player decides its input from the graybox page's state; the graybox |
| page steps 1/24 s and renders with its own camera, then the three styled pages |
| step with the same input and that camera. The simulation is deterministic, so |
| all four show the same moment; every page's state is compared each frame and |
| the largest drift is recorded. |
| |
| plan writes DIR/jobs.jsonl: per map --clips clips spread over the walkable |
| --block-metre blocks (every block before any repeats), 20% at landmarks, eight |
| yaw sectors and three pitch tiers, team sizes, fight distance and the side the |
| enemy comes from, all seeded by the clip number. |
| |
| DIR/<map>/clip_NNNNNN/ |
| original.mp4 cartoon.mp4 war.mp4 punk.mp4 960x544, 24 fps, x264 yuv444 crf 12 |
| cond/LLLLLL.png graybox frame for latent L (frame 0, then 4L-3), lossless |
| frames.jsonl per frame: input, graybox state, camera, drift across pages |
| meta.json the job, renderer, drift and fight statistics |
| """ |
| import argparse |
| import asyncio |
| import base64 |
| import json |
| import math |
| import os |
| import random |
| import subprocess |
| import time |
| from pathlib import Path |
|
|
| from playwright.async_api import async_playwright |
|
|
| LOOKS = ('original', 'cartoon', 'war', 'punk') |
| MAPS = ('paris', 'tsinghua') |
| FPS = 24 |
| SECONDS = 10 |
| STYLED_FRAMES = FPS * SECONDS |
| W, H = 960, 544 |
| CHROME = os.environ.get('GBDATA_CHROMIUM', '/root/autodl-tmp/edenlive/cache/chrome-156/chrome-linux64/chrome') |
| |
| |
| GPU_ARGS = ['--enable-gpu', '--ignore-gpu-blocklist', '--use-angle=vulkan', '--enable-features=Vulkan'] |
| |
| RIGHT = {False: 1.9, True: 1.1} |
| UP = {'paris': .65, 'tsinghua': .15} |
| TURN = 2.6 / FPS |
| TAPS = {'shift', ' ', 'r', 'f', '1', '2', '3'} |
| |
| TEAMS = [((0, 1), 15), ((0, 2), 10), ((0, 3), 10), ((1, 2), 15), ((2, 3), 15), ((3, 4), 8), ((4, 5), 7), |
| ((1, 5), 10), ((0, 0), 10)] |
| DISTANCE = {'close': ((3, 6), (4, 9), .3), 'mid': ((8, 15), (10, 20), .35), 'far': ((16, 26), (20, 32), .35)} |
| |
| APPROACH = {'close': (3, 8), 'mid': (6, 16), 'far': (14, 28)} |
| SIDES = {'front': (0.0, .5), 'left': (-math.pi / 2, .15), 'right': (math.pi / 2, .15), 'back': (math.pi, .2)} |
|
|
|
|
| def wrap(a): |
| return (a + math.pi) % (2 * math.pi) - math.pi |
|
|
|
|
| def pick(rng, table): |
| keys = list(table) |
| return rng.choices(keys, [table[k][-1] for k in keys])[0] |
|
|
|
|
| class CombatPolicy: |
| """A third-person shooter player: face the nearest enemy, line the crosshair |
| up (the camera sits right of the eye, so the aim leads by asin(offset/d)), |
| fire in bursts, strafe, dash, jump, crouch, switch guns, reload, throw |
| grenades, look away and back, and back off when a wall squeezes the camera.""" |
|
|
| def __init__(self, seed, map_id, pitch_bias, distance='mid'): |
| self.rng = random.Random(seed) |
| self.map = map_id |
| self.retreat, self.approach = APPROACH[distance] |
| self.pitch_bias = pitch_bias |
| self.held = set() |
| self.strafe, self.strafe_left = None, 0 |
| self.look_away, self.look_left = 0.0, 0 |
| self.burst_left, self.pause_left = 0, 0 |
| self.ads_left, self.squeezed, self.escape_left = 0, 0, 0 |
| self.yaw = self.pitch = None |
|
|
| def step(self, state, camera, frame): |
| rng = self.rng |
| p = state['player'] |
| if self.yaw is None: |
| self.yaw, self.pitch = state['yaw'], state['pitch'] |
| foes = [a for a in state['actors'] if a['team'] == 1 and a['alive']] |
| want, taps = set(), [] |
| ads = self.ads_left > 0 |
| self.ads_left = max(0, self.ads_left - 1) |
| |
| close = camera and camera.get('s') and camera['s']['distance'] < 2.4 |
| self.squeezed = self.squeezed + 1 if close else 0 |
| if self.squeezed > 8 and self.escape_left <= 0: |
| self.escape_left, self.look_away = 30, rng.choice([-1, 1]) * rng.uniform(2.2, 3.1) |
| self.look_left = 30 |
| if self.look_left <= 0 and rng.random() < .005: |
| self.look_away = rng.choice([-1, 1]) * rng.uniform(.8, 2.6) |
| self.look_left = rng.randint(18, 40) |
| away = 0.0 |
| if self.look_left > 0: |
| self.look_left -= 1 |
| away = self.look_away |
| if foes and self.escape_left <= 0: |
| foe = min(foes, key=lambda a: math.hypot(a['x'] - p['x'], a['z'] - p['z'])) |
| d = math.hypot(foe['x'] - p['x'], foe['z'] - p['z']) |
| bearing = math.atan2(foe['x'] - p['x'], foe['z'] - p['z']) |
| aim_yaw = bearing - math.asin(min(.95, RIGHT[ads] / max(d, 1e-3))) + away |
| aim_pitch = math.atan2(foe['y'] + 1.1 - (p['y'] + p['h'] + 1.65 + UP[self.map]), d + 4.5) |
| if away: |
| aim_pitch += self.pitch_bias |
| err = wrap(aim_yaw - self.yaw) |
| self.yaw += max(-TURN, min(TURN, err)) |
| self.pitch += max(-.04, min(.04, aim_pitch - self.pitch)) |
| |
| if self.burst_left > 0: |
| self.burst_left -= 1 |
| if abs(err) < .2: |
| want.add('j') |
| elif self.pause_left > 0: |
| self.pause_left -= 1 |
| elif abs(err) < .12: |
| self.burst_left, self.pause_left = rng.randint(10, 40), rng.randint(2, 10) |
| if ads or (abs(err) < .1 and d > 15 and rng.random() < .03): |
| want.add('k') |
| if not ads: |
| self.ads_left = rng.randint(24, 60) |
| if d > self.approach: |
| want.add('w') |
| elif d < self.retreat: |
| want.add('s') |
| if rng.random() < .012 and 6 < d < 22 and state['grenades'] == 0: |
| taps.append('f') |
| else: |
| |
| self.escape_left = max(0, self.escape_left - 1) |
| target = self.yaw + (away if self.escape_left else TURN * .5 * math.sin(frame / 31)) |
| self.yaw += max(-TURN, min(TURN, wrap(target - self.yaw))) |
| self.pitch += max(-.03, min(.03, self.pitch_bias + .08 * math.sin(frame / 23) - self.pitch)) |
| want.add('w') |
| if foes and rng.random() < .5: |
| want.add('j') |
| if self.strafe_left <= 0: |
| self.strafe = rng.choice([None, 'a', 'd', 'a', 'd']) |
| self.strafe_left = rng.randint(10, 36) |
| self.strafe_left -= 1 |
| if self.strafe and not self.escape_left: |
| want.add(self.strafe) |
| if rng.random() < .008: |
| taps.append('shift') |
| if rng.random() < .006: |
| taps.append(' ') |
| if ('c' in self.held and rng.random() < .97) or ('c' not in self.held and rng.random() < .004): |
| want.add('c') |
| if frame == 0 or rng.random() < .015: |
| taps.append(rng.choice(['1', '2', '3'])) |
| if state['ammo'] <= 2 and not state['reloading'] and rng.random() < .2: |
| taps.append('r') |
| down = sorted((want | set(taps)) - self.held) |
| up = sorted((self.held - want) | (self.held & TAPS)) |
| self.held = (self.held - set(up)) | set(down) |
| return dict(down=down, up=up, yaw=self.yaw, pitch=self.pitch) |
|
|
|
|
| def make_jobs(map_id, blocks, points, clips, offset): |
| jobs = [] |
| order = random.Random(offset).sample(range(len(blocks)), len(blocks)) |
| for i in range(clips): |
| rng = random.Random(offset + i) |
| if rng.random() < .2 and points: |
| point = rng.choice(points) |
| x, z, anchor = point['x'], point['z'], point['id'] |
| yaw = (point['yaw'] if point['yaw'] is not None else 0.0) + rng.uniform(-.6, .6) |
| else: |
| x, z = blocks[order[i % len(blocks)]] |
| x, z, anchor = x + rng.uniform(-3, 3), z + rng.uniform(-3, 3), None |
| yaw = ((i * 3 + i // len(blocks)) % 8) * math.pi / 4 + rng.uniform(-.3, .3) |
| (allies, enemies) = rng.choices([t for t, _ in TEAMS], [w for _, w in TEAMS])[0] |
| distance, side = pick(rng, DISTANCE), pick(rng, SIDES) |
| engage, spread = DISTANCE[distance][0], DISTANCE[distance][1] |
| bearing = yaw + SIDES[side][0] |
| enemy_at = [] |
| for _ in range(enemies): |
| r, lateral = rng.uniform(*spread), rng.uniform(-5, 5) |
| enemy_at.append([x + math.sin(bearing) * r + math.cos(bearing) * lateral, z + math.cos(bearing) * r - math.sin(bearing) * lateral]) |
| ally_at = [] |
| for _ in range(allies): |
| b, r = yaw + rng.choice([-1, 1]) * rng.uniform(.9, 1.9), rng.uniform(2.5, 6) |
| ally_at.append([x + math.sin(b) * r, z + math.cos(b) * r]) |
| jobs.append(dict(map=map_id, clip=i, seed=offset + i, anchor=anchor, teams=[allies + 1, enemies], distance=distance, side=side, |
| pitch_bias=rng.choice([-.22, -.04, .18, .3]), |
| setup=dict(seed=offset + i + 1, width=W, height=H, engage=list(engage), |
| player=dict(x=x, z=z, yaw=yaw, pitch=-.04), allies=ally_at, enemies=enemy_at))) |
| return jobs |
|
|
|
|
| async def open_pages(browser, port, map_id): |
| pages = [] |
| for look in LOOKS: |
| page = await browser.new_page(viewport={'width': W, 'height': H}, device_scale_factor=1) |
| await page.add_init_script('window.__CAPTURE=true') |
| page.on('pageerror', lambda e, look=look: print(f'[{look}] pageerror {e}', flush=True)) |
| for _ in range(100): |
| try: |
| await page.goto(f'http://127.0.0.1:{port}/?map={map_id}&style={look}') |
| break |
| except Exception: |
| await asyncio.sleep(.3) |
| pages.append(page) |
| for page in pages: |
| await page.wait_for_function('window.__cap && window.__cap.ready()', timeout=300_000) |
| return pages |
|
|
|
|
| def encoder(path): |
| return subprocess.Popen(['ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-f', 'image2pipe', '-c:v', 'png', '-r', str(FPS), |
| '-i', '-', '-c:v', 'libx264', '-preset', 'veryfast', '-crf', '12', '-pix_fmt', 'yuv444p', |
| '-threads', '2', str(path)], stdin=subprocess.PIPE) |
|
|
|
|
| def drift(states): |
| ref, worst = states[0], 0.0 |
| for s in states[1:]: |
| worst = max(worst, abs(s['player']['x'] - ref['player']['x']), abs(s['player']['z'] - ref['player']['z']), |
| abs(s['bullets'] - ref['bullets'])) |
| for a, b in zip(s['actors'], ref['actors']): |
| worst = max(worst, abs(a['x'] - b['x']), abs(a['z'] - b['z']), abs(a['health'] - b['health']), |
| float(a['alive'] != b['alive']) * 100) |
| return worst |
|
|
|
|
| async def capture(pages, job, out, renderer): |
| tmp = out.with_name(out.name + '.part') |
| if tmp.exists(): |
| for f in sorted(tmp.rglob('*'), reverse=True): |
| f.unlink() if f.is_file() else f.rmdir() |
| (tmp / 'cond').mkdir(parents=True) |
| setups = [await page.evaluate('cfg => window.__cap.setup(cfg)', job['setup']) for page in pages] |
| policy = CombatPolicy(job['seed'], job['map'], job['pitch_bias'], job['distance']) |
| state, camera = setups[0], None |
| encoders = {look: encoder(tmp / f'{look}.mp4') for look in LOOKS} |
| cond = {0} | {4 * l - 3 for l in range(1, 1 + (FPS * SECONDS - 1) // 4)} |
| log = (tmp / 'frames.jsonl').open('w') |
| worst, started, bullet_frames, actor_frames = 0.0, time.perf_counter(), 0, [] |
| for frame in range(FPS * SECONDS): |
| action = policy.step(state, camera, frame) |
| gray = await pages[0].evaluate('a => window.__cap.step(a)', dict(action, capture=True)) |
| |
| styled = await asyncio.gather(*[p.evaluate('a => window.__cap.step(a)', dict(action, capture=True, cam=gray['camera'])) |
| for p in pages[1:]]) if frame < STYLED_FRAMES else [] |
| for look, result in zip(LOOKS, [gray, *styled]): |
| png = base64.b64decode(result['image'].split(',', 1)[1]) |
| encoders[look].stdin.write(png) |
| if look == 'original' and frame in cond: |
| (tmp / 'cond' / f'{sorted(cond).index(frame):06d}.png').write_bytes(png) |
| state, camera = gray['state'], gray['camera'] |
| gap = drift([gray['state']] + [s['state'] for s in styled]) |
| worst = max(worst, gap) |
| bullet_frames += state['bullets'] > 0 |
| actor_frames.append(sum(a['alive'] for a in state['actors'])) |
| log.write(json.dumps(dict(frame=frame, input=action, state=state, camera=camera, drift=round(gap, 6))) + '\n') |
| log.close() |
| for proc in encoders.values(): |
| proc.stdin.close() |
| proc.wait() |
| meta = dict(job, frames=FPS * SECONDS, fps=FPS, width=W, height=H, looks=LOOKS, styled_frames=STYLED_FRAMES, renderer=renderer, |
| cond_frames=sorted(cond), max_drift=worst, bullet_frame_share=round(bullet_frames / (FPS * SECONDS), 3), |
| mean_alive_bots=round(sum(actor_frames) / len(actor_frames), 2), shots=state['shots'], hits=state['hits'], |
| kills=state['kills'], wall_s=round(time.perf_counter() - started, 1)) |
| (tmp / 'meta.json').write_text(json.dumps(meta, indent=1, ensure_ascii=False)) |
| tmp.rename(out) |
| return meta |
|
|
|
|
| async def run(args): |
| server = subprocess.Popen(['python3', str(args.game / 'start.py'), '--port', str(args.port)], |
| stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) |
| try: |
| async with async_playwright() as pw: |
| browser = await pw.chromium.launch(executable_path=CHROME, headless=True, args=GPU_ARGS) |
| if args.mode == 'plan': |
| jobs = [] |
| for m, map_id in enumerate(MAPS): |
| pages = await open_pages(browser, args.port, map_id) |
| blocks = await pages[0].evaluate('s => window.__cap.blocks(s)', args.block) |
| points = await pages[0].evaluate('window.__cap.points()') |
| jobs += make_jobs(map_id, blocks, points, args.clips, args.seed_base * (m + 1)) |
| print(json.dumps(dict(map=map_id, blocks=len(blocks), landmarks=len(points), clips=args.clips)), flush=True) |
| for page in pages: |
| await page.close() |
| args.out.mkdir(parents=True, exist_ok=True) |
| (args.out / 'jobs.jsonl').write_text(''.join(json.dumps(j) + '\n' for j in jobs)) |
| return |
| jobs = [json.loads(l) for l in (args.out / 'jobs.jsonl').read_text().splitlines() if l.strip()] |
| jobs = [j for j in jobs if j['map'] == args.map][args.shard::args.shards] |
| pages = await open_pages(browser, args.port, args.map) |
| renderer = await pages[0].evaluate("(() => {const g=document.createElement('canvas').getContext('webgl2');" |
| "const e=g.getExtension('WEBGL_debug_renderer_info');return g.getParameter(e.UNMASKED_RENDERER_WEBGL);})()") |
| print(json.dumps(dict(map=args.map, shard=args.shard, jobs=len(jobs), renderer=renderer)), flush=True) |
| for job in jobs: |
| out = args.out / job['map'] / f"clip_{job['clip']:06d}" |
| if out.exists(): |
| continue |
| meta = await capture(pages, job, out, renderer) |
| print(json.dumps(dict(clip=f"{job['map']}/{out.name}", teams=job['teams'], distance=job['distance'], |
| drift=meta['max_drift'], bullets=meta['bullet_frame_share'], shots=meta['shots'], |
| kills=meta['kills'], wall_s=meta['wall_s'])), flush=True) |
| await browser.close() |
| finally: |
| server.terminate() |
|
|
|
|
| def main(): |
| global SECONDS, STYLED_FRAMES |
| parser = argparse.ArgumentParser() |
| parser.add_argument('mode', choices=('plan', 'work')) |
| parser.add_argument('--game', type=Path, required=True) |
| parser.add_argument('--out', type=Path, required=True) |
| parser.add_argument('--map', choices=MAPS) |
| parser.add_argument('--clips', type=int, default=300) |
| parser.add_argument('--block', type=float, default=10) |
| parser.add_argument('--seed-base', type=int, default=100_000, help='job seeds are seed_base * (map index + 1) + clip') |
| parser.add_argument('--shard', type=int, default=0) |
| parser.add_argument('--shards', type=int, default=1) |
| parser.add_argument('--port', type=int, default=8131) |
| parser.add_argument('--looks', default=','.join(LOOKS), |
| help="pages to render, the graybox first (the sakura package adds 'sakura')") |
| parser.add_argument('--seconds', type=int, default=SECONDS, help='clip length (long-rollout clips: 120)') |
| parser.add_argument('--styled-frames', type=int, default=None, |
| help='record the styled looks only for the first N frames (long graybox clips need just frame 0)') |
| args = parser.parse_args() |
| SECONDS = args.seconds |
| STYLED_FRAMES = args.styled_frames or FPS * SECONDS |
| globals()['LOOKS'] = tuple(args.looks.split(',')) |
| if LOOKS[0] != 'original': |
| parser.error('--looks must start with original (the graybox page drives the others)') |
| asyncio.run(run(args)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
