mirror of
https://github.com/storytold/lightcraft.git
synced 2026-10-14 19:27:01 -04:00
Guinther KovalskiandClaude Opus 5.5
6468bfc06b
AI masks from Meta's SAM 3, run in LightCraft itself: a new
`lightcraft-segment` crate ports the Hugging Face transformers
implementation (Apache-2.0) to candle, on the GPU through Metal on macOS
(CPU elsewhere), with no C/C++ and no `unsafe`.
- Object: click to include, ⌥-click to leave a part out (SAM 3's
interactive prompt encoder + two-way mask decoder).
- Describe: a text prompt selects every instance ("sky", "the red car";
"car, road" merges several) via the CLIP text encoder and DETR detector.
- Both as new masks and as Add / Subtract / Intersect components; + / −
buttons under the selected mask combine selections.
- The segmentation is stored with the mask (288² quantized logits, ~10-30
KB), so renders, exports and the web build never need the model.
- A zoomed-in detail pass (background, ~5 s on an M4 Pro) re-segments the
area around a selection from a sharper source: 5-10x finer edges on
small objects.
- Per-selection Edge slider (hard ↔ soft); hovering a mask in the list
shows it in red (the always-on overlay is now off by default, O toggles).
- The engine feature `sam` (on in the desktop app) gates the model; the
image is encoded once per photo in the background, clicks take ~35 ms.
- Weights are not bundled (SAM License): tools/install-sam3.sh / .ps1
download and verify facebook/sam3; docs/ai-masks.md.
- Clearer errors (missing file, missing model) and a stderr logger in the
app (LIGHTCRAFT_LOG=info).
candle is pinned at 0.9.2: later releases make candle-core depend on
tokenizers with the Oniguruma C library. Every stage is checked against
the Python reference (tests/reference.rs, tools/sam3_reference.py): on
Metal the backbone matches to 7e-5 relative error, masks to ~5e-5.
Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01C8BkS7JDqa1EMcjW6spDGA
72 lines
3.7 KiB
Python
72 lines
3.7 KiB
Python
"""Reference tensors for the Rust SAM 3 port (crates/segment/tests/reference.rs).
|
|
|
|
Runs the Hugging Face transformers implementation (fp32, CPU) on one photo and writes the
|
|
shared pixel input plus intermediate and final outputs of the click (tracker) and text
|
|
(detector) paths. Usage (needs torch, torchvision, transformers, pillow, safetensors):
|
|
LIGHTCRAFT_SAM3_DIR=<model dir> python tools/sam3_reference.py photo.jpg ref.safetensors
|
|
The text prompts below suit a photo of the moon in a blue sky; change them for other photos.
|
|
"""
|
|
import os, sys, time
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from safetensors.torch import save_file
|
|
from transformers import Sam3TrackerModel, Sam3Model, CLIPTokenizer
|
|
|
|
MODEL = os.environ.get("LIGHTCRAFT_SAM3_DIR", os.path.expanduser("~/Library/Application Support/LightCraft/models/sam3"))
|
|
img_path, out_path = sys.argv[1], sys.argv[2]
|
|
torch.set_grad_enabled(False)
|
|
|
|
img = Image.open(img_path).convert("RGB")
|
|
W, H = img.size
|
|
# Same preprocessing as the Rust side: bilinear resize to 1008², /255, (x-0.5)/0.5.
|
|
r = img.resize((1008, 1008), Image.BILINEAR)
|
|
px = torch.from_numpy(np.asarray(r, dtype=np.float32) / 255.0).permute(2, 0, 1)[None]
|
|
px = (px - 0.5) / 0.5
|
|
out = {"pixel_values": px.contiguous()}
|
|
|
|
t = time.time()
|
|
trk = Sam3TrackerModel.from_pretrained(MODEL).eval()
|
|
vis = trk.vision_encoder(px)
|
|
out["backbone"] = vis.last_hidden_state.contiguous()
|
|
for i, f in enumerate(vis.fpn_hidden_states):
|
|
out[f"trk_fpn{i}"] = f.contiguous()
|
|
# points in 1008 space: one positive, one negative
|
|
pts = torch.tensor([[[[0.5 * 1008, 0.55 * 1008], [0.2 * 1008, 0.2 * 1008]]]], dtype=torch.float32)
|
|
lbl = torch.tensor([[[1, 0]]], dtype=torch.int32)
|
|
o = trk(pixel_values=px, input_points=pts, input_labels=lbl, multimask_output=False)
|
|
out["trk_points"] = pts.contiguous()
|
|
out["trk_masks"] = o.pred_masks.contiguous()
|
|
out["trk_iou"] = o.iou_scores.contiguous()
|
|
o3 = trk(pixel_values=px, input_points=pts[:, :, :1], input_labels=lbl[:, :, :1], multimask_output=True)
|
|
out["trk_multi_masks"] = o3.pred_masks.contiguous()
|
|
out["trk_multi_iou"] = o3.iou_scores.contiguous()
|
|
print("tracker", time.time() - t, o.pred_masks.shape, o.iou_scores, file=sys.stderr)
|
|
del trk
|
|
|
|
t = time.time()
|
|
det = Sam3Model.from_pretrained(MODEL).eval()
|
|
tok = CLIPTokenizer.from_pretrained(MODEL)
|
|
for i, text in enumerate(["moon", "sky", "bird"]):
|
|
enc = tok(text, padding="max_length", max_length=32, return_tensors="pt")
|
|
out[f"det{i}_ids"] = enc.input_ids.to(torch.int64).contiguous()
|
|
out[f"det{i}_mask"] = enc.attention_mask.to(torch.int64).contiguous()
|
|
tf = det.get_text_features(input_ids=enc.input_ids, attention_mask=enc.attention_mask)
|
|
out[f"det{i}_text"] = tf.pooler_output.contiguous()
|
|
o = det(pixel_values=px, input_ids=enc.input_ids, attention_mask=enc.attention_mask)
|
|
out[f"det{i}_logits"] = o.pred_logits.contiguous()
|
|
out[f"det{i}_presence"] = o.presence_logits.contiguous()
|
|
out[f"det{i}_boxes"] = o.pred_boxes.contiguous()
|
|
out[f"det{i}_masks"] = o.pred_masks.contiguous()
|
|
sc = (o.pred_logits.sigmoid() * o.presence_logits.sigmoid())[0]
|
|
print(text, "presence", o.presence_logits.sigmoid().item(), "top", sc.topk(3).values.tolist(), file=sys.stderr)
|
|
for i, f in enumerate(det.vision_encoder(px).fpn_hidden_states):
|
|
out[f"det_fpn{i}"] = f.contiguous()
|
|
print("detector", time.time() - t, file=sys.stderr)
|
|
# tokenizer cases for the Rust BPE
|
|
cases = ["person", "Red Car", "a dog's toy", "sky, clouds!", "café 2 people", "tree-line"]
|
|
for i, c in enumerate(cases):
|
|
out[f"tok{i}"] = tok(c, padding="max_length", max_length=32, return_tensors="pt").input_ids.to(torch.int64).contiguous()
|
|
save_file(out, out_path, metadata={"image": img_path, "W": str(W), "H": str(H), "cases": "|".join(cases)})
|
|
print("saved", out_path, file=sys.stderr)
|