🆕 Simplified SAM segmentation in TIAViz

[measty]

This PR takes the sam-viz PR and fixes and substantially simplifies it.

I think the original PR was going for a more general prompt segmentor engine that would in theory cope with prompts dispersed over a whole WSI, but I think that approach has issues in practice (and it would also need some significant work to fix from its current state, I think). This PR focusses on segmenting prompts made within a specific view window, which I think is more practical and allows for a much simplified approach.

To use, draw a few points (and/or a box) on some object(s) in the field of view in TIAViz using this widget:

Then select SAM in the model select box and click run. The segmetations should be displayed after a short time.

Examples:

[Copilot]

Pull request overview

This PR adds simplified SAM (Segment Anything Model) integration to TIAViz, enabling interactive segmentation using point and bounding box prompts within a specific view window. The implementation includes a new SAM architecture wrapper, a PromptSegmentor engine, and UI integration for interactive use.

Changes:

• Adds SAM model architecture and PromptSegmentor engine for interactive segmentation with point/box prompts
• Integrates SAM into TIAViz UI with callback for running segmentation on selected regions
• Extends utility functions to support offset parameter for annotation coordinate transformations

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
tiatoolbox/models/architecture/sam.py	New SAM model architecture using transformers library for zero-shot segmentation
tiatoolbox/models/engine/prompt_segmentor.py	New PromptSegmentor engine that handles SAM inference and output conversion
tiatoolbox/visualization/bokeh_app/main.py	Adds sam_segment callback and SAM model option to UI
tiatoolbox/utils/misc.py	Extends process_contours and dict_to_store_semantic_segmentor to support offset parameter
tiatoolbox/visualization/tileserver.py	Filters out temporary bokeh databases from commit operations
tiatoolbox/models/init.py	Exports SAM and PromptSegmentor classes
requirements/requirements.txt	Adds transformers>=4.51.1 dependency
.pre-commit-config.yaml	Sets default Python version to 3.12
tests/models/test_prompt_segmentor.py	Unit tests for PromptSegmentor functionality
tests/models/test_arch_sam.py	Unit tests for SAM architecture
tests/test_app_bokeh.py	Integration test for SAM in TIAViz UI

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[shaneahmed]

Thanks @measty

This looks good. A few minor comments if you could address before we merge.

• I was not able to test it though on my workstation using --base-path argument.
• Also it would be good to show any progress bar or message on the visualization rather than only in the terminal if possible.
• Can we reset the model weights after the first run if possible?

[measty]

Thanks @measty

This looks good. A few minor comments if you could address before we merge.

• I was not able to test it though on my workstation using --base-path argument.
• Also it would be good to show any progress bar or message on the visualization rather than only in the terminal if possible.
• Can we reset the model weights after the first run if possible?

I've addressed most of the coments and added the clear button to get rid of existing box/dot annotations. Regarding testing issues using --base-path, can you reproduce that?

A progress bar might be useful, but its a 'nice to have ' that maybe can be skipped for now in the interests of time. It is usually not taking that long to run, so shouldn't be a huge problem.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 14 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The test only tests points and boxes separately, but doesn't test the case where both point_coords and box_coords are provided together. Given the potential issue with how the forward method handles both types of prompts simultaneously (concatenating results), a test case should be added to verify the behavior when both are provided.

Suggested change

	assert len(mask_output) > 0, "Output should have at least one element"
	assert len(score_output) > 0, "Output should have at least one element"
	assert mask_output is not None, "Output should not be None"
	assert len(mask_output) > 0, "Output should have at least one element"
	assert len(score_output) > 0, "Output should have at least one element"
	# test inference with both points and boxes
	mask_output, score_output = model.infer_batch(
	model, patch, points, box_coords=boxes, device=select_device(on_gpu=ON_GPU)
	)
	assert mask_output is not None, "Output should not be None"
	assert len(mask_output) > 0, "Output should have at least one element"
	assert len(score_output) > 0, "Output should have at least one element"

[Copilot]

The test doesn't test with point prompts as the docstring suggests "Test running SAM on points and a box." The test only sets up box_source data but never sets pt_source data. Consider adding a test case that includes point coordinates, or update the docstring to accurately reflect what is being tested.

[Copilot]

The type hint for the model parameter should be more specific. In the init signature, it's typed as torch.nn.Module, but the docstring and the code logic assume it should be a SAM model or None. Consider changing the type hint to SAM | None to be more accurate, or if it should accept any model, update the docstring accordingly.

Suggested change

	model (SAM):
	Model architecture to use. If None, defaults to SAM.
	"""
	def __init__(
	self,
	model: torch.nn.Module = None,
	model (SAM | None):
	Model architecture to use. If None, defaults to SAM.
	"""
	def __init__(
	self,
	model: SAM | None = None,

[Copilot]

The model.to(device) call is made, but then infer_batch also calls model.eval().to(device) again. This is redundant. Consider removing the model.to(device) call on line 73 since infer_batch already handles device placement.

Suggested change

self.model.to(device)

[Copilot]

The sam_segment function doesn't validate that at least one of point_coords or box_coords is provided before calling prompt_segmentor.run(). If both UI["pt_source"].data["x"] and UI["box_source"].data["x"] are empty, both point_coords and box_coords will be None, which will cause a ValueError in SAM.infer_batch. Consider adding validation to check that at least one source has data before proceeding.

[Copilot]

The to method overrides the parent ModelABC.to() method and adds device tracking for self.device and self.model (the internal SamModel). However, the return type annotation shows it returns ModelABC or DataParallel[ModelABC], but this method doesn't call the parent's logic for handling DataParallel wrapping. The super().to() call returns the parent, but the method then returns self without checking if DataParallel should be applied. This could cause inconsistency with the base class behavior when multiple GPUs are available.

Suggested change

	super().to(device, dtype=dtype, non_blocking=non_blocking)
	self.device = device
	self.model.to(device)
	return self
	result = super().to(device, dtype=dtype, non_blocking=non_blocking)
	# Ensure SAM-specific device tracking is applied to the underlying module,
	# even if it has been wrapped in DataParallel by the base class.
	module = result.module if isinstance(result, torch.nn.DataParallel) else result
	if isinstance(module, SAM):
	module.device = device
	module.model.to(device=device, dtype=dtype, non_blocking=non_blocking)
	return result

[Copilot]

The coordinate transformation for box coordinates appears inconsistent. On line 1538, the box bottom-right corner is calculated as [x[i] + width[i], height[i] + y[i]], but standard bounding box format is [x1, y1, x2, y2] where x2 = x1 + width and y2 = y1 + height. The current code has the y-coordinate addition order swapped (height[i] + y[i] instead of y[i] + height[i]), though they're mathematically equivalent. For clarity and consistency with the x-coordinate, consider using y[i] + height[i].

Suggested change

	[x[i], y[i], x[i] + width[i], height[i] + y[i]]
	[x[i], y[i], x[i] + width[i], y[i] + height[i]]

[Copilot]

The temporary directory created here is never cleaned up, unlike in the segment_on_box function which calls rmtree on temporary directories. This will cause temporary files to accumulate over time. Consider adding cleanup code similar to segment_on_box, such as calling rmtree(tmp_save_dir) after the annotations are loaded.

Suggested change

	tmp_save_dir = Path(tempfile.mkdtemp(suffix="bokeh_temp"))
	tmp_dir = tempfile.TemporaryDirectory(suffix="bokeh_temp")
	tmp_save_dir = Path(tmp_dir.name)