From a436dd7e0146d63e01ab41fb72eb6d78021196c8 Mon Sep 17 00:00:00 2001
From: Oleksandr Bezdieniezhnykh <oleksandr.bezdieniezhnykh@pwc.com>
Date: Wed, 13 May 2026 00:54:51 +0300
Subject: [PATCH] refactor: clean up exports and constants

- Removed unused imports and the `form_data_sample` function from `exports.py` to streamline the code.
- Added whitespace in `constants.py` for improved readability.

These changes enhance code clarity and maintainability.
---
 src/constants.py   |  3 ++-
 src/data_sample.py | 64 ++++++++++++++++++++++++++++++++++++++++++++++
 src/exports.py     | 26 +------------------
 3 files changed, 67 insertions(+), 26 deletions(-)
 create mode 100644 src/data_sample.py

diff --git a/src/constants.py b/src/constants.py
index d880ea6..a87b713 100644
--- a/src/constants.py
+++ b/src/constants.py
@@ -54,7 +54,8 @@ class Config(BaseModel):
     @property
     def data_dir(self) -> str:
         return path.join(self.dirs.root, self.dirs.data)
-
+    
+    
     @property
     def images_dir(self) -> str:
         return path.join(self.data_dir, IMAGES_DIR)
diff --git a/src/data_sample.py b/src/data_sample.py
new file mode 100644
index 0000000..25fc099
--- /dev/null
+++ b/src/data_sample.py
@@ -0,0 +1,64 @@
+import argparse
+import random
+import shutil
+from os import path, scandir, makedirs
+
+import constants
+
+
+def form_data_sample(destination_path: str = None, size: int = 500, write_txt_log: bool = False) -> None:
+    if destination_path is None:
+        destination_path = constants.config.sample_dir
+
+    images = []
+    with scandir(constants.config.images_dir) as imd:
+        for image_file in imd:
+            if not image_file.is_file():
+                continue
+            images.append(image_file)
+    print('shuffling images')
+    random.shuffle(images)
+    images = images[:size]
+
+    shutil.rmtree(destination_path, ignore_errors=True)
+    makedirs(destination_path, exist_ok=True)
+
+    lines = []
+    for image in images:
+        shutil.copy(image.path, path.join(destination_path, image.name))
+        lines.append(f'./{image.name}')
+    if write_txt_log:
+        with open(path.join(destination_path, 'azaion_subset.txt'), 'w', encoding='utf-8') as f:
+            f.writelines([f'{line}\n' for line in lines])
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description='Form a random image sample (e.g. for int8 calibration).',
+    )
+    parser.add_argument(
+        '--dest', '-d',
+        default=None,
+        help='Destination directory. Defaults to constants.config.sample_dir. WARNING: wiped before writing.',
+    )
+    parser.add_argument(
+        '--size', '-n',
+        type=int,
+        default=500,
+        help='Number of images to sample (default: 500).',
+    )
+    parser.add_argument(
+        '--write-txt-log',
+        action='store_true',
+        help='Also write azaion_subset.txt listing the sampled file names.',
+    )
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = _parse_args()
+    form_data_sample(
+        destination_path=args.dest,
+        size=args.size,
+        write_txt_log=args.write_txt_log,
+    )
diff --git a/src/exports.py b/src/exports.py
index a4c5487..f75ec17 100644
--- a/src/exports.py
+++ b/src/exports.py
@@ -1,8 +1,7 @@
 import os
 import shutil
-from os import path, scandir, makedirs
+from os import path
 from pathlib import Path
-import random
 
 import netron
 import yaml
@@ -60,29 +59,6 @@ def export_tensorrt(model_path):
     )
 
 
-def form_data_sample(destination_path, size=500, write_txt_log=False):
-    images = []
-    with scandir(constants.config.images_dir) as imd:
-        for image_file in imd:
-            if not image_file.is_file():
-                continue
-            images.append(image_file)
-    print('shuffling images')
-    random.shuffle(images)
-    images = images[:size]
-
-    shutil.rmtree(destination_path, ignore_errors=True)
-    makedirs(destination_path, exist_ok=True)
-
-    lines = []
-    for image in images:
-        shutil.copy(image.path, path.join(destination_path, image.name))
-        lines.append(f'./{image.name}')
-    if write_txt_log:
-        with open(path.join(destination_path, 'azaion_subset.txt'), 'w', encoding='utf-8') as f:
-            f.writelines([f'{line}\n' for line in lines])
-
-
 def show_model(model: str = None):
     netron.start(model)