scripts/phase1_segment.py

import json
import os

import cv2
import numpy as np
import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor

import transformers.models.sam3.modeling_sam3 as sam3_mod
_orig = sam3_mod.Sam3Model.get_text_features
def _patched(self, input_ids, attention_mask=None, **kwargs):
    result = _orig(self, input_ids, attention_mask=attention_mask, **kwargs)
    if hasattr(result, 'pooler_output') and result.pooler_output is not None:
        return result.pooler_output
    return result
sam3_mod.Sam3Model.get_text_features = _patched

ROOT = "/root/bdmc_pipeline"
FRAMES_DIR = f"{ROOT}/frames"
MASKS_DIR = f"{ROOT}/masks/road_surface"
MODEL_ID = "jetjodh/sam3"
N_FRAMES = 300
H, W = 1920, 1080

os.makedirs(f"{ROOT}/outputs", exist_ok=True)
os.makedirs(MASKS_DIR, exist_ok=True)

device = "cuda"
print("Loading SAM 3.1...")
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).to(device).eval()

def stream_masks_for_prompt(text_prompt, n_frames=300):
    first_im = cv2.imread(f"{FRAMES_DIR}/f_0001.jpg")
    first_pil = Image.fromarray(cv2.cvtColor(first_im, cv2.COLOR_BGR2RGB))

    session = processor.init_video_session(
        video=[first_pil], inference_device=device, inference_state_device=device)
    processor.add_text_prompt(session, text_prompt)

    all_masks = []
    with torch.no_grad():
        for i in range(n_frames):
            im = cv2.imread(f"{FRAMES_DIR}/f_{i+1:04d}.jpg")
            pil = Image.fromarray(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
            pv = processor(images=pil, return_tensors="pt").pixel_values.to(device)

            if i == 0:
                outputs = model(inference_session=session, frame_idx=0, frame=pv)
            else:
                outputs = model(inference_session=session, frame=pv)

            obj_to_mask = outputs.get("obj_id_to_mask", {})
            if obj_to_mask:
                best_oid = max(obj_to_mask.keys(),
                               key=lambda k: outputs.get("obj_id_to_score", {}).get(k, 0))
                sig = torch.sigmoid(obj_to_mask[best_oid])
                binary = (sig > 0.5).float()
                if binary.dim() == 3:
                    binary = binary[0]
                m = binary.cpu().numpy().astype(np.uint8)
                if m.shape != (H, W):
                    m = cv2.resize(m, (W, H), interpolation=cv2.INTER_NEAREST)
                all_masks.append(m * 255)
            else:
                all_masks.append(np.zeros((H, W), dtype=np.uint8))

            if (i + 1) % 50 == 0:
                area = (all_masks[-1] > 127).sum() / (H * W)
                print(f"  Frame {i+1}/{n_frames}: area={area:.3f}")

    return all_masks

def compute_temporal_iou(masks):
    ious = []
    for i in range(1, len(masks)):
        a, b = masks[i-1] > 127, masks[i] > 127
        inter, union = (a & b).sum(), (a | b).sum()
        if union > 0:
            ious.append(float(inter) / float(union))
    return float(np.mean(ious)) if ious else 0.0

def compute_area_stats(masks):
    fracs = [(m > 127).sum() / (H * W) for m in masks]
    return {
        "mean_area_frac": round(float(np.mean(fracs)), 4),
        "median_area_frac": round(float(np.median(fracs)), 4),
        "min_area_frac": round(float(np.min(fracs)), 4),
    }

print("\n=== Prompt comparison ===")
prompts = ["road surface", "asphalt pavement", "road"]
prompt_results = {}
best_prompt, best_iou = None, -1

for p in prompts:
    print(f"\nPrompt: '{p}'")
    masks = stream_masks_for_prompt(p)
    iou = compute_temporal_iou(masks)
    stats = compute_area_stats(masks)
    stats["mean_temporal_iou"] = round(iou, 4)
    prompt_results[p] = stats
    print(f"  IoU={iou:.4f}, area={stats['mean_area_frac']:.4f}, min_area={stats['min_area_frac']:.4f}")
    if iou > best_iou:
        best_iou, best_prompt = iou, p

prompt_results["_selected"] = best_prompt
print(f"\nSelected: '{best_prompt}' (IoU={best_iou:.4f})")

print(f"\n=== Final segmentation: '{best_prompt}' ===")
final_masks = stream_masks_for_prompt(best_prompt)

for i, mask in enumerate(final_masks):
    cv2.imwrite(f"{MASKS_DIR}/f_{i+1:04d}.png", mask)

area_fracs = [(m > 127).sum() / (H * W) for m in final_masks]
summary = {
    "method": "SAM 3.1 (jetjodh/sam3) video tracking / Object Multiplex",
    "prompt": best_prompt,
    "prompts_tested": prompts,
    "prompt_comparison": prompt_results,
    "n_frames": N_FRAMES,
    "image_size": [H, W],
    "mean_area_frac": round(float(np.mean(area_fracs)), 4),
    "min_area_frac": round(float(np.min(area_fracs)), 4),
}
with open(f"{ROOT}/outputs/phase1_summary.json", "w") as f:
    json.dump(summary, f, indent=1)
with open(f"{ROOT}/outputs/prompt_comparison.json", "w") as f:
    json.dump(prompt_results, f, indent=1)

print(f"\nPhase 1 complete. {N_FRAMES} masks saved.")
print(json.dumps(summary, indent=1))