From 712c15bd595f83b599c8e44a9aeb362653dbc50c Mon Sep 17 00:00:00 2001
From: Songki Choi <songki.choi@intel.com>
Date: Tue, 4 Apr 2023 13:33:37 +0900
Subject: [PATCH] Fix classification model download logic to resolve zip issue
 (#1967)

* Replace model download logic by pytorchcv to resolve zip issue

* Update change log

* Fix pre-commit

* Configure checkpoint before model

* Add pytorchcv to cls/seg requirement

---------
Signed-off-by: Songki Choi <songki.choi@intel.com>
---
 CHANGELOG.md                                                   | 3 ++-
 otx/algorithms/classification/adapters/mmcls/tasks/stage.py    | 2 +-
 .../common/adapters/mmcv/models/backbones/efficientnet.py      | 3 ++-
 requirements/classification.txt                                | 1 +
 requirements/segmentation.txt                                  | 1 +
 5 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b57cc57e657..4e7fe78ec5a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -28,7 +28,8 @@ All notable changes to this project will be documented in this file.
 
 - Fix exception -> warning for anomaly dump_feature option
 - Remove `dataset.with_empty_annotations()` to keep original input structure (<https://github.com/openvinotoolkit/training_extensions/pull/1964>)
-- Fix OV batch inference (saliency map generation)  (<https://github.com/openvinotoolkit/training_extensions/pull/1965>)
+- Fix OV batch inference (saliency map generation) (<https://github.com/openvinotoolkit/training_extensions/pull/1965>)
+- Replace EfficentNetB0 model download logic by pytorchcv to resolve zip issue (<https://github.com/openvinotoolkit/training_extensions/pull/1967>)
 
 ## \[v1.1.1\]
 
diff --git a/otx/algorithms/classification/adapters/mmcls/tasks/stage.py b/otx/algorithms/classification/adapters/mmcls/tasks/stage.py
index 11dd3c9aefc..b2504f11daa 100644
--- a/otx/algorithms/classification/adapters/mmcls/tasks/stage.py
+++ b/otx/algorithms/classification/adapters/mmcls/tasks/stage.py
@@ -30,8 +30,8 @@ def configure(self, model_cfg, model_ckpt, data_cfg, training=True, **kwargs):
 
         # Recipe + model
         cfg = self.cfg
-        self.configure_model(cfg, model_cfg, **kwargs)
         self.configure_ckpt(cfg, model_ckpt, kwargs.get("pretrained", None))
+        self.configure_model(cfg, model_cfg, **kwargs)
         self.configure_data(cfg, data_cfg, training)
         self.configure_task(cfg, training)
         return cfg
diff --git a/otx/algorithms/common/adapters/mmcv/models/backbones/efficientnet.py b/otx/algorithms/common/adapters/mmcv/models/backbones/efficientnet.py
index e16541f6327..c7ca005c13b 100644
--- a/otx/algorithms/common/adapters/mmcv/models/backbones/efficientnet.py
+++ b/otx/algorithms/common/adapters/mmcv/models/backbones/efficientnet.py
@@ -18,6 +18,7 @@
 from mmcv.cnn import build_activation_layer
 from mmcv.cnn.bricks import ConvModule
 from mmcv.runner import load_checkpoint
+from pytorchcv.models.model_store import download_model
 from torch import nn
 from torch.nn import init
 
@@ -1281,5 +1282,5 @@ def init_weights(self, pretrained=None):
             load_checkpoint(self, pretrained)
             logger.info(f"init weight - {pretrained}")
         elif pretrained is not None:
-            load_checkpoint(self, pretrained_urls[self.model_name])
+            download_model(net=self, model_name=self.model_name)
             logger.info(f"init weight - {pretrained_urls[self.model_name]}")
diff --git a/requirements/classification.txt b/requirements/classification.txt
index 3efc94c8a9e..f65df21b120 100644
--- a/requirements/classification.txt
+++ b/requirements/classification.txt
@@ -4,3 +4,4 @@ mmcv-full==1.7.0
 mmcls==0.25.0
 timm
 otxdeploy==0.12.1  # FIXME: mmdeploy==? when PyPI packages are available
+pytorchcv
diff --git a/requirements/segmentation.txt b/requirements/segmentation.txt
index fcf033eb799..34980f90081 100644
--- a/requirements/segmentation.txt
+++ b/requirements/segmentation.txt
@@ -5,3 +5,4 @@ mmsegmentation==0.30.0
 scikit-image
 otxdeploy==0.12.1  # FIXME: mmdeploy==? when PyPI packages are available
 timm
+pytorchcv