Skip to content

Introduce Classification Factory and Simplify Model Imports #4456

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

Conversation

kprokofi
Copy link
Contributor

@kprokofi kprokofi commented Jul 10, 2025

Summary

  • Hide internal model classes like EfficientNetMulticlassCls, EfficientNetHLabelCls, MobileNetV3MultilabelCls, etc., in favor of user-friendly factory usage. Using @overload helps provide docstrings, parameter hints, and autocomplete support in IDEs without actually hiding the arguments in **kwargs.
  • Simplify OTX API:
from otx.models import EfficientNet, ATSS

cls_model = EfficientNet(task="multi_class", ...)
od_model = ATSS(..)

from otx.data import OTXDataModule, OTXSegmentationDataset
datamodule = OTXDataModule()
dataset = OTXSegmentationDataset()
  • Move some utils from one place to another to avoid circular imports

How to test

Checklist

  • I have added unit tests to cover my changes.​
  • I have added integration tests to cover my changes.​
  • I have ran e2e tests and there is no issues.
  • I have added the description of my changes into CHANGELOG in my target branch (e.g., CHANGELOG in develop).​
  • I have updated the documentation in my target branch accordingly (e.g., documentation in develop).
  • I have linked related issues.

License

  • I submit my code changes under the same Apache License that covers the project.
    Feel free to contact the maintainers if that's a concern.
  • I have updated the license header for each file (see an example below).
# Copyright (C) 2025 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

@kprokofi kprokofi added the API Any changes in OTX API label Jul 10, 2025
@kprokofi kprokofi added this to the 2.5.0 milestone Jul 10, 2025
@github-actions github-actions bot added the TEST Any changes in tests label Jul 11, 2025
Copy link
Contributor

@eugene123tw eugene123tw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Kirill for the great work on enhancing the classification model factories!

I added some comments related to

  • Two EfficientNet classes (factory vs backbone) created confusion
  • Using Literal types for model names
    Something like:
MobileNetV3: Literal["mobilenetv3_small", "mobilenetv3_large"] (2 variants)
VisionTransformer: Literal["vit-tiny", "vit-small", "vit-base", "vit-large", "dinov2-small", "dinov2-base", "dinov2-large", "dinov2-giant"] (8 variants)
...
  • Improve docstring to expose 50+ TorchVision models are hidden from users: I think either list_models and write a example docstring or just explain in the docstring where could user find those available models

Lastly, could you please add usage examples in the docstrings to demonstrate the improvements? I think this would help users understand how to use the enhanced factory APIs:

Something like:

class MobileNetV3:
    """Factory class for MobileNetV3 models.
    
    Args:
        model_name (Literal["mobilenetv3_small", "mobilenetv3_large"]): Model variant.
    
    Examples:
        >>> # Basic usage with type safety
        >>> model = MobileNetV3(
        ...     task="multi_class",
        ...     model_name="mobilenetv3_small",  # IDE shows both options
        ...     label_info=10
        ... )
        
        >>> # Multi-label classification  
        >>> model = MobileNetV3(
        ...     task="multi_label", 
        ...     model_name="mobilenetv3_large",
        ...     label_info=[1, 5, 10]  # Multi-label setup
        ... )
"""

But also for TVModel, VisionTransformer, etc.

@kprokofi kprokofi requested a review from eugene123tw July 22, 2025 12:19
@kprokofi
Copy link
Contributor Author

@eugene123tw, thank you for your review! I applied your suggestions.

Copy link
Contributor

@eugene123tw eugene123tw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great! Thank you for your efforts! :)

@sovrasov sovrasov merged commit 21508cd into open-edge-platform:release/2.5 Jul 22, 2025
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
API Any changes in OTX API TEST Any changes in tests
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants