AI_spine/sam_segment.py

#!/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()