Push unit tests and some bug fixes to develop (#3224)

* Migrate release branch to develop

* Fix file name

* Merge latest commit

* Add #3240

README.md

@@ -22,7 +22,6 @@
 [![Codecov](https://codecov.io/gh/openvinotoolkit/training_extensions/branch/develop/graph/badge.svg?token=9HVFNMPFGD)](https://codecov.io/gh/openvinotoolkit/training_extensions)
 [![OpenSSF Scorecard](https://api.securityscorecards.dev/projects/github.com/openvinotoolkit/training_extensions/badge)](https://securityscorecards.dev/viewer/?uri=github.com/openvinotoolkit/training_extensions)
 [![Pre-Merge Test](https://github.com/openvinotoolkit/training_extensions/actions/workflows/pre_merge.yaml/badge.svg)](https://github.com/openvinotoolkit/training_extensions/actions/workflows/pre_merge.yaml)
-[![Nightly Test](https://github.com/openvinotoolkit/training_extensions/actions/workflows/daily.yaml/badge.svg)](https://github.com/openvinotoolkit/training_extensions/actions/workflows/daily.yaml)
 [![Build Docs](https://github.com/openvinotoolkit/training_extensions/actions/workflows/docs.yaml/badge.svg)](https://github.com/openvinotoolkit/training_extensions/actions/workflows/docs.yaml)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 [![Downloads](https://static.pepy.tech/personalized-badge/otx?period=total&units=international_system&left_color=grey&right_color=green&left_text=PyPI%20Downloads)](https://pepy.tech/project/otx)
@@ -34,12 +33,12 @@
 ## Introduction
 
 OpenVINO™ Training Extensions is a low-code transfer learning framework for Computer Vision.
-The CLI commands of the framework allows users to train, infer, optimize and deploy models easily and quickly even with low expertise in the deep learning field.
+The API & CLI commands of the framework allows users to train, infer, optimize and deploy models easily and quickly even with low expertise in the deep learning field.
 OpenVINO™ Training Extensions offers diverse combinations of model architectures, learning methods, and task types based on [PyTorch](https://pytorch.org) and [OpenVINO™ toolkit](https://software.intel.com/en-us/openvino-toolkit).
 
-OpenVINO™ Training Extensions provides a "model template" for every supported task type, which consolidates necessary information to build a model.
+OpenVINO™ Training Extensions provides a "recipe" for every supported task type, which consolidates necessary information to build a model.
 Model templates are validated on various datasets and serve one-stop shop for obtaining the best models in general.
-If you are an experienced user, you can configure your own model based on [torchvision](https://pytorch.org/vision/stable/index.html), [mmcv](https://github.com/open-mmlab/mmcv), [timm](https://github.com/huggingface/pytorch-image-models) and [OpenVINO Model Zoo (OMZ)](https://github.com/openvinotoolkit/open_model_zoo).
+If you are an experienced user, you can configure your own model based on [torchvision](https://pytorch.org/vision/stable/index.html), [mmcv](https://github.com/open-mmlab/mmcv) and [OpenVINO Model Zoo (OMZ)](https://github.com/openvinotoolkit/open_model_zoo).
 
 Furthermore, OpenVINO™ Training Extensions provides automatic configuration for ease of use.
 The framework will analyze your dataset and identify the most suitable model and figure out the best input size setting and other hyper-parameters.
@@ -55,59 +54,145 @@ OpenVINO™ Training Extensions supports the following computer vision tasks:
 - **Instance segmentation** including tiling algorithm support
 - **Action recognition** including action classification and detection
 - **Anomaly recognition** tasks including anomaly classification, detection and segmentation
+- **Visual Prompting** tasks including segment anything model, zero-shot visual prompting
 
 OpenVINO™ Training Extensions supports the [following learning methods](https://openvinotoolkit.github.io/training_extensions/latest/guide/explanation/algorithms/index.html):
 
-- **Supervised**, incremental training, which includes class incremental scenario and contrastive learning for classification and semantic segmentation tasks
-- **Semi-supervised learning**
-- **Self-supervised learning**
+- **Supervised**, incremental training, which includes class incremental scenario.
 
 OpenVINO™ Training Extensions provides the following usability features:
 
-- [Auto-configuration](https://openvinotoolkit.github.io/training_extensions/latest/guide/explanation/additional_features/auto_configuration.html). OpenVINO™ Training Extensions analyzes provided dataset and selects the proper task and model with appropriate input size to provide the best accuracy/speed trade-off. It will also make a random auto-split of your dataset if there is no validation set provided.
+- [Auto-configuration](https://openvinotoolkit.github.io/training_extensions/latest/guide/explanation/additional_features/auto_configuration.html). OpenVINO™ Training Extensions analyzes provided dataset and selects the proper task and model to provide the best accuracy/speed trade-off.
 - [Datumaro](https://openvinotoolkit.github.io/datumaro/stable/index.html) data frontend: OpenVINO™ Training Extensions supports the most common academic field dataset formats for each task. We are constantly working to extend supported formats to give more freedom of datasets format choice.
 - **Distributed training** to accelerate the training process when you have multiple GPUs
 - **Mixed-precision training** to save GPUs memory and use larger batch sizes
 - Integrated, efficient [hyper-parameter optimization module (HPO)](https://openvinotoolkit.github.io/training_extensions/latest/guide/explanation/additional_features/hpo.html). Through dataset proxy and built-in hyper-parameter optimizer, you can get much faster hyper-parameter optimization compared to other off-the-shelf tools. The hyperparameter optimization is dynamically scheduled based on your resource budget.
 
 ---
 
-## Getting Started
-
-### Installation
+## Installation
 
 Please refer to the [installation guide](https://openvinotoolkit.github.io/training_extensions/latest/guide/get_started/installation.html).
+If you want to make changes to the library, then a local installation is recommended.
+
+<details>
+<summary>Install from PyPI</summary>
+Installing the library with pip is the easiest way to get started with otx.
+
+```bash
+pip install otx
+```
+
+This will install OTX CLI. OTX requires torch and lightning by default to provide training. To use the full pipeline, you need the commands below:
+
+```bash
+# Get help for the installation arguments
+otx install -h
+
+# Install the full package
+otx install
+
+# Install with verbose output
+otx install -v
+
+# Install with docs option only.
+otx install --option docs
+```
+
+</details>
+
+<details>
+<summary>Install from source</summary>
+To install from source, you need to clone the repository and install the library using pip via editable mode.
+
+```bash
+# Use of virtual environment is highy recommended
+# Using conda
+yes | conda create -n otx_env python=3.10
+conda activate otx_env
+
+# Or using your favorite virtual environment
+# ...
+
+# Clone the repository and install in editable mode
+git clone https://github.com/openvinotoolkit/training_extensions.git
+cd training_extensions
+pip install -e .
+```
+
+This will install OTX CLI. OTX requires torch and lightning by default to provide training. To use the full pipeline, you need the commands below:
+
+```bash
+# Get help for the installation arguments
+otx install -h
+
+# Install the full package
+otx install
+
+# Install with verbose output
+otx install -v
+
+# Install with docs option only.
+otx install --option docs
+```
+
+</details>
+
+---
+
+## Quick-Start
+
+OpenVINO™ Training Extensions supports both API and CLI-based training. The API is more flexible and allows for more customization, while the CLI training utilizes command line interfaces, and might be easier for those who would like to use OpenVINO™ Training Extensions off-the-shelf.
+
+For the CLI, the commands below provide subcommands, how to use each subcommand, and more:
+
+```bash
+# See available subcommands
+otx --help
+
+# Print help messages from the train subcommand
+otx train --help
+```
+
+You can find details with examples in the [CLI Guide](https://openvinotoolkit.github.io/training_extensions/latest/guide/get_started/cli_commands.html). and [API Quick-Guide](https://openvinotoolkit.github.io/training_extensions/latest/guide/get_started/api_tutorial.html).
+
+Below is how to train with auto-configuration, which is provided to users with datasets and tasks:
+
+<details>
+<summary>Training via API</summary>
+
+```python
+# Training with Auto-Configuration via Engine
+from otx.engine import Engine
 
-Note: Python 3.8, 3.9 and 3.10 were tested, along with Ubuntu 18.04, 20.04 and 22.04.
+engine = Engine(data_root="data/wgisd", task="DETECTION")
+engine.train()
+```
 
-### OpenVINO™ Training Extensions CLI Commands
+For more examples, see documentation: [CLI Guide](https://openvinotoolkit.github.io/training_extensions/latest/guide/get_started/cli_commands.html)
 
-- `otx find` helps you quickly find the best pre-configured models templates as well as a list of supported backbones
-- `otx build` creates the workspace folder with all necessary components to start training. It can help you configure your own model with any supported backbone and even prepare a custom split for your dataset
-- `otx train` actually starts training on your dataset
-- `otx eval` runs evaluation of your trained model in PyTorch or OpenVINO™ IR format
-- `otx optimize` runs an optimization algorithm to quantize and prune your deep learning model with help of [NNCF](https://github.com/openvinotoolkit/nncf) and [POT](https://docs.openvino.ai/latest/pot_introduction.html) tools.
-- `otx export` starts exporting your model to the OpenVINO™ IR format
-- `otx deploy` outputs the exported model together with the self-contained python package, a demo application to port and infer it outside of this repository.
-- `otx demo` allows one to apply a trained model on the custom data or the online footage from a web camera and see how it will work in a real-life scenario.
-- `otx explain` runs explain algorithm on the provided data and outputs images with the saliency maps to show how your model makes predictions.
+</details>
 
-You can find more details with examples in the [CLI command intro](https://openvinotoolkit.github.io/training_extensions/latest/guide/get_started/cli_commands.html).
+<details>
+<summary>Training via CLI</summary>
+
+```bash
+otx train --data_root data/wgisd --task DETECTION
+```
+
+For more examples, see documentation: [API Quick-Guide](https://openvinotoolkit.github.io/training_extensions/latest/guide/get_started/api_tutorial.html)
+
+</details>
+
+In addition to the examples above, please refer to the documentation for tutorials on using custom models, training parameter overrides, and [tutorial per task types](https://openvinotoolkit.github.io/training_extensions/latest/guide/tutorials/base/how_to_train/index.html), etc.
 
 ---
 
 ## Updates
 
-### v1.4.0 (3Q23)
+### v2.0.0 (1Q24)
 
-- Support encrypted dataset training (<https://github.com/openvinotoolkit/training_extensions/pull/2209>)
-- Add custom max iou assigner to prevent CPU OOM when large annotations are used (<https://github.com/openvinotoolkit/training_extensions/pull/2228>)
-- Auto train type detection for Semi-SL, Self-SL and Incremental: "--train-type" now is optional (<https://github.com/openvinotoolkit/training_extensions/pull/2195>)
-- Add per-class XAI saliency maps for Mask R-CNN model (<https://github.com/openvinotoolkit/training_extensions/pull/2227>)
-- Add new object detector Deformable DETR (<https://github.com/openvinotoolkit/training_extensions/pull/2249>)
-- Add new object detector DINO (<https://github.com/openvinotoolkit/training_extensions/pull/2266>)
-- Add new visual prompting task (<https://github.com/openvinotoolkit/training_extensions/pull/2203>, <https://github.com/openvinotoolkit/training_extensions/pull/2274>, <https://github.com/openvinotoolkit/training_extensions/pull/2311>, <https://github.com/openvinotoolkit/training_extensions/pull/2354>, <https://github.com/openvinotoolkit/training_extensions/pull/2318>)
-- Add new object detector ResNeXt101-ATSS (<https://github.com/openvinotoolkit/training_extensions/pull/2309>)
+TBD
 
 ### Release History
 
@@ -151,3 +236,15 @@ See Intel's [Global Human Rights Principles](https://www.intel.com/content/www/u
 Intel's products and software are intended only to be used in applications that do not cause or contribute to a violation of an internationally recognized human right.
 
 ---
+
+## Contributing
+
+For those who would like to contribute to the library, see [CONTRIBUTING.md](CONTRIBUTING.md) for details.
+
+Thank you! we appreciate your support!
+
+<a href="https://github.com/openvinotoolkit/training_extensions/graphs/contributors">
+  <img src="https://contrib.rocks/image?repo=openvinotoolkit/training_extensions" />
+</a>
+
+---

pyproject.toml

@@ -30,7 +30,7 @@ dependencies = [
     "datumaro==1.6.0rc1",
     "omegaconf==2.3.0",
     "rich==13.7.1",
-    "jsonargparse==4.27.6",
+    "jsonargparse==4.27.1",
     "psutil==5.9.8", # Mem cache needs system checks
     "ftfy==6.1.3",
     "regex==2023.12.25",
@@ -89,18 +89,7 @@ mmlab = [
     "oss2==2.17.0",
 ]
 anomaly = [
-    # [FIXME] @ashwinvaidya17: Install using a temporary hot-fix commit due to a torchmetrics version conflict.
-    "anomalib @ git+https://github.com/openvinotoolkit/anomalib.git@e78091883a620229c277a79674a904d9f785f8d5",
-    # This is a dependency to avoid conflicts with installing the anomalib[core] option.
-    "av>=10.0.0",
-    "einops>=0.3.2",
-    "freia>=0.2",
-    "imgaug==0.4.0",
-    "kornia>=0.6.6,<0.6.10",
-    "matplotlib>=3.4.3",
-    "opencv-python>=4.5.3.56",
-    "pandas>=1.1.0",
-    "open-clip-torch>=2.23.0",
+    "anomalib[core]==1.0.1",
 ]
 
 [project.scripts]

src/otx/algo/visual_prompting/backbones/tiny_vit.py

@@ -9,7 +9,7 @@
 
 import torch
 import torch.nn.functional as F  # noqa: N812
-from timm.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
 from torch import Tensor, nn
 
 from otx.algo.visual_prompting.utils.layer_norm_2d import LayerNorm2d

src/otx/cli/cli.py

@@ -454,6 +454,7 @@ def save_config(self, work_dir: Path) -> None:
         The configuration is saved as a YAML file in the engine's working directory.
         """
         self.config[self.subcommand].pop("workspace", None)
+        self.config[self.subcommand]["work_dir"] = str(self.workspace.work_dir.parent)
         # TODO(vinnamki): Revisit it after changing the optimizer and scheduler instantiating.
         cfg = deepcopy(self.config.get(str(self.subcommand), self.config))
         cfg.model.init_args.pop("optimizer")

src/otx/core/model/classification.py

@@ -13,7 +13,6 @@
 import torch
 from torchmetrics import Accuracy
 
-from otx.algo.hooks.recording_forward_hook import feature_vector_fn
 from otx.core.data.entity.base import (
     OTXBatchLossEntity,
     T_OTXBatchDataEntity,
@@ -84,6 +83,8 @@ def head_forward_fn(self, x: torch.Tensor) -> torch.Tensor:
 
     def forward_explain(self, inputs: T_OTXBatchDataEntity) -> T_OTXBatchPredEntity:
         """Model forward function."""
+        from otx.algo.hooks.recording_forward_hook import feature_vector_fn
+
         self.model.feature_vector_fn = feature_vector_fn
         self.model.explain_fn = self.get_explain_fn()
 
@@ -143,6 +144,8 @@ def get_explain_fn(self) -> Callable:
         return explainer.func
 
     def _reset_model_forward(self) -> None:
+        from otx.algo.hooks.recording_forward_hook import feature_vector_fn
+
         if not self.explain_mode:
             return
 

src/otx/core/model/instance_segmentation.py

@@ -17,7 +17,6 @@
 from openvino.model_api.tilers import InstanceSegmentationTiler
 from torchvision import tv_tensors
 
-from otx.algo.hooks.recording_forward_hook import MaskRCNNRecordingForwardHook, feature_vector_fn
 from otx.core.config.data import TileConfig
 from otx.core.data.entity.base import OTXBatchLossEntity
 from otx.core.data.entity.instance_segmentation import InstanceSegBatchDataEntity, InstanceSegBatchPredEntity
@@ -241,6 +240,8 @@ def forward_explain(
         inputs: InstanceSegBatchDataEntity,
     ) -> InstanceSegBatchPredEntity:
         """Model forward function."""
+        from otx.algo.hooks.recording_forward_hook import feature_vector_fn
+
         self.model.feature_vector_fn = feature_vector_fn
         self.model.explain_fn = self.get_explain_fn()
 
@@ -298,6 +299,8 @@ def _forward_explain_inst_seg(
 
     def get_explain_fn(self) -> Callable:
         """Returns explain function."""
+        from otx.algo.hooks.recording_forward_hook import MaskRCNNRecordingForwardHook
+
         explainer = MaskRCNNRecordingForwardHook(num_classes=self.num_classes)
         return explainer.func
 

src/otx/engine/engine.py

@@ -80,23 +80,26 @@ class Engine:
     Example:
         The following examples show how to use the Engine class.
 
-        Auto-Configuration with data_root
-        >>> engine = Engine(
-        ...     data_root=<dataset/path>,
-        ... )
-
-        Create Engine with Custom OTXModel
-        >>> engine = Engine(
-        ...     data_root=<dataset/path>,
-        ...     model=OTXModel(...),
-        ...     checkpoint=<checkpoint/path>,
-        ... )
-
-        Create Engine with Custom OTXDataModule
-        >>> engine = Engine(
-        ...     model = OTXModel(...),
-        ...     datamodule = OTXDataModule(...),
-        ... )
+        Auto-Configuration with data_root::
+
+            engine = Engine(
+                data_root=<dataset/path>,
+            )
+
+        Create Engine with Custom OTXModel::
+
+            engine = Engine(
+                data_root=<dataset/path>,
+                model=OTXModel(...),
+                checkpoint=<checkpoint/path>,
+            )
+
+        Create Engine with Custom OTXDataModule::
+
+            engine = Engine(
+                model = OTXModel(...),
+                datamodule = OTXDataModule(...),
+            )
     """
 
     _EXPORTED_MODEL_BASE_NAME: ClassVar[str] = "exported_model"
@@ -546,6 +549,7 @@ def optimize(
             ...     datamodule=OTXDataModule(),
             ...     checkpoint=<checkpoint/path>,
             ... )
+
         CLI Usage:
             To optimize a model, run
                 ```python
@@ -755,17 +759,17 @@ def from_model_name(
             ...     data_root=<dataset/path>,
             ... )
 
-            If you want to override configuration from default config
-            >>> overriding = {
-            ...     "data.config.train_subset.batch_size": 2,
-            ...     "data.config.test_subset.subset_name": "TESTING",
-            ... }
-            >>> engine = Engine(
-            ...     model_name="atss_mobilenetv2",
-            ...     task="DETECTION",
-            ...     data_root=<dataset/path>,
-            ...     **overriding,
-            ... )
+            If you want to override configuration from default config:
+                >>> overriding = {
+                ...     "data.config.train_subset.batch_size": 2,
+                ...     "data.config.test_subset.subset_name": "TESTING",
+                ... }
+                >>> engine = Engine(
+                ...     model_name="atss_mobilenetv2",
+                ...     task="DETECTION",
+                ...     data_root=<dataset/path>,
+                ...     **overriding,
+                ... )
         """
         default_config = DEFAULT_CONFIG_PER_TASK.get(task)
         model_path = str(default_config).split("/")