Coverage for collectra/tasks/image_segmentation/sam3.py: 89.74%

1"""SAM 3/3.1 image instance segmentation and mask-supervised training data.

3Uses Meta's sam3 package, installed separately from its official repository.

4Dataset preparation requires only Collectra's ordinary image dependencies.

5"""

7from __future__ import annotations

9import json

10import os

11import shutil

12import subprocess

13import sys

14import time

15import uuid

16from dataclasses import dataclass

17from datetime import datetime

18from pathlib import Path

19from typing import Annotated

21from cappa import Arg

22import numpy as np

23import yaml

24from PIL import Image as PILImage

26from collectra.cli import command

27from collectra.tasks.base import Task

28from collectra.types.images import Image, ImageSegmentation

29from collectra.utils import threading_locked

32def training_progress(message):

33 """Keep progress visible in terminals and redirected cluster job logs."""

34 print(f"[SAM3] {message}", flush=True)

37def _sam3():

38 try:

39 import sam3

40 except ModuleNotFoundError as error:

41 if error.name != "sam3":

42 raise

43 raise ImportError(

44 'Install Meta SAM3 separately: pip install "sam3[train] @ '

45 'git+https://github.com/facebookresearch/sam3.git". '

46 "See Meta prerequisites (Python 3.12+, PyTorch, CUDA) and checkpoint access."

47 ) from error

48 return sam3

51def _bounds(image: Image):

52 return tuple(round(value) for value in image.display_bounds())

55def mask_in_view(image: Image, mask: PILImage.Image):

56 """Apply the same crop, region and rotation as image.pil()."""

57 if mask.size != (image.raw_width, image.raw_height):

58 raise ValueError("Annotation mask dimensions differ from the source image")

59 if isinstance(image, ImageSegmentation):

60 pixels = (np.asarray(mask) != 0) & (np.asarray(image._mask_image()) != 0)

61 mask = PILImage.fromarray(pixels.astype("uint8"))

62 return mask.crop(_bounds(image)).rotate(image.orientation.to_degree(), expand=True)

65def mask_to_source(image: Image, mask: np.ndarray):

66 """Invert the displayed image's rotation/crop into source-sized 0/1 pixels."""

67 local = PILImage.fromarray(mask.astype("uint8"))

68 local = local.rotate(-image.orientation.to_degree(), expand=True)

69 left, top, right, bottom = _bounds(image)

70 if local.size != (right - left, bottom - top):

71 raise ValueError("Predicted mask dimensions differ from the input image view")

72 source = PILImage.new("L", (image.raw_width, image.raw_height))

73 source.paste(local, (left, top))

74 pixels = np.asarray(source) != 0

75 if isinstance(image, ImageSegmentation):

76 pixels &= np.asarray(image._mask_image()) != 0

77 return pixels.astype("uint8")

80def encode_rle(mask: np.ndarray):

81 """Lossless COCO uncompressed RLE, including holes and disconnected pixels."""

82 values = np.asarray(mask, dtype=np.uint8).ravel(order="F")

83 boundaries = np.concatenate(

84 ([0], np.flatnonzero(values[1:] != values[:-1]) + 1, [len(values)])

85 )

86 counts = np.diff(boundaries).tolist()

87 if len(values) and values[0]:

88 counts.insert(0, 0)

89 return {"size": list(mask.shape), "counts": counts}

92@dataclass

93class SegmentationSample:

94 image: Image

95 instances: list[ImageSegmentation]

98def prepare_segmentation_inputs(

99 processed_inputs, processed_parents, input_maps, parent_input_maps

100):

101 samples = []

102 for filename, instances in input_maps.items():

103 parents = {

104 image.id: image

105 for image in parent_input_maps.get(filename, [])

106 if isinstance(image, Image)

107 }

108 grouped = {}

109 for instance in instances:

110 if not isinstance(instance, ImageSegmentation):

111 raise TypeError(

112 "SAM training targets must be ImageSegmentation instances"

113 )

114 ids = (

115 [instance.parents]

116 if isinstance(instance.parents, str)

117 else instance.parents

118 )

119 candidates = [parents[key] for key in ids if key in parents]

120 if len(candidates) != 1:

121 raise ValueError(

122 f"Mask {instance.id!r} must identify exactly one input image parent"

123 )

124 parent = candidates[0]

125 if parent.get_path().resolve() != instance.get_path().resolve():

126 raise ValueError(

127 f"Mask {instance.id!r} refers to a different source image"

128 )

129 grouped.setdefault(

130 parent.id, SegmentationSample(parent, [])

131 ).instances.append(instance)

132 samples.extend(grouped.values())

133 return samples

134

135

136def _sample_partition(sample):

137 partitions = {item.partition for item in sample.instances if item.partition}

138 if sample.image.partition:

139 partitions.add(sample.image.partition)

140 if len(partitions) > 1:

141 raise ValueError("Instances of one image cannot belong to different partitions")

142 return next(iter(partitions), "")

143

144

145def export_dataset(

146 samples, directory: Path, *, prompts=None, validation="", exclude="", max_items=None

147):

148 if max_items is not None and max_items < 1:

149 raise ValueError("max_items must be positive")

150 started = time.monotonic()

151 directory = Path(directory).resolve()

152 selected = []

153 for sample in samples:

154 partition = _sample_partition(sample)

155 if exclude and partition == exclude:

156 continue

157 selected.append(

158 (sample, "valid" if validation and partition == validation else "train")

159 )

160 if max_items is not None and len(selected) >= max_items:

161 break

162 labels = sorted({item.name for sample, _ in selected for item in sample.instances})

163 if not labels or not any(split == "train" for _, split in selected):

164 raise ValueError("No labelled training instances were found")

165 categories = [

166 {"id": index + 1, "name": (prompts or {}).get(label, label.replace("_", " "))}

167 for index, label in enumerate(labels)

168 ]

169 label_ids = {label: index + 1 for index, label in enumerate(labels)}

170 training_progress(

171 f"Exporting {len(selected)} images and "

172 f"{sum(len(sample.instances) for sample, _ in selected)} masks to {directory} "

173 f"({sum(split == 'train' for _, split in selected)} train, "

174 f"{sum(split == 'valid' for _, split in selected)} validation images)"

175 )

176 datasets = {

177 split: {"images": [], "annotations": [], "categories": categories}

178 for split in ("train", "valid")

179 }

180 for split in datasets:

181 (directory / split).mkdir(parents=True, exist_ok=True)

182 for image_id, (sample, split) in enumerate(selected, 1):

183 training_progress(

184 f"Export {image_id}/{len(selected)} [{split}]: {sample.image.id} "

185 f"({len(sample.instances)} masks; {time.monotonic() - started:.1f}s elapsed)"

186 )

187 image = sample.image.pil().convert("RGB")

188 filename = f"{image_id:08d}.png"

189 image.save(directory / split / filename)

190 data = datasets[split]

191 data["images"].append(

192 {

193 "id": image_id,

194 "file_name": filename,

195 "width": image.width,

196 "height": image.height,

197 }

198 )

199 for mask_index, instance in enumerate(sample.instances, 1):

200 mask = np.asarray(mask_in_view(sample.image, instance._mask_image())) != 0

201 ys, xs = np.nonzero(mask)

202 if not len(xs):

203 raise ValueError(

204 f"Mask {instance.id!r} has no pixels inside its input image"

205 )

206 if mask.shape != (image.height, image.width):

207 raise ValueError(

208 "Image subclass crop differs from its declared display bounds"

209 )

210 data["annotations"].append(

211 {

212 "id": len(data["annotations"]) + 1,

213 "image_id": image_id,

214 "category_id": label_ids[instance.name],

215 "iscrowd": 0,

216 "bbox": [

217 int(xs.min()),

218 int(ys.min()),

219 int(xs.max() - xs.min() + 1),

220 int(ys.max() - ys.min() + 1),

221 ],

222 "area": int(mask.sum()),

223 "segmentation": encode_rle(mask),

224 }

225 )

226 if mask_index % 25 == 0:

227 training_progress(

228 f" Image {image_id}/{len(selected)}: encoded mask "

229 f"{mask_index}/{len(sample.instances)}"

230 )

231 for split, data in datasets.items():

232 training_progress(

233 f"Writing {split} annotations: {len(data['images'])} images, "

234 f"{len(data['annotations'])} masks"

235 )

236 (directory / split / "_annotations.coco.json").write_text(

237 json.dumps(data), encoding="utf-8"

238 )

239 training_progress(f"Dataset export complete ({time.monotonic() - started:.1f}s): {directory}")

240 return directory

241

242

243class ImageSegmentationSam3(Task):

244 """Find all instances matching text prompts and return one mask per instance."""

245

246 def __init__(

247 self,

248 name: str,

249 model: str | Path = "",

250 prompt: str = "",

251 prompts: dict[str, str] | None = None,

252 version: str = "sam3.1",

253 threshold: float = 0.5,

254 device: str = "cuda",

255 **kwargs,

256 ):

257 super().__init__(name, **kwargs)

258 if version not in {"sam3", "sam3.1"}:

259 raise ValueError("version must be sam3 or sam3.1")

260 if not 0 <= threshold <= 1:

261 raise ValueError("threshold must be between 0 and 1")

262 self.model, self.prompt, self.prompts = model, prompt, prompts or {}

263 self.version, self.threshold, self.device = version, threshold, device

264 self._processor = None

265

266 def attributes_to_ignore(self):

267 return super().attributes_to_ignore() | {"_processor"}

268

269 def input_type(self) -> type:

270 return Image

271

272 def output_type(self) -> type:

273 return ImageSegmentation

274

275 def _prompts(self):

276 if self.prompts:

277 mapping = self.prompts

278 else:

279 outputs = getattr(self, "output", None) or [self.get_output_name()]

280 outputs = outputs if isinstance(outputs, list) else [outputs]

281 mapping = {

282 label: self.prompt or label.replace("_", " ") for label in outputs

283 }

284 if any(not label or not prompt.strip() for label, prompt in mapping.items()):

285 raise ValueError("SAM prompts and output labels must be nonempty")

286 return mapping

287

288 def _checkpoint(self):

289 _sam3()

290 if self.model:

291 path = Path(self.model).expanduser()

292 if not path.is_absolute() and self.pipeline is not None:

293 path = Path(self.pipeline.path) / path

294 path = path.resolve()

295 if not path.is_file():

296 raise FileNotFoundError(path)

297 return path

298 from sam3.model_builder import download_ckpt_from_hf

299

300 return Path(download_ckpt_from_hf(version=self.version))

301

302 def _init_model(self):

303 if self._processor is not None:

304 return

305 _sam3()

306 import torch

307 from sam3.model_builder import build_sam3_image_model

308 from sam3.model.sam3_image_processor import Sam3Processor

309

310 checkpoint = self._checkpoint()

311 state = torch.load(checkpoint, map_location="cpu", weights_only=True)

312 state = state.get("model", state)

313 from .sam3_training import load_image_state

314

315 model = build_sam3_image_model(

316 checkpoint_path=None,

317 load_from_HF=False,

318 device=self.device,

319 enable_segmentation=True,

320 )

321 load_image_state(model, state)

322 model.to(self.device).eval()

323 self._processor = Sam3Processor(

324 model, device=self.device, confidence_threshold=self.threshold

325 )

326

327 @threading_locked()

328 def run(self, *images: Image) -> list[ImageSegmentation]:

329 if not images or any(not isinstance(image, Image) for image in images):

330 raise TypeError(

331 "SAM requires one or more Image artefacts (including subclasses)"

332 )

333 prompts = self._prompts()

334 self._init_model()

335 directory = Path(self.context.file_path or Path.cwd()) / "masks"

336 directory.mkdir(parents=True, exist_ok=True)

337 results = []

338 for image in images:

339 state = self._processor.set_image(image.pil().convert("RGB"))

340 for label, prompt in prompts.items():

341 self._processor.reset_all_prompts(state)

342 prediction = self._processor.set_text_prompt(prompt=prompt, state=state)

343 masks = prediction["masks"]

344 if hasattr(masks, "detach"):

345 masks = masks.detach().cpu().numpy()

346 masks = np.asarray(masks)

347 if masks.ndim == 4 and masks.shape[1] == 1:

348 masks = masks[:, 0]

349 if masks.ndim != 3:

350 raise ValueError("SAM returned masks with an unexpected shape")

351 for mask in masks:

352 pixels = mask_to_source(image, mask > 0.5)

353 if not pixels.any():

354 continue

355 path = directory / f"sam3-{uuid.uuid4().hex}.png"

356 PILImage.fromarray(pixels).save(path)

357 results.append(

358 ImageSegmentation(

359 name=label,

360 data=image.get_path(),

361 mask=path,

362 parents=[image.id],

363 orientation=image.orientation,

364 partition=image.partition,

365 )

366 )

367 return results

368

369 def _samples(self, inputs, max_items=None, exclude=""):

370 from collectra.tasks.machine_learning.training import (

371 _training_files,

372 _prepare_data,

373 )

374

375 if self.pipeline is None:

376 raise RuntimeError(

377 "Attach the task to a pipeline before preparing training data"

378 )

379 started = time.monotonic()

380 training_progress("Discovering training result files (CPU/I/O preparation)")

381 self.pipeline._ensure_workflow_connected()

382 node = self.pipeline.node_manager.resolve_node(self.name)

383 files = []

384 for value in inputs:

385 path = Path(value).expanduser().resolve()

386 training_progress(f"Scanning {path}")

387 if path.name == "results.yaml" and path.is_file():

388 path = path.parent

389 if path.is_dir() and (path / "results.yaml").is_file():

390 files.append(path)

391 else:

392 files.extend(_training_files(self, [str(path)]))

393 samples = []

394 files = list(dict.fromkeys(files))

395 training_progress(f"Found {len(files)} result files; loading images and validating masks")

396 for index, path in enumerate(files, 1):

397 training_progress(

398 f"Load {index}/{len(files)}: {path} "

399 f"({time.monotonic() - started:.1f}s elapsed)"

400 )

401 # Read each results directory independently. Different datasets can

402 # have the same page basename and even the same artefact ids.

403 targets, target_maps = _prepare_data(

404 self, [path], self.pipeline.node_manager.get_children_artefact(node)

405 )

406 parents, parent_maps = _prepare_data(

407 self, [path], self.pipeline.node_manager.get_ancestor_artefacts(node)

408 )

409 for sample in prepare_segmentation_inputs(targets, parents, target_maps, parent_maps):

410 if max_items is not None and exclude and _sample_partition(sample) == exclude:

411 continue

412 samples.append(sample)

413 if max_items is not None and len(samples) >= max_items:

414 break

415 if max_items is not None and len(samples) >= max_items:

416 training_progress(f"Reached max_items={max_items}; stopping source loading")

417 break

418 training_progress(

419 f"Loaded {len(samples)} labelled images and "

420 f"{sum(len(sample.instances) for sample in samples)} masks "

421 f"({time.monotonic() - started:.1f}s)"

422 )

423 return samples

424

425 def prepare_training_data(

426 self, inputs: list[str], output: Path, validation: str = "", exclude: str = "",

427 max_items: int | None = None,

428 ) -> Path:

429 """Export annotated masks and input images as SAM-compatible COCO datasets."""

430 return export_dataset(

431 self._samples(inputs, max_items=max_items, exclude=exclude),

432 output,

433 prompts=self._prompts(),

434 validation=validation,

435 exclude=exclude,

436 max_items=max_items,

437 )

438

439 @command

440 def train(

441 self,

442 inputs: Annotated[

443 list[str],

444 Arg(

445 required=False,

446 default=[],

447 help="Results YAML files, result directories, or parent directories containing annotated images and masks. May be omitted when reusing an existing export.",

448 ),

449 ],

450 output: Annotated[

451 Path | None,

452 Arg(help="Directory for prepared data, metrics.csv, logs and best model weights. Defaults to a timestamped directory; required with --reuse-data."),

453 ] = None,

454 validation: Annotated[

455 str,

456 Arg(help="Partition used for validation and best-model selection by mask IoU. Omit to train without validation. Ignored when reusing an existing export."),

457 ] = "",

458 exclude: Annotated[

459 str,

460 Arg(help="Partition excluded from both training and validation during data preparation. Ignored when reusing an existing export."),

461 ] = "",

462 epochs: Annotated[

463 int, Arg(help="Number of training epochs; validation runs after every epoch when available."),

464 ] = 20,

465 batch: Annotated[

466 int, Arg(help="Number of training images per batch on each GPU."),

467 ] = 1,

468 num_gpus: Annotated[

469 int, Arg(help="Number of GPUs used for training on this node."),

470 ] = 1,

471 keep_log: Annotated[

472 bool,

473 Arg(help="Keep prepared data, metrics, logs and checkpoints after copying the best weights into the pipeline. --no-keep-log removes the output directory."),

474 ] = True,

475 reuse_data: Annotated[

476 bool,

477 Arg(help="Reuse OUTPUT/dataset when both split annotation files exist; otherwise prepare it from the input paths. Requires --output. Reuse ignores input paths, partitions and --max-items. Does not resume optimizer state."),

478 ] = False,

479 wandb: Annotated[

480 bool,

481 Arg(help="Log epoch metrics to Weights & Biases. Project: <pipeline-directory>-<task>; run: output directory name; entity: W&B settings or WANDB_ENTITY."),

482 ] = False,

483 max_items: Annotated[

484 int | None,

485 Arg(help="Maximum total images across training and validation, after exclusions, in input order. Omit for all images. Ignored when reusing an existing export."),

486 ] = None,

487 ) -> Path:

488 """Prepare instance masks and fine-tune using Meta's image trainer.

489

490 --reuse-data uses OUTPUT/dataset directly when present, otherwise

491 prepares it from inputs. It requires --output.

492 """

493 if min(epochs, batch, num_gpus) < 1:

494 raise ValueError("epochs, batch and num_gpus must be positive")

495 if reuse_data and output is None:

496 raise ValueError("--reuse-data requires --output")

497 output = (

498 Path(output).expanduser().resolve()

499 if output is not None

500 else Path.cwd() / datetime.now().strftime("%Y%m%d_%H%M%S")

501 )

502 reuse_prepared = reuse_data and all(

503 (output / "dataset" / split / "_annotations.coco.json").is_file()

504 for split in ("train", "valid")

505 )

506 if not reuse_prepared and max_items is not None and max_items < 1:

507 raise ValueError("max_items must be positive")

508 if not reuse_prepared and not inputs:

509 raise ValueError(

510 "Supply training inputs to prepare the dataset; --reuse-data without "

511 f"inputs requires an existing export at {output / 'dataset'}"

512 )

513 output.mkdir(parents=True, exist_ok=True)

514 started = time.monotonic()

515 training_progress(

516 f"Training {self.name}: output={output}, epochs={epochs}, batch={batch}, GPUs={num_gpus}"

517 )

518 if reuse_prepared:

519 dataset = output / "dataset"

520 training_progress(f"Reusing prepared dataset directly: {dataset}")

521 training_progress("Skipping source checks and export; using existing train/validation splits")

522 if max_items is not None:

523 training_progress("Ignoring --max-items because --reuse-data uses the existing dataset unchanged")

524 else:

525 if reuse_data:

526 training_progress("Prepared dataset is missing or incomplete; falling back to preparation from inputs")

527 training_progress("Preparing the dataset before GPU training starts")

528 dataset = self.prepare_training_data(

529 inputs, output / "dataset", validation, exclude, max_items=max_items

530 )

531 from .sam3_training import training_config

532

533 training_progress("Resolving the SAM checkpoint (a download may be needed)")

534 initial_checkpoint = self._checkpoint()

535 training_progress(f"Using checkpoint: {initial_checkpoint}")

536 training_progress("Building the SAM training configuration")

537 config = training_config(

538 dataset,

539 output,

540 initial_checkpoint,

541 epochs=epochs,

542 batch=batch,

543 num_gpus=num_gpus,

544 wandb=wandb,

545 run_name=output.name,

546 project=f"{self.pipeline.path.name}-{self.name}",

547 confidence_threshold=self.threshold,

548 )

549 config_path = output / "sam3-training.yaml"

550 config_path.write_text(

551 yaml.safe_dump(config, sort_keys=False), encoding="utf-8"

552 )

553 training_progress(f"Configuration saved: {config_path}")

554 training_progress(

555 f"Starting SAM trainer ({time.monotonic() - started:.1f}s preparation); "

556 "trainer output follows"

557 )

558 subprocess.run(

559 [

560 sys.executable,

561 "-u",

562 "-c",

563 "from collectra.tasks.image_segmentation.sam3_training import launch; import sys; launch(sys.argv[1])",

564 str(config_path),

565 ],

566 check=True,

567 env={**os.environ, "PYTHONUNBUFFERED": "1"},

568 )

569 training_progress("SAM trainer exited successfully; checking its checkpoint")

570 checkpoint = output / "training" / "checkpoints" / "best.pt"

571 if not checkpoint.is_file():

572 raise FileNotFoundError(

573 f"SAM training produced no checkpoint at {checkpoint}"

574 )

575 destination = self.pipeline.path / f"{self.name}-{uuid.uuid4().hex[:8]}.pt"

576 shutil.copy2(checkpoint, destination)

577 self.model = destination

578 self._processor = None

579 self.pipeline.data[self.name]["model"] = destination.name

580 self.pipeline.save(self.name)

581 training_progress(f"Saved trained checkpoint: {destination}")

582 if not keep_log:

583 shutil.rmtree(output, ignore_errors=True)

584 training_progress(f"Removed training directory: {output}")

585 training_progress(f"Training complete ({time.monotonic() - started:.1f}s)")

586 return destination