pms-nvidia-processor 2.6.0__tar.gz → 2.6.1__tar.gz

{pms_nvidia_processor-2.6.0 → pms_nvidia_processor-2.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pms-nvidia-processor
-Version: 2.6.0
+Version: 2.6.1
 Summary: lib for pms nvidia processors
 Author: HyeongSeok Kim
 Author-email: tiryul@gmail.com

{pms_nvidia_processor-2.6.0 → pms_nvidia_processor-2.6.1}/pms_nvidia_processor/__init__.py

@@ -29,4 +29,4 @@ __all__: list[str] = [
     "CDRUF2Processor",
 ]
 
-__version__ = "2.6.0"
+__version__ = "2.6.1"

{pms_nvidia_processor-2.6.0 → pms_nvidia_processor-2.6.1}/pms_nvidia_processor/processor/col_en/config.py

@@ -6,11 +6,11 @@ class ColEnConfig:
     NUMBER_OF_INPUT_CHANNELS: int = 3 * NUMBER_OF_FRAMES
     NUMBER_OF_OUTPUT_CHANNELS: int = 3
     UPSCALE_RATIO: int = 1
-    PATCH_SIZE = 1536
+    PATCH_SIZE = 1280
     MAX_BATCH_SIZE = 1
     MIN_BATCH_SIZE = 1
     OPT_BATCH_SIZE = 1
-    INPUT_OVERLAB_LENGTH = PATCH_SIZE // 2
+    INPUT_OVERLAB_LENGTH = 640  # PATCH_SIZE // 2
 
     PATCHER_CONFIG = PatcherWIOConfig(
         patch_size=PATCH_SIZE,

pms_nvidia_processor-2.6.1/pms_nvidia_processor/utility/builder.py

@@ -0,0 +1,361 @@
+import os
+import tempfile as tf
+import pms_tensorrt as trt
+import pms_model_manager as pmm
+from ..processor.dpir.config import DPIRConfig
+from ..processor.dru_asm_sr_f3.config import DRUASMSRF3Config
+from ..processor.dru_rbpn_deinter_f3_gloss.config import (
+    DRURBPNDEINTERF3GLOSSConfig,
+)
+from ..processor.dru_rbpn_deinter_pc_f2.config import (
+    DRURBPNDEINTERPCF2Config,
+)
+from ..processor.dru_rbpn_sr_f3.config import DRURBPNSRF3Config
+from ..processor.dru_rbpn_sr_f5.config import DRURBPNSRF5Config
+from ..processor.color_resnet.config import (
+    ColorResnetPostConfig,
+    ColorResnetPreConfig,
+)
+from ..processor.fisf.config import FISFConfig
+from ..processor.gg.config import GGConfig
+from ..processor.col_en.config import ColEnConfig
+from ..processor.cdru_f2.config import CDRUF2Config
+
+__MODEL_CONFIG_MAP = {
+    "DPIR": DPIRConfig,
+    "DRU-RBPN-SR-F3": DRURBPNSRF3Config,
+    "DRU-RBPN-SR-F5": DRURBPNSRF5Config,
+    "DRU-ASM-SR-F3": DRUASMSRF3Config,
+    "DRU-RBPN-DEINTER-F3-GLOSS": DRURBPNDEINTERF3GLOSSConfig,
+    "DRU-RBPN-DEINTER-PC-F2": DRURBPNDEINTERPCF2Config,
+    "FISF": FISFConfig,
+    "COLOR-RESNET-PRE": ColorResnetPreConfig,
+    "COLOR-RESNET-POST": ColorResnetPostConfig,
+    "GG": GGConfig,
+    "COL-EN": ColEnConfig,
+    "CDRU-F2": CDRUF2Config,
+}
+
+
+def create_build_params(
+    model_name: str,
+    device: int = 0,
+    downlaod_model_alias="newest",
+):  # -> dict[str, Any]:
+    import pms_tensorrt as trt
+
+    assert (
+        model_name in __MODEL_CONFIG_MAP
+    ), f"ERROR!, model_name={model_name} is not in the list. {__MODEL_CONFIG_MAP.keys()}"
+    params: dict = {
+        "model_name": model_name,
+        "device": device,
+        "downlaod_model_alias": downlaod_model_alias,
+    }
+    config = __MODEL_CONFIG_MAP[model_name]
+    if config is FISFConfig:
+        params.update(
+            {
+                "build_config_flags": [],
+                "shape_profile": {
+                    "input0": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                    },
+                    "input1": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                    },
+                    "output": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                    },
+                },
+            }
+        )
+    elif config is ColorResnetPreConfig:
+        params.update(
+            {
+                "build_config_flags": [],
+                "shape_profile": {
+                    "input": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_SHAPE_INPUT,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_SHAPE_INPUT,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_SHAPE_INPUT,
+                        ],
+                    },
+                    "output": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_SHAPE_OUTPUT,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_SHAPE_OUTPUT,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_SHAPE_OUTPUT,
+                        ],
+                    },
+                },
+            }
+        )
+    elif config is ColorResnetPostConfig:
+        params.update(
+            {
+                "build_config_flags": [],
+                "shape_profile": {
+                    "input": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                    },
+                    "model_output": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_SHAPE_MODEL_OUTPUT,
+                        ],
+                        "opt_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_SHAPE_MODEL_OUTPUT,
+                        ],
+                        "max_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_SHAPE_MODEL_OUTPUT,
+                        ],
+                    },
+                    "output": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                    },
+                },
+            }
+        )
+    elif config is GGConfig:
+        params.update(
+            {
+                "build_config_flags": [],
+                "shape_profile": {
+                    "input_image": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                    },
+                    "grain_level": {
+                        "min_shape": [config.MIN_BATCH_SIZE, 1, 1, 1],
+                        "opt_shape": [config.OPT_BATCH_SIZE, 1, 1, 1],
+                        "max_shape": [config.MAX_BATCH_SIZE, 1, 1, 1],
+                    },
+                    "grain_noise": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                    },
+                    "output": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                    },
+                },
+            }
+        )
+    else:
+        params.update(
+            {
+                "build_config_flags": [
+                    trt.TRTBuilderFlag.FP16,
+                ],
+                "shape_profile": {
+                    "input": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.input_shape,
+                        ],
+                    },
+                    "output": {
+                        "min_shape": [
+                            config.MIN_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                        "opt_shape": [
+                            config.OPT_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                        "max_shape": [
+                            config.MAX_BATCH_SIZE,
+                            *config.TRT_CONFIG.output_shape,
+                        ],
+                    },
+                },
+            }
+        )
+    return params
+
+
+def build_and_upload_plan(
+    model_name: str,
+    shape_profile: dict[str, dict[str, list[int]]],
+    build_config_flags: list[trt.TRTBuilderFlag],
+    device: int = 0,
+    onnx_model_file_name="model.onnx",
+    plan_model_file_name="model.plan",
+    downlaod_model_alias="newest",
+):
+    download_model_name = f"{model_name}#onnx"
+    device_name = "nvidia-" + "-".join(
+        s for s in trt.get_device_list()[0].split(":")[1].lower().split(" ")[2:-1]
+    )
+    if device_name.endswith("-d"):
+        device_name = device_name[:-2]
+    upload_model_name = f"{model_name}#{device_name}"
+    upload_model_alias = "newest"
+
+    number_of_devices = trt.get_device_count()
+    assert number_of_devices > 0, "There is no NVIDIA GPUs."
+    assert number_of_devices > device, f"device {device} is not exist."
+    assert device > -1, f"device_id {device} is not exist."
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(device)
+
+    upload_tag = {
+        f"{io_name}.{shape_name}": str(shape_profile[io_name][shape_name])
+        for io_name in shape_profile
+        for shape_name in shape_profile[io_name]
+    }
+    for flag in build_config_flags:
+        upload_tag.update({flag.name: ""})
+    print(upload_tag)
+    with tf.TemporaryDirectory() as temp_dir:
+        print(f"Temporary Directory is {temp_dir}")
+        model_manager = pmm.ModelManager(temp_dir)
+        model_manager.download(
+            model_name=download_model_name, alias=downlaod_model_alias
+        )
+        meta_data = pmm.MLFlowModel.load_metadata(
+            model_manager.get_metadata_path(
+                model_name=download_model_name, alias=downlaod_model_alias
+            )
+        )
+        upload_tag["source_version"] = str(meta_data["version"])
+        upload_tag["source_aliases"] = str(meta_data["aliases"])
+        upload_tag["source_run_id"] = str(meta_data["run_id"])
+        model_onnx_dir = model_manager.get_local_model_dir(
+            model_name=download_model_name,
+            alias=downlaod_model_alias,
+        )
+        model_plan_dir = model_manager.get_local_model_dir(
+            model_name=upload_model_name,
+            alias=upload_model_alias,
+        )
+        os.makedirs(model_plan_dir, exist_ok=False)
+        onnx_model_path = os.path.join(model_onnx_dir, onnx_model_file_name)
+        plan_model_path = os.path.join(model_plan_dir, plan_model_file_name)
+        builder = trt.EngineBuilder()
+        builder.build_from_onnx(
+            onnx_path=onnx_model_path,
+            plan_path=plan_model_path,
+            config_flags=build_config_flags,
+            shape_profiles=[
+                trt.ShapeProfile(name=io_name, **shape_profile[io_name])
+                for io_name in shape_profile
+            ],
+        )
+        os.system(f"ls {os.path.dirname(plan_model_path)}")
+        model_manager.upload(
+            model_name=upload_model_name,
+            model_dir=model_plan_dir,
+            aliases=[upload_model_alias],
+            tag=upload_tag,
+        )

{pms_nvidia_processor-2.6.0 → pms_nvidia_processor-2.6.1}/pms_nvidia_processor/utility/patcher/_patcher_w.py

@@ -32,7 +32,7 @@ class PatcherW:
         self._input_overlap_length = input_overlap_length
 
         self._patch_pos_list = self._get_patch_pos_list()
-        self._weight_mask = self._make_weight_mask(output_patch_shape)
+        self._weight_mask = self._make_hanning_weight_mask(output_patch_shape)
 
     def _get_patch_pos_list(self):
         input_vector_width = self._input_vector_shape[1]
@@ -52,7 +52,7 @@ class PatcherW:
                 patchs.append((slice(y0, y1), slice(x0, x1)))
         return patchs
 
-    def _make_weight_mask(
+    def _make_gaussian_weight_mask(
         self,
         shape: tuple[int, int, int],
         min_val=0.0,
@@ -63,7 +63,6 @@ class PatcherW:
         h, w: 패치 높이/너비
         sigma_scale: 중심에서 분포하는 정도 (0.25이면 가우시안 표준편차가 크기 대비 1/4)
         """
-        import cv2
 
         h, w, c = shape
         y = np.linspace(-1, 1, h)
@@ -75,12 +74,44 @@ class PatcherW:
         np.subtract(weights, weights.min(), out=weights)
         np.divide(weights, weights.max(), out=weights)
         weights = weights * (max_val - min_val) + min_val
-        cv2.imwrite(
-            "/home/hskim/pms-nvidia-processor/.pytest_cache/weights.png",
-            (weights * 255).astype(np.uint8),
-        )
-        return weights.astype(np.float32)[:, :, None]  # type cast and add dimension
-        # return np.stack([weights for _ in range(c)], axis=-1, dtype=np.float32)
+        return weights.astype(np.float32)[:, :, None]
+
+    def _make_hanning_weight_mask(
+        self,
+        shape: tuple[int, int, int],
+        min_val=0.0,
+        max_val=1.0,
+    ):
+        """
+        Hanning window 기반의 weight mask를 생성합니다.
+
+        Parameters:
+            shape (tuple[int, int, int]): (h, w, c) - 패치의 높이, 너비, 채널 수
+            min_val (float): weight의 최소값 (스케일 조정용)
+            max_val (float): weight의 최대값 (스케일 조정용)
+
+        Returns:
+            np.ndarray: shape = (h, w, c), float32 타입의 Hanning weight mask
+        """
+        h, w, c = shape
+
+        hann_y = np.hanning(h)
+        hann_x = np.hanning(w)
+        hann_2d = np.outer(hann_y, hann_x).astype(np.float32)
+        # import cv2
+        # cv2.imwrite(
+        #     ".pytest_cache/hanning_wight_mask.png", (hann_2d * 255).astype(np.uint8)
+        # )
+
+        # 정규화 및 스케일 조정
+        # hann_2d -= hann_2d.min()
+        # hann_2d /= hann_2d.max()
+        # hann_2d = hann_2d * (max_val - min_val) + min_val
+
+        # (h, w, 1)로 reshape
+        weight = hann_2d[:, :, None]
+
+        return weight.astype(np.float32)
 
     def slice(self, input_vector: np.ndarray) -> list[np.ndarray]:
         patchs = []

pms_nvidia_processor-2.6.1/pms_nvidia_processor/utility/pip.py

@@ -0,0 +1,76 @@
+import subprocess
+import tomllib  # Python 3.11+
+import re
+from pathlib import Path
+from packaging.version import Version
+
+
+def get_pip_list_from_poetry_lock(direcotry: str) -> list[str]:
+    req_txt = subprocess.check_output(
+        [
+            "poetry",
+            "export",
+            "--direcotry",
+            direcotry,
+            "--without-hashes",
+            "-f",
+            "requirements.txt",
+        ],
+        text=True,
+    )
+    pip_packages = [
+        line.strip()
+        for line in req_txt.splitlines()
+        if line.strip() and not line.startswith("#")
+    ]
+    return pip_packages
+
+
+def _convert_caret_version(ver: str) -> str:
+    try:
+        v = Version(ver.strip("^"))
+        if v.major > 0:
+            return f">={v},<{v.major + 1}.0.0"
+        else:
+            return f">={v},<{v.major}.{v.minor + 1}.0"
+    except Exception:
+        return ver  # fallback
+
+
+def _convert_tilde_version(ver: str) -> str:
+    try:
+        v = Version(ver.strip("~"))
+        return f">={v},<{v.major}.{v.minor + 1}.0"
+    except Exception:
+        return ver
+
+
+def _format_dependancy(pkg, ver):
+    if isinstance(ver, str):
+        ver = ver.strip()
+        if ver in {"*", ""} or "*" in ver:
+            return pkg
+        elif ver.startswith("^"):
+            return f"{pkg}{_convert_caret_version(ver)}"
+        elif ver.startswith("~"):
+            return f"{pkg}{_convert_tilde_version(ver)}"
+        else:
+            return f"{pkg}{ver}"
+    else:
+        return pkg  # 복잡한 경우 이름만
+
+
+def get_pip_list_from_pyproject_toml(pyproject_path: str) -> list[str]:
+    with open(pyproject_path, "rb") as f:
+        pyproject = tomllib.load(f)
+
+    poetry_tool = pyproject.get("tool", {}).get("poetry", {})
+    dependencies = poetry_tool.get("dependencies", {})
+
+    # python 자체는 제외
+    deps = [
+        _format_dependancy(pkg, ver)
+        for pkg, ver in dependencies.items()
+        if pkg.lower() != "python"
+    ]
+    return deps

{pms_nvidia_processor-2.6.0 → pms_nvidia_processor-2.6.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pms-nvidia-processor"
-version = "2.6.0"
+version = "2.6.1"
 description = "lib for pms nvidia processors"
 authors = ["HyeongSeok Kim <tiryul@gmail.com>"]
 readme = "README.md"