File size: 11,853 Bytes
00c7586 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | import asyncio
from fractions import Fraction
import hashlib
import json
from pathlib import Path
import tempfile
import time
from .bootstrap import production as a
from .config import MODELS, PROFILE_POOLS
from .media import frames_for, check_synchronized, InputError
from .models import Session, generate, runtime_event
from .preprocessing import temporal_inputs, human_transport, instruction_suffix, three_view_caption_inputs
class Predictor:
def __init__(self, videos, source, pool, bound, model, proc, instruction, progress):
self.videos, self.source, self.pool, self.bound = (videos, source, pool, bound)
self.model, self.proc, self.instruction, self.progress = (model, proc, instruction, progress)
self.calls = []
def predict(self, request):
from flora_runtime.video_segmentation.fine_subtask import FINE_RESPONSE_SCHEMA
if request['source']['source_id'] != self.source.source_id or request['model_binding'] != self.bound:
raise RuntimeError('Production request binding changed')
predictions = []
for payload in request['windows']:
window = a.window_from_payload(payload, self.source)
plan = a.temporal_frame_plan(window, list(self.videos['main'].pts_us))
mosaic = 'Dex' if self.pool in ('dex', 'tianji') else 'YAM'
images = frames_for(self.videos, set(plan['source_frame_indices']), mosaic=mosaic)
if self.pool == 'human':
images = human_transport(images)
frames = [images[i] for i in plan['source_frame_indices']]
inputs, prompt = temporal_inputs(frames, window, self.pool, self.proc, list(self.videos), self.instruction)
self.progress(0.25, desc=f'Finding actions 路 window {len(self.calls) + 1}')
text = generate(self.model, self.proc, inputs, 1024)
predictions.append(dict(scope_id=window.scope_id, prediction_json=text))
self.calls.append(dict(window=window.worker_payload(), frame_plan=plan, prompt_sha256=hashlib.sha256(prompt.encode()).hexdigest(), output=text))
return dict(schema_version=FINE_RESPONSE_SCHEMA, source_id=self.source.source_id, model_io_contract=self.bound['model_io_contract'], model_binding=self.bound, predictions=predictions)
def segment(videos, pool, session, instruction, progress):
model, proc = session.load('tempseg')
main = videos['main']
source = a.AnnotationSource(source_id='upload_' + main.sha256[:24], location=str(main.path), duration_s=main.duration_s, fps=float(main.rate))
bound = dict(candidate_id='flora_final', checkpoint_sha256=MODELS['tempseg']['sha256'], model_io_contract=a.fine_window_profile(a.PROFILES[pool]).model_io_contract)
predictor = Predictor(videos, source, pool, bound, model, proc, instruction, progress)
if len(videos) == 3:
fine = a.FineSubtaskSegmenter(predictor=predictor, model_binding=bound, profile_id='robot_atomic_intent_v1', window_profile_id=a.PROFILES[pool], fail_on_unresolved=True, batch_concurrency=1, unresolved_stitch_tolerance_ticks=2, quality_policy=a.FineSegmentationQualityPolicy(drop_guard_only_segments=True))
engine = a.HierarchicalSubtaskSegmenter(coarse_segmenter=a.NoCoarse(), fine_segmenter=fine, fine_prior_mode='none', fine_failure_policy='adaptive_fine_then_conservative_completion', fine_output_coverage_policy='allow_gaps', adaptive_fine_retry_max_windows=32)
else:
engine = a.production_segmenter(predictor, bound, pool)
segments = asyncio.run(engine.segment_async(source))
return ([dict(start_s=s.span.start_s, end_s=s.span.end_s, metadata=dict(s.metadata)) for s in segments], predictor.calls)
def single_plans(videos, segments, profile):
if profile in ('Human', 'Human actions'):
from pts_input_adapter import native_blocks, microseconds
from flora_runtime.subtask_captioning import u2_contract
main = videos['main']
pts = [float(t / 1000000) for t in main.pts_us]
plans, dispositions = ([], [])
for i, s in enumerate(segments):
record = dict(event_id=f'event_{i:04d}', event_index=i, start_s=s['start_s'], end_s=s['end_s'])
try:
blocks = native_blocks(pts, microseconds(s['start_s']), microseconds(s['end_s']), main.duration_us, main.rate, u2_contract)
except ValueError:
dispositions.append(dict(record, status='unsupported_target_frames'))
continue
plans.append(dict(record, blocks=blocks, source_fps_num=main.rate.numerator, source_fps_den=main.rate.denominator))
dispositions.append(dict(record, status='caption_input'))
return (plans, dispositions)
main = videos['main']
if not main.cfr:
raise InputError('Use the Human actions profile for videos with variable frame timestamps.')
return a.caption_plans(segments, dict(fps_num=main.rate.numerator, fps_den=main.rate.denominator, frame_count=len(main.pts_us)))
def caption_single(videos, segments, profile, session, instruction, progress):
from flora_unified.unified_input import processor_inputs, prompt as native_prompt
from transformers.video_utils import VideoMetadata
plans, dispositions = single_plans(videos, segments, profile)
if not plans:
return ([dict(d, caption='', status=d['status']) for d in dispositions], [])
model, proc = session.load('captioner')
rows, audit = ({}, [])
with tempfile.TemporaryDirectory(prefix='flora-frames-') as temporary:
root = Path(temporary)
wanted = {i for plan in plans for block in plan['blocks'] for i in block['frame_indices']}
images = frames_for(videos, wanted)
receipts = {}
for i, image in images.items():
path = root / f'{i:012d}.png'
image.save(path, format='PNG')
receipts[i] = dict(path=path.name, png_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), size=list(image.size))
for n, plan in enumerate(plans):
progress(0.55 + 0.4 * n / len(plans), desc=f'Writing captions 路 {n + 1}/{len(plans)}')
sample = a.caption_sample(plan, receipts, ['main'])
if instruction.strip():
from flora_unified.unified_input import load_frames
frames, _ = load_frames(sample, root)
text = proc.apply_chat_template([dict(role='user', content=native_prompt(sample, proc.video_token) + instruction_suffix(instruction))], tokenize=False, add_generation_prompt=True, enable_thinking=False)
metadata = [VideoMetadata(total_num_frames=len(f), fps=float(videos['main'].rate), frames_indices=list(b['block_local_timestamp_indices'])) for f, b in zip(frames, sample['blocks'], strict=True)]
inputs = proc(text=[text], videos=frames, video_metadata=metadata, do_sample_frames=False, return_metadata=True, return_tensors='pt')
else:
inputs, _, text = processor_inputs(sample, root, proc)
caption = generate(model, proc, inputs, 512)
rows[plan['event_id']] = dict(start_s=plan['start_s'], end_s=plan['end_s'], caption=caption, status='complete')
audit.append(dict(event_id=plan['event_id'], blocks=plan['blocks'], prompt_sha256=hashlib.sha256(text.encode()).hexdigest()))
return ([rows.get(d['event_id'], dict(start_s=d['start_s'], end_s=d['end_s'], caption='', status=d['status'])) for d in dispositions], audit)
def caption_three(videos, segments, session, instruction, progress, profile='YAM'):
from plans import plans_for_episode
from robot_caption_geometry import normalized, source_frames
main = videos['main']
from math import lcm
scale = lcm(*(t.denominator for t in main.pts_us)) * 1000000
ticks = [int(t * scale / 1000000) for t in main.pts_us]
clock = dict(ticks=ticks, time_base=str(Fraction(1, scale)), fps=str(main.rate), duration_s=main.duration_s, first_pts=0)
group = 'Dex' if profile in ('Dex', 'Tianji') else 'YAM'
episode = dict(episode_id='upload', source_id=main.sha256, group=group, units=[])
prepared, rows, audit = ([], [], [])
for i, s in enumerate(segments):
unit = dict(target_id=f'event_{i:04d}', unit_id=f'prediction:{i}', native_scope_s=[s['start_s'], s['end_s']])
try:
plan = plans_for_episode(dict(episode, units=[unit]), {role: clock for role in videos})[0]
prepared.append((i, plan))
except AssertionError:
rows.append((i, dict(start_s=s['start_s'], end_s=s['end_s'], caption='', status='unsupported_target_frames')))
if prepared:
model, proc = session.load('captioner')
for n, (i, plan) in enumerate(prepared):
progress(0.55 + 0.4 * n / len(prepared), desc=f'Writing captions 路 {n + 1}/{len(prepared)}')
wanted = {j for block in plan['blocks'] for j in block['indices']}
images = frames_for(videos, wanted, mosaic=group)
geometry = normalized(plan, images)
inputs, text = three_view_caption_inputs(plan, source_frames(geometry), proc, instruction)
caption = generate(model, proc, inputs, 512)
rows.append((i, dict(start_s=plan['native_scope_s'][0], end_s=plan['native_scope_s'][1], caption=caption, status='complete')))
audit.append(dict(event_id=plan['target_id'], blocks=plan['blocks'], common_geometry=geometry['common_geometry'], camera_layout=group, prompt_sha256=hashlib.sha256(text.encode()).hexdigest()))
return ([row for _, row in sorted(rows)], audit)
def annotate(videos, profile, long_windows, instruction, progress):
check_synchronized(videos)
instruction_suffix(instruction)
if profile not in PROFILE_POOLS:
raise InputError('Choose a supported action profile.')
if len(videos) == 3 and profile in ('Human', 'Human actions'):
raise InputError('Choose YAM, Dex or Tianji when uploading wrist cameras.')
if long_windows and len(videos) != 3:
raise InputError('The 32-second profile requires three synchronized camera views.')
if long_windows and profile not in ('YAM', 'Robot actions'):
raise InputError('The 32-second profile is available for YAM videos.')
pool = 'yam32' if long_windows else PROFILE_POOLS[profile]
start = time.monotonic()
session = Session()
try:
progress(0.1, desc='Loading the action segmenter')
segments, temporal_audit = segment(videos, pool, session, instruction, progress)
runtime_event('segmentation_complete', windows=len(temporal_audit), segments=len(segments))
progress(0.5, desc='Stitching finished 路 loading the captioner')
if len(videos) == 3:
captions, caption_audit = caption_three(videos, segments, session, instruction, progress, profile=profile)
else:
captions, caption_audit = caption_single(videos, segments, profile, session, instruction, progress)
for row, segment_row in zip(captions, segments, strict=True):
row['segmentation_provenance'] = segment_row['metadata'].get('segmentation_provenance', 'fine')
runtime_event('annotation_complete', segments=len(captions), complete_captions=sum((r['status'] == 'complete' for r in captions)), elapsed_s=round(time.monotonic() - start, 3))
return dict(schema_version='flora_demo_annotation_v1', duration_s=videos['main'].duration_s, segments=captions, elapsed_s=round(time.monotonic() - start, 2), source_clock='seconds relative to the first native source presentation timestamp', camera_roles=list(videos), profile=profile, model_revisions=MODELS, input_sha256={role: video.sha256 for role, video in videos.items()}, pipeline=dict(temporal_windows=temporal_audit, caption_inputs=caption_audit))
finally:
session.close()
|