#!/usr/bin/env python3 """One-shot GPU worker: semantic ROI prompts -> real SAM masks -> previews.""" import argparse from contextlib import nullcontext import json from pathlib import Path import sys import time import traceback import numpy as np from PIL import Image from sam_runtime import check_environment, load_detector from segmentation_outputs import assess_candidate, render_outputs, write_report def choose_candidate(masks, scores, alpha, box, points): """Rank real model candidates; ROI and seeds are hints, not ground truth.""" ranked = [] for index, (mask, score) in enumerate(zip(masks, scores)): qa = assess_candidate(mask, alpha, box, points) reasons = [] if qa['area_px'] < 4: reasons.append('EMPTY_OR_TINY_FOREGROUND') if qa['inside_roi_fraction'] < .1: reasons.append('OUTSIDE_TARGET_ROI') rank = float(score) * (.25 + .75 * qa['inside_roi_fraction']) * (1 - qa['outside_alpha_ratio']) rank *= .8 + .2 * qa['seed_hit_fraction'] ranked.append({'index': index, 'model_score': float(score), 'rank_score': rank, 'qa': qa, 'rejected': reasons}) eligible = [c for c in ranked if not c['rejected']] return max(eligible, key=lambda c: c['rank_score']) if eligible else None, ranked def validate_prompts(plan, prompts, image): w, h = image.size if (prompts.get('width'), prompts.get('height')) != (w, h): raise ValueError('Prompt canvas does not match source image') parts = {p['id']: p for p in plan['parts']} entries = prompts['parts'] if len(entries) != len(parts) or {p['id'] for p in entries} != set(parts): raise ValueError('Prompt IDs must match the semantic plan exactly') alpha = np.array(image)[:, :, 3] for entry in entries: part = parts[entry['id']] if entry['status'] not in ['candidate', 'hidden', 'blocked']: raise ValueError('Unknown prompt status') if entry['status'] != 'candidate': continue if part['visibility'] == 'hidden' or entry.get('seed_source') != 'alpha_roi_distance_transform': raise ValueError('Only visible, program-generated candidate prompts are supported') box = entry['bbox_px_xyxy'] if len(box) != 4 or any(type(v) is not int for v in box) or not (0 <= box[0] < box[2] <= w and 0 <= box[1] < box[3] <= h): raise ValueError('Invalid pixel ROI') if not entry['positive_points_px']: raise ValueError('Candidate prompt has no foreground seeds') for point in entry['positive_points_px']: if len(point) != 2 or any(type(v) is not int for v in point): raise ValueError('Invalid pixel seed') x, y = point if not (box[0] <= x < box[2] and box[1] <= y < box[3]) or alpha[y, x] == 0: raise ValueError('Candidate seed is outside foreground/ROI') return parts def run(args, runtime): import torch from sam3.model.sam3_image_processor import Sam3Processor out = Path(args.output) plan = json.loads((out / 'semantic-parts.json').read_text(encoding='utf-8')) prompts = json.loads((out / 'segmentation-prompts.json').read_text(encoding='utf-8')) image = Image.open(out / 'source.png').convert('RGBA') parts_by_id = validate_prompts(plan, prompts, image) alpha = np.array(image)[:, :, 3] print('Loading local SAM detector weights...', flush=True) model, loading = load_detector(args.checkpoint, args.precision) processor = Sam3Processor(model, device='cuda', confidence_threshold=args.confidence) rgb = Image.alpha_composite(Image.new('RGBA', image.size, 'white'), image).convert('RGB') candidates_dir = out / 'candidates' candidates_dir.mkdir(exist_ok=True) selected_parts, skipped, diagnostics = [], [], [] started = time.monotonic() context = torch.autocast('cuda', dtype=torch.bfloat16) if args.precision == 'bf16' else nullcontext() with torch.inference_mode(), context: print('Encoding reference image on GPU...', flush=True) state = processor.set_image(rgb) for number, entry in enumerate(prompts['parts'], 1): name = entry['id'] if entry['status'] != 'candidate': skipped.append({'id': name, 'reason': entry['status']}) continue part = parts_by_id[name] print(f'[{number}/{len(prompts["parts"])}] {name}: text / text + coarse ROI', flush=True) processor.reset_all_prompts(state) # The official image concept API accepts text and boxes, not points. # Alpha seeds are used only to rank the returned mask candidates. processor.set_text_prompt(part['semantic'], state) x0, y0, x1, y1 = entry['bbox_px_xyxy'] w, h = image.size # Preserve text-only candidates: an inaccurate coarse box can # steer the concept detector toward an unrelated fragment. text_masks = state['masks'].detach().cpu().numpy().reshape(-1, h, w).copy() text_scores = state['scores'].detach().float().cpu().numpy().reshape(-1).copy() processor.add_geometric_prompt([(x0+x1)/(2*w), (y0+y1)/(2*h), (x1-x0)/w, (y1-y0)/h], True, state) box_masks = state['masks'].detach().cpu().numpy().reshape(-1, h, w) box_scores = state['scores'].detach().float().cpu().numpy().reshape(-1) masks = np.concatenate([text_masks, box_masks]) scores = np.concatenate([text_scores, box_scores]) best, ranked = choose_candidate(masks, scores, alpha, entry['bbox_px_xyxy'], entry['positive_points_px']) for candidate in ranked: candidate['prompt_mode'] = 'text_only' if candidate['index'] < len(text_masks) else 'text_and_box' diagnostic = {'id': name, 'text_prompt': part['semantic'], 'candidates': ranked, 'selected_index': best['index'] if best else None} diagnostics.append(diagnostic) if best is None: skipped.append({'id': name, 'reason': 'NO_ACCEPTABLE_SAM_MASK', 'candidate_count': len(ranked)}) continue mask = masks[best['index']] Image.fromarray(mask.astype(np.uint8)*255).save(candidates_dir / f'{name}-raw.png') warnings = ['SEMANTIC_REVIEW_REQUIRED'] if best['qa']['outside_alpha_ratio'] > .1: warnings.append('MASK_OUTSIDE_CHARACTER') if best['qa']['inside_roi_fraction'] < .5: warnings.append('MASK_EXTENDS_BEYOND_COARSE_ROI') if best['qa']['seed_hit_fraction'] < 1: warnings.append('CANDIDATE_SEED_NOT_COVERED') if best['qa']['components_at_least_4px'] > 10: warnings.append('FRAGMENTED_MASK') selected_parts.append({'id': name, 'semantic': part['semantic'], 'rig_role': part['rig_role'], 'mask': mask, 'model_score': best['model_score'], 'qa': best['qa'], 'prompt_mode': best['prompt_mode'], 'warnings': warnings, 'semantic_verified': False}) if not selected_parts: (out / 'candidate-diagnostics.json').write_text(json.dumps(diagnostics, indent=2), encoding='utf-8') raise RuntimeError('SAM produced no acceptable masks; see candidate-diagnostics.json') result = render_outputs(image, selected_parts, out) result.update(schema_version=1, quality_status='needs_review', skipped=skipped, runtime=runtime, model_loading=loading, elapsed_seconds=time.monotonic()-started, peak_cuda_memory_bytes=torch.cuda.max_memory_allocated(), prompt_strategy='rank text_only and text_and_box candidates; alpha seeds rank masks only') (out / 'segmentation.json').write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding='utf-8') (out / 'candidate-diagnostics.json').write_text(json.dumps(diagnostics, ensure_ascii=False, indent=2), encoding='utf-8') write_report(out, result) print(f'SAM finished: {len(selected_parts)} masks, {len(skipped)} skipped. GPU worker exiting.', flush=True) def main(): parser = argparse.ArgumentParser() parser.add_argument('--code-dir', required=True) parser.add_argument('--checkpoint', required=True) parser.add_argument('--output') parser.add_argument('--probe', action='store_true') parser.add_argument('--confidence', type=float, default=.2) parser.add_argument('--precision', choices=['bf16', 'fp32'], default='bf16') args = parser.parse_args() if not Path(args.checkpoint).is_file(): parser.error('Local checkpoint does not exist') if not 0 < args.confidence < 1: parser.error('--confidence must be between 0 and 1') try: runtime = check_environment(args.code_dir) if args.probe: print(json.dumps(runtime)) return if not args.output: parser.error('--output is required') run(args, runtime) except Exception as error: if args.output and Path(args.output).is_dir(): (Path(args.output) / 'worker-error.json').write_text(json.dumps({'type': type(error).__name__, 'message': str(error), 'traceback': traceback.format_exc()}, indent=2), encoding='utf-8') traceback.print_exc() raise SystemExit(1) if __name__ == '__main__': main()