feat: Implement adaptive superpixel generation for SuiT #2
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This commit introduces an adaptive superpixel generation mechanism into the
SuiT architecture.
Key changes:
1. **Parameter Predictor (`models/parameter_predictor.py`)**:
A lightweight CNN is implemented to predict superpixel parameters
(number of segments `K` and compactness `m`) dynamically based on the
input image.
2. **Superpixel Generation (`utils/superpixels.py`)**:
A new utility function `generate_superpixels` uses `skimage.segmentation.slic`
to generate superpixel maps based on the predicted `K` and `m`.
It supports batch processing and handles cases where scikit-image
might not be available.
3. **Adaptive SuiT Model (`models/suit_adaptive.py`)**:
The `SuperpixelVisionTransformer` has been refactored into `SuitAdaptive`.
This new model incorporates the `ParameterPredictor` and
`generate_superpixels` to dynamically create superpixels for each input
image. The predicted parameters `K` and `m` are clipped to predefined
ranges (K: 50-1000, m: 1-40). Internal methods have been updated
to use this dynamic superpixel generation. Model registration functions
have also been updated.
4. **Main Script (`main.py`)**:
The main training script has been updated to use the new `SuitAdaptive`
models. Command-line arguments for fixed superpixel parameters
(`--compactness`, `--n-spix-segments`) have been removed as these are
now predicted.
5. **Unit Tests (`tests/test_adaptive.py`)**:
Comprehensive unit tests have been added for the `ParameterPredictor`,
`generate_superpixels` function (including skimage availability checks
and fallback testing), and the `SuitAdaptive` model integration
(including forward pass and parameter clipping tests).
6. **Visualization Utilities (`visualize/adaptive_debug.py`)**:
Debugging utilities are provided to visualize superpixel overlays on
images and to compare segmentations generated with adaptive parameters
versus fixed parameters. These include handling for optional
dependencies like matplotlib and skimage.
This refactoring allows for more flexible and image-specific superpixel
generation, aiming to enhance semantic granularity. The training strategy
for the parameter predictor (heuristic loss or end-to-end) is planned
as a subsequent step.
This commit refactors the codebase to integrate adaptive superpixel
functionality directly into the existing `suit.py` and `utils.py` files,
removing the `models/` and `utils/` subdirectories that were
previously introduced. This addresses your feedback and resolves
`ModuleNotFoundError` issues related to package imports.
Key changes:
1. **File Structure Flattening**:
* The `ParameterPredictor` class definition has been moved directly
into `suit.py`.
* The `generate_superpixels` function and its `SKIMAGE_AVAILABLE`
helper constant have been moved directly into `utils.py`.
* The `SuitAdaptive` model definition (previously
`models/suit_adaptive.py`) now constitutes the main content of
`suit.py`, replacing the original `SuperpixelVisionTransformer` logic.
Registered model names (e.g., `suit_tiny_224_adaptive`) are
preserved.
* The `models/` and `utils/` directories and their specific files
(`parameter_predictor.py`, `suit_adaptive.py`, `superpixels.py`)
have been removed.
2. **Import Adjustments**:
* All import statements in `main.py`, `tests/test_adaptive.py`, and
`visualize/adaptive_debug.py` have been updated to reflect the
new locations of `ParameterPredictor` (in `suit.py`),
`generate_superpixels` (in `utils.py`), and the `SuitAdaptive`
model variants (in `suit.py`).
* Mock paths in tests have also been updated.
This resolves the previously reported `ModuleNotFoundError` and aligns
the codebase with the preferred flat directory structure.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This commit introduces an adaptive superpixel generation mechanism into the SuiT architecture.
Key changes:
Parameter Predictor (
models/parameter_predictor.py): A lightweight CNN is implemented to predict superpixel parameters (number of segmentsKand compactnessm) dynamically based on the input image.Superpixel Generation (
utils/superpixels.py): A new utility functiongenerate_superpixelsusesskimage.segmentation.slicto generate superpixel maps based on the predictedKandm. It supports batch processing and handles cases where scikit-image might not be available.Adaptive SuiT Model (
models/suit_adaptive.py): TheSuperpixelVisionTransformerhas been refactored intoSuitAdaptive. This new model incorporates theParameterPredictorandgenerate_superpixelsto dynamically create superpixels for each input image. The predicted parametersKandmare clipped to predefined ranges (K: 50-1000, m: 1-40). Internal methods have been updated to use this dynamic superpixel generation. Model registration functions have also been updated.Main Script (
main.py): The main training script has been updated to use the newSuitAdaptivemodels. Command-line arguments for fixed superpixel parameters (--compactness,--n-spix-segments) have been removed as these are now predicted.Unit Tests (
tests/test_adaptive.py): Comprehensive unit tests have been added for theParameterPredictor,generate_superpixelsfunction (including skimage availability checks and fallback testing), and theSuitAdaptivemodel integration (including forward pass and parameter clipping tests).Visualization Utilities (
visualize/adaptive_debug.py): Debugging utilities are provided to visualize superpixel overlays on images and to compare segmentations generated with adaptive parameters versus fixed parameters. These include handling for optional dependencies like matplotlib and skimage.This refactoring allows for more flexible and image-specific superpixel generation, aiming to enhance semantic granularity. The training strategy for the parameter predictor (heuristic loss or end-to-end) is planned as a subsequent step.