🆕 Integrate Foundation Models Available VIA timm: UNI, Prov-GigaPath, H-optimus-0

[GeorgeBatch]

Making a pull request as discussed in issue #855

---

Copied from the issue:

I think it would be useful to integrate pre-trained foundation models from other labs into tiatoolbox.models.architecture.vanilla.py.

Currently, the _get_architecture() function allows the use of models from torchvision.models.

But another function _get_timm_architecture() could be made to incorporate foundation models which are available from timm with weights on HuggingFace Hub. All the models from timm that I've used require users to sign the licence agreement with the authors, so the licencing question seems to be solved itself since there is no way users will get access to the model weights just through Tiatoolbox without getting the access request approved by the authors first.

[GeorgeBatch]

Only added UNI and Prov-GigaPath for now. Will add more after initial comments.

I do not like that TimmBackbone pretty much repeats the CNNBackbone. The only difference is the absence of global average pooling in TimmBackbone.

[GeorgeBatch]

I found this file for testing: https://github.com/TissueImageAnalytics/tiatoolbox/blob/develop/tests/models/test_arch_vanilla.py

There might be a problem regarding memory and compute resources when running some of the larger feature extractors through GitHub actions, e.g. Prov-GigaPath needs a considerable amount of memory just to be loaded.

[shaneahmed]

I found this file for testing: https://github.com/TissueImageAnalytics/tiatoolbox/blob/develop/tests/models/test_arch_vanilla.py

There might be a problem regarding memory and compute resources when running some of the larger feature extractors through GitHub actions, e.g. Prov-GigaPath needs a considerable amount of memory just to be loaded.

As this is just testing the functionality and not loading weights, I hope this would work.

[shaneahmed]

Only added UNI and Prov-GigaPath for now. Will add more after initial comments.

I do not like that TimmBackbone pretty much repeats the CNNBackbone. The only difference is the absence of global average pooling in TimmBackbone.

In that case, you can inherit CNNBackbone and just define the functions which change like we are doing here

tiatoolbox/tiatoolbox/models/engine/patch_predictor.py

Line 25 in e4e6f22

class PatchPredictor(EngineABC):

for PatchPredictor. The PatchPredictor uses all the functionalities of EngineABC other than the ones defined explicitly.

[shaneahmed]

I have updated this branch to make sure that tests pass on Ubuntu-24 before we merge it with develop.

[GeorgeBatch]

I think we should add a notebook to show how to

1. Extract and save metadata from your slides
2. Make and save thumbnails and masks for the WSI-s
3. Extract features using the masks saved before using either CNNBackbone or TimmBackbone

This is precisely what I wanted to know how to do, and there was no end-to-end example notebook.

---

Alternatively, we can add code for saving a mask into the Masking Notebook, which currently does not show how to save the masks:
https://github.com/TissueImageAnalytics/tiatoolbox/blob/develop/examples/03-tissue-masking.ipynb

wsi = WSIReader.open(slide_path)
mask = wsi.tissue_mask(resolution=1.25, units="power")
mask_thumbnail = mask.slide_thumbnail(resolution=1.25, units="power",)
mask_thumbnail_path = os.path.join(f"{slide_name}_mask.png")
imwrite(mask_thumbnail_path, np.uint8(mask_thumbnail * 255))

And add model = TimmBackbone("UNI") as an alternative to model = CNNBackbone("resnet50") in Slide Graph pipeline since it has the feature extraction code. Note, in Slide Graph notebooks, masks are assumed to be already saved.
https://github.com/TissueImageAnalytics/tiatoolbox/blob/develop/examples/inference-pipelines/slide-graph.ipynb

[GeorgeBatch]

There are 3 more families of models that can be integrated, but they require their own files to be created: Hibou -b and -L, Phikon v1 and v2, Virchow v1 and v2

Should I add those files and call them from _get_timm_network()?

Hibou requires to trust remote code when creating the model, which I do not really like.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[GeorgeBatch]

@shaneahmed, I think this branch is ready to be merged into develop.

I added code for saving a mask into the Masking Notebook, which currently does not show how to save the masks:
https://github.com/TissueImageAnalytics/tiatoolbox/blob/develop/examples/03-tissue-masking.ipynb

I also added a comment showing that TimmBackbone could be used as an alternative:

model = CNNBackbone("resnet50")  # TimmBackbone(backbone="UNI", pretrained=True)

in both slide graph notebooks:
https://github.com/TissueImageAnalytics/tiatoolbox/blob/develop/examples/inference-pipelines/slide-graph.ipynb
https://github.com/TissueImageAnalytics/tiatoolbox/blob/develop/examples/full-pipelines/slide-graph.ipynb

[shaneahmed]

Thanks @GeorgeBatch This would be very helpful.