Also use Enums for the node and edge attributes. I wanted to do this anyways and then thought I made a typo in "segmentation" so I just went for it to debug.
Checklist
Go through these things before merge. Actions should run automatically to test them, but for information on how to run locally, see CONTRIBUTING.md.
[x] I have added tests that prove that my feature works in various situations or tests the bugfix (if applicable).
[x] I have checked that the tests pass and I maintained or improved test coverage (if applicable).
[x] I have written docstrings and checked that they render correctly in the documentation build.
[x] I have checked that mypy type checking passes.
Proposed Change
Checklist
Go through these things before merge. Actions should run automatically to test them, but for information on how to run locally, see CONTRIBUTING.md.