pms-nvidia-processor 2.10.2__tar.gz → 2.11.0__tar.gz

{pms_nvidia_processor-2.10.2 → pms_nvidia_processor-2.11.0}/PKG-INFO

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

{pms_nvidia_processor-2.10.2 → pms_nvidia_processor-2.11.0}/pms_nvidia_processor/__init__.py

@@ -13,6 +13,7 @@ from .processor.gg.processor import GGProcessor
 from .processor.col_en.processor import COLENProcessor
 from .processor.cdru_f2.processor import CDRUF2Processor
 from .processor.do_f3.processor import DOF3Processor
+from .processor.tc.processor import TCProcessor
 
 
 __all__: list[str] = [
@@ -29,6 +30,7 @@ __all__: list[str] = [
     "COLENProcessor",
     "CDRUF2Processor",
     "DOF3Processor",
+    "TCProcessor",
 ]
 
-__version__ = "2.10.2"
+__version__ = "2.11.0"

{pms_nvidia_processor-2.10.2 → pms_nvidia_processor-2.11.0}/pms_nvidia_processor/processor/base/processor.py

@@ -68,8 +68,9 @@ class NVIDIAProcessorBase(IEngineProcessor[InputTypeT, OutputTypeT], metaclass=A
         self.session.run()
 
     def __del__(self):
-        if self._session:
-            del self._session
+        session = getattr(self, "_session", None)
+        if session:
+            del session
 
     @property
     def session(self) -> TRT.TRTSession:

pms_nvidia_processor-2.11.0/pms_nvidia_processor/processor/tc/config.py

@@ -0,0 +1,27 @@
+from ..base.config import *
+
+
+class TCConfig:
+    NUMBER_OF_FRAMES = 1
+    NUMBER_OF_INPUT_CHANNELS: int = 3 * NUMBER_OF_FRAMES
+    NUMBER_OF_OUTPUT_CHANNELS: int = 3
+    UPSCALE_RATIO: int = 1
+    PATCH_SIZE = 512
+    MAX_BATCH_SIZE = 8
+    MIN_BATCH_SIZE = 1
+    OPT_BATCH_SIZE = 4
+    INPUT_OVERLAB_LENGTH = 256  # PATCH_SIZE // 2
+
+    PATCHER_CONFIG = PatcherWIOConfig(
+        patch_size=PATCH_SIZE,
+        upscale_ratio=UPSCALE_RATIO,
+        number_of_input_channels=NUMBER_OF_INPUT_CHANNELS,
+        number_of_output_channels=NUMBER_OF_OUTPUT_CHANNELS,
+        input_overlab_length=INPUT_OVERLAB_LENGTH,
+    )
+    TRT_CONFIG = TRTIOConfig(
+        patch_size=PATCH_SIZE,
+        upscale_ratio=UPSCALE_RATIO,
+        number_of_input_channels=NUMBER_OF_INPUT_CHANNELS,
+        number_of_output_channels=NUMBER_OF_OUTPUT_CHANNELS,
+    )

pms_nvidia_processor-2.11.0/pms_nvidia_processor/processor/tc/context.py

@@ -0,0 +1,38 @@
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher, normalizer
+
+
+class TCContext:
+
+    def __init__(
+        self,
+        input_vector: np.ndarray,
+        input_patcher: patcher.PatcherW,
+    ) -> None:
+        config = TCConfig()
+
+        self.input_vector = input_vector
+        self.input_pathcer = input_patcher
+        self.padded_input_image = patcher.pad_vector(
+            vector=self.input_vector,
+            overlap_length=config.PATCHER_CONFIG.input_overlab_length,
+            mode="edge",
+        )
+        self.input_patches: list[np.ndarray] = self.input_pathcer.slice(
+            self.padded_input_image
+        )
+        self.output_patches: np.ndarray = np.zeros(
+            (
+                len(self.input_patches),
+                config.PATCH_SIZE * config.UPSCALE_RATIO,
+                config.PATCH_SIZE * config.UPSCALE_RATIO,
+                3,
+            ),
+            np.float32,
+        )
+        self.output_vector: np.ndarray = np.zeros(
+            (*self.padded_input_image.shape[:-1], 3),
+            np.float32,
+        )

pms_nvidia_processor-2.11.0/pms_nvidia_processor/processor/tc/processor.py

@@ -0,0 +1,186 @@
+from .config import *
+from .context import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher, normalizer, batch, caster
+
+
+def _pre_processing(
+    batch_input_images: list[np.ndarray],
+    input_buffer: np.ndarray,
+) -> None:
+    b = len(batch_input_images)
+
+    for batch_idx in range(b):
+        image = batch_input_images[batch_idx]
+        h, w, c = image.shape
+        for channel_idx in range(c):
+            input_buffer[batch_idx, channel_idx, :h, :w] = image[:, :, channel_idx]
+
+
+def _post_processing(
+    output_buffer: np.ndarray,  # BxCxHxW
+    output_image: np.ndarray,  # BxHxWxC
+) -> None:
+
+    b, h, w, c = output_image.shape
+    for i in range(3):
+        np.copyto(src=output_buffer[:b, i, :h, :w], dst=output_image[:, :, :, i])
+
+
+@register
+class TCProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
+
+    def __init__(
+        self,
+        concurrency: int,
+        index: int,
+        model_path: str,
+        device: Literal["auto"] | int = "auto",
+        normalization: str = "color-basic",
+    ):
+        # super
+        super().__init__(
+            concurrency=concurrency,
+            index=index,
+            model_path=model_path,
+            device=device,
+        )
+        self.config = TCConfig
+        self.normalizer = normalizer.NormalizerFactory.create_multi(normalization)
+        self.patcher = None
+
+    async def inference(
+        self,
+        batch_input_data: list[np.ndarray],
+        batch_output_data: np.ndarray,
+    ):
+        session = self.session
+        batch = len(batch_input_data)
+
+        _pre_processing(
+            batch_input_images=batch_input_data,
+            input_buffer=self.input_buffer,
+        )
+        session.run()
+        _post_processing(
+            output_buffer=self.output_buffer,
+            output_image=batch_output_data,
+        )
+
+    async def _run(self, input_data: EngineIOData) -> EngineIOData:
+        max_batch_size = self.io_shapes["input"][0][0]
+
+        # 여기서 patching
+        input_vector = self.caster.cast(input_data.frame)
+        [
+            self.normalizer.normalize(input_vector[:, :, i * 3 : (i + 1) * 3], 2)
+            for i in range(self.config.NUMBER_OF_FRAMES)
+        ]
+        context = TCContext(input_vector, self.patcher)
+
+        # batch inference
+        for batch_input_patches, batch_output_patches in zip(
+            batch.BatchTool.batch_list(context.input_patches, max_batch_size),
+            batch.BatchTool.batch_vector(context.output_patches, max_batch_size),
+        ):
+            await self.inference(
+                batch_input_data=batch_input_patches,
+                batch_output_data=batch_output_patches,
+            )
+
+        self.patcher.merge(
+            patches=context.output_patches,
+            output_vector=context.output_vector,
+        )
+        context.output_vector = patcher.crop_vector(
+            context.output_vector, context.input_pathcer._input_overlap_length
+        )
+
+        self.normalizer.denormalize(context.output_vector, 2)
+        result = EngineIOData(
+            frame_id=input_data.frame_id,
+            frame=self.caster.uncast(context.output_vector),
+        )
+        return result
+
+    def _ready_processor(self) -> bool:
+        return True
+
+    def _bind_io(self, input_data: EngineIOData):
+        self.caster = caster.Caster(input_data.frame)
+
+        patcher_config = self.config.PATCHER_CONFIG
+        trt_config = self.config.TRT_CONFIG
+
+        input_image: np.ndarray = input_data.frame  # type: ignore
+        padded_input_image = patcher.pad_vector(
+            input_image,
+            overlap_length=patcher_config.input_overlab_length,
+        )
+        output_image: np.ndarray = np.zeros(
+            (
+                input_image.shape[0] * self.config.UPSCALE_RATIO,
+                input_image.shape[1] * self.config.UPSCALE_RATIO,
+                self.config.NUMBER_OF_OUTPUT_CHANNELS,
+            )
+        )
+        temp_patcher = patcher.PatcherW(
+            **patcher_config.build_patcher_params(
+                input_vector=padded_input_image,
+                output_vector=output_image,
+            )
+        )
+        n_patches = len(temp_patcher.slice(input_vector=padded_input_image))
+
+        # set io shape
+        self.batch_size = min(n_patches, self.config.MAX_BATCH_SIZE)
+        self.io_shapes = {
+            "input": (
+                [self.batch_size, *trt_config.input_shape],
+                np.float32,
+            ),
+            "output": (
+                [self.batch_size, *trt_config.output_shape],
+                np.float32,
+            ),
+        }
+
+        # init trt engine
+        self.initialize_trt_session(
+            required_batch_size=self.batch_size,
+            io_shape=self.io_shapes,
+        )
+        self.patcher = patcher.PatcherW(
+            **patcher_config.build_patcher_params(
+                input_vector=padded_input_image,
+                output_vector=output_image,
+            )
+        )
+
+        # self.session._trt_proxy._context.push()
+
+        # set io buffer
+        self.input_buffer = self.session._input_bindings[0].host_buffer.reshape(
+            self.io_shapes["input"][0]
+        )
+        self.output_buffer = self.session._output_bindings[0].host_buffer.reshape(
+            *self.io_shapes["output"][0]
+        )
+        self.input_buffer.fill(0)
+        self.output_buffer.fill(0)
+
+        return True
+
+    def _get_live(self) -> bool:
+        return True
+
+    def _get_concurrency(self) -> int:
+        return self._concurrency
+
+    # def __del__(self):
+    #     if self._session:
+    #         self._session._trt_proxy._context.pop()
+    #         del self._session
+    #     if self.patcher:
+    #         del self.patcher

{pms_nvidia_processor-2.10.2 → pms_nvidia_processor-2.11.0}/pms_nvidia_processor/utility/builder.py

@@ -21,6 +21,7 @@ from ..processor.gg.config import GGConfig
 from ..processor.col_en.config import ColEnConfig
 from ..processor.cdru_f2.config import CDRUF2Config
 from ..processor.do_f3.config import DOF3Config
+from ..processor.tc.config import TCConfig
 
 __MODEL_CONFIG_MAP = {
     "DPIR": DPIRConfig,
@@ -37,6 +38,7 @@ __MODEL_CONFIG_MAP = {
     "CDRU-F2": CDRUF2Config,
     "DO-L-F3": DOF3Config,
     "DO-H-F3": DOF3Config,
+    "TC": TCConfig,
 }
 
 
@@ -106,6 +108,42 @@ def create_build_params(
                 },
             }
         )
+    elif config is TCConfig:
+        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,
+                        ],
+                    },
+                    "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(
             {

{pms_nvidia_processor-2.10.2 → pms_nvidia_processor-2.11.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pms-nvidia-processor"
-version = "2.10.2"
+version = "2.11.0"
 description = "lib for pms nvidia processors"
 authors = ["HyeongSeok Kim <tiryul@gmail.com>"]
 readme = "README.md"
@@ -20,7 +20,7 @@ pytest-order = "^1.2.0"
 pms-model-manager = "^1.3.1"
 tensorrt = "^10.13"
 cuda-python = "^13.0"
-tensorrt-cu13-libs = "^10.13"
+tensorrt-cu13-libs = "^10.13.2"
 tensorrt-cu13-bindings = "^10.13"