From 94d4f3fd7b64e21340c3e50c1cc08ea68f5c52b3 Mon Sep 17 00:00:00 2001
From: Vladislav Sovrasov <sovrasov.vlad@gmail.com>
Date: Mon, 16 Dec 2024 09:42:20 +0100
Subject: [PATCH] BC improvement (#4154)

* Add OTX version to exported models

* Forward data format to tiling dataset to fix arrow handling

* Workaround missing label_ids

* Version up

* Update release notes

* Update changelog
---
 CHANGELOG.md                              |  7 +++++++
 docs/source/guide/release_notes/index.rst | 11 +++++++++++
 src/otx/__init__.py                       |  2 +-
 src/otx/core/data/dataset/tile.py         |  1 +
 src/otx/core/model/base.py                |  5 +++++
 src/otx/core/types/export.py              |  2 ++
 tests/unit/core/types/test_export.py      |  1 +
 7 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cfbed5b2df5..45c1096ed4f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,13 @@
 
 All notable changes to this project will be documented in this file.
 
+## \[2.2.2\]
+
+### Bug fixes
+
+- BC improvement
+  (<https://github.com/openvinotoolkit/training_extensions/pull/4154>)
+
 ## \[2.2.1\]
 
 ### Bug fixes
diff --git a/docs/source/guide/release_notes/index.rst b/docs/source/guide/release_notes/index.rst
index 182d908ec35..6a826256dcf 100644
--- a/docs/source/guide/release_notes/index.rst
+++ b/docs/source/guide/release_notes/index.rst
@@ -4,6 +4,17 @@ Releases
 .. toctree::
   :maxdepth: 1
 
+v2.2.2 (2024.12)
+----------------
+
+Enhancements
+^^^^^^^^^^^^
+
+Bug fixes
+^^^^^^^^^
+
+- BC improvement
+
 v2.2.1 (2024.12)
 ----------------
 
diff --git a/src/otx/__init__.py b/src/otx/__init__.py
index a52507dcfda..d241ccd04b0 100644
--- a/src/otx/__init__.py
+++ b/src/otx/__init__.py
@@ -3,7 +3,7 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "2.2.1"
+__version__ = "2.2.2"
 
 import os
 from pathlib import Path
diff --git a/src/otx/core/data/dataset/tile.py b/src/otx/core/data/dataset/tile.py
index b94931487d2..97ff44b913f 100644
--- a/src/otx/core/data/dataset/tile.py
+++ b/src/otx/core/data/dataset/tile.py
@@ -263,6 +263,7 @@ def __init__(self, dataset: OTXDataset, tile_config: TileConfig) -> None:
             dataset.image_color_channel,
             dataset.stack_images,
             dataset.to_tv_image,
+            data_format=dataset.data_format,
         )
         self.tile_config = tile_config
         self._dataset = dataset
diff --git a/src/otx/core/model/base.py b/src/otx/core/model/base.py
index d4176fd0cbb..d65e420b90c 100644
--- a/src/otx/core/model/base.py
+++ b/src/otx/core/model/base.py
@@ -391,6 +391,11 @@ def load_state_dict_incrementally(self, ckpt: dict[str, Any], *args, **kwargs) -
             msg = "Checkpoint should have `label_info`."
             raise ValueError(msg, ckpt_label_info)
 
+        if not hasattr(ckpt_label_info, "label_ids"):
+            msg = "Loading checkpoint from OTX < 2.2.1, label_ids are assigned automatically"
+            logger.info(msg)
+            ckpt_label_info.label_ids = [str(i) for i, _ in enumerate(ckpt_label_info.label_names)]
+
         if ckpt_label_info != self.label_info:
             msg = (
                 "Load model state dictionary incrementally: "
diff --git a/src/otx/core/types/export.py b/src/otx/core/types/export.py
index a66541d0873..18a48e87e03 100644
--- a/src/otx/core/types/export.py
+++ b/src/otx/core/types/export.py
@@ -9,6 +9,7 @@
 from dataclasses import dataclass, fields
 from enum import Enum
 
+import otx
 from otx.core.config.data import TileConfig
 from otx.core.types.label import HLabelInfo, LabelInfo
 
@@ -122,6 +123,7 @@ def to_metadata(self) -> dict[tuple[str, str], str]:
             ("model_info", "labels"): all_labels.strip(),
             ("model_info", "label_ids"): all_label_ids.strip(),
             ("model_info", "optimization_config"): json.dumps(self.optimization_config),
+            ("model_info", "otx_version"): otx.__version__,
         }
 
         if isinstance(self.label_info, HLabelInfo):
diff --git a/tests/unit/core/types/test_export.py b/tests/unit/core/types/test_export.py
index 72add6f3c31..70a4aa1aa2f 100644
--- a/tests/unit/core/types/test_export.py
+++ b/tests/unit/core/types/test_export.py
@@ -52,3 +52,4 @@ def test_wrap(fxt_label_info, task_type):
     assert ("model_info", "tile_size") in metadata
     assert ("model_info", "tiles_overlap") in metadata
     assert ("model_info", "max_pred_number") in metadata
+    assert ("model_info", "otx_version") in metadata