Add UI for Interactive Attention Exploration

[DEschweiler]

As part of the TransformLiver project, a user interface for interactively exploring attention (class token and tile-to-tile) was added. A hook was installed to extract attention from the ViT, the command "STAMP attention_ui" was added to the CLI functionality, and the default config was extended accordingly.

[FWao]

I would personally like to make this feature optional as it introduces a bunch of new dependencies for the UI part.

My proposal would be to add a optional dependency group to the pyproject.toml which can then be installed via
pip install stamp[attentionui] or similar. So the default installation would be without all of the UI dependencies. Especially so as pyqt5 seems to be no longer developed / supported.
When using this feature, try to import the dependencies and inform the user if they are not installed.

[DEschweiler]

Thank you, good point! Made the UI dependencies optional, added the error message on failed dependency imports and unified the naming to "attentionui".

[georg-wolflein]

Hi @DEschweiler, thank you so much for your contribution, excellent work!

I slightly modified your PR to fix some typing issues that mypy showed and I changed last_attn_weights to not be an instance attribute but rather a variable within the forward() function (as this is the only place it was used). If you are happy with these changes (44386d9), your PR is ready to be merged!

By the way, could you provide a screenshot of the UI? This could be a nice addition to the main README.md ;)

[DEschweiler]

Hi @georg-wolflein, thanks a lot for your modifications and approval! Attached you can find a screenshot of the napari-based UI. Do you want me to also upload it to the repository and add it to the readme?

[EzicStar]

Hello @DEschweiler,
First of all, sorry for not reviewing your PR yet!

By the end of this week I will review it and if ready, merge it into main. Unfortunately, I have to merge one quick fix for current gradcams and create a new stable version, 2.2.3, which should be done by Wednesday. Once ready, could you bring those changes to this branch as you have been doing with the latest updates?

I am looking forward to try it as the screenshots from the UI seem really cool! Thank you very much for contributing.

[DEschweiler]

Hello @EzicStar,
no worries, all good. Once you are finished with creating the new stable version, I'll try to bring the updates to this branch as soon as possible.
Thank you, I'm looking forward to the integration and seeing the UI in use.

Cheers

[EzicStar]

Hello @DEschweiler,
sorry again for taking so long to review this. I've just tried it and it looks great! I think that users will love it. I made some comments on the code for further improvements, but in general it looks really good. The next steps would be to decide whether to merge it directly into main branch or keep it in a separate one (still in STAMP repo and documented).

These are the difficulties for merging it directly into main:

• STAMP is used frequently on headless Linux systems that don't support user interfaces, and its development is focused on this. A solution to this approach would be for example to have this UI web with a client-server model, as Jupyter does. However, this would take a lot of rework as it is already built with pyqt5.
• The ViT model code becomes more complex and therefore more difficult to maintain, specially now that we have plans of adding regression and survival prediction.

Merging this on a separate branch which is referenced by the documentation on main's branch would still give visibility to this feature and make it easier to maintain, not requiring to takle this difficulties. What do you think @FWao, @georg-wolflein and @s1787956?

P.S: please bring changes from v2.3 and add documentation on the getting-started.md file so it is easier for users to use it.

Thank you for contributing :)

Best,
Juan Pablo

[EzicStar]

Attention UI logic should be in a separate script, as it happens with crossval and train.

[EzicStar]

The params are identical to the ones provided for gradcam heatmaps (current ones) so this is not needed. To decide whether to generate gradcam heatmaps or attention maps, a better option could be to add an optional field in heatmaps config to let the user decide which approach to use. It could be like this:

heatmaps:
  # Whether to use GRADCAM based heatmaps or (introduce description of this implementation here, also please inform that you require and extra dependency).
  attention_ui: False

[EzicStar]

As mentioned in the other comment, this new section could be integrated into heatmaps one adding one more optional field

[EzicStar]

I find that adding 2 new forward functions to self attention increases code complexity, making it more difficult to maintain. The same goes for the return weights logics in MultiHeadALiBi. Is there a way to make it simpler and more decoupled?