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

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

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

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

@@ -10,7 +10,8 @@ from .processor.fisf.processor import FISFProcessor
 from .processor.color_resnet.processor_pre import COLORRESNETPREProcessor
 from .processor.color_resnet.processor_post import COLORRESNETPOSTProcessor
 from .processor.gg.processor import GGProcessor
-from .processor.col_en.processor import ColEnProcessor
+from .processor.col_en.processor import COLENProcessor
+from .processor.cdru_f2.processor import CDRUF2Processor
 
 
 __all__: list[str] = [
@@ -24,7 +25,8 @@ __all__: list[str] = [
     "COLORRESNETPREProcessor",
     "COLORRESNETPOSTProcessor",
     "GGProcessor",
-    "ColEnProcessor",
+    "COLENProcessor",
+    "CDRUF2Processor",
 ]
 
-__version__ = "2.5.0"
+__version__ = "2.6.0"

pms_nvidia_processor-2.6.0/pms_nvidia_processor/processor/cdru_f2/config.py

@@ -0,0 +1,27 @@
+from ..base.config import *
+
+
+class CDRUF2Config:
+    NUMBER_OF_FRAMES = 2
+    NUMBER_OF_INPUT_CHANNELS: int = 3 * NUMBER_OF_FRAMES
+    NUMBER_OF_OUTPUT_CHANNELS: int = 3
+    UPSCALE_RATIO: int = 2
+    PATCH_SIZE = 512
+    MAX_BATCH_SIZE = 8
+    MIN_BATCH_SIZE = 1
+    OPT_BATCH_SIZE = MAX_BATCH_SIZE // 2
+    INPUT_OVERLAB_LENGTH = 16
+
+    PATCHER_CONFIG = PatcherIOConfig(
+        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.6.0/pms_nvidia_processor/processor/cdru_f2/context.py

@@ -0,0 +1,42 @@
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher, normalizer
+
+
+class CDRUF2Context:
+
+    def __init__(
+        self,
+        input_vector: np.ndarray,
+        input_patcher: patcher.Patcher,
+    ) -> None:
+        config = CDRUF2Config
+
+        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(
+            (
+                input_vector.shape[0] * config.UPSCALE_RATIO,
+                input_vector.shape[1] * config.UPSCALE_RATIO,
+                config.NUMBER_OF_OUTPUT_CHANNELS,
+            ),
+            np.float32,
+        )

pms_nvidia_processor-2.6.0/pms_nvidia_processor/processor/cdru_f2/processor.py

@@ -0,0 +1,160 @@
+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
+        input_buffer[batch_idx, 0:c, :h, :w] = image.transpose(2, 0, 1)
+
+
+def _post_processing(
+    output_buffer: np.ndarray,  # BxCxHxW
+    output_image: np.ndarray,  # BxHxWxC
+) -> None:
+    b, h, w, c = output_image.shape
+    x = output_buffer[:b, :, :h, :w]
+    for i in range(3):
+        np.copyto(
+            src=x[:b, i, :h, :w],
+            dst=output_image[:, :, :, i],
+            casting="safe",
+        )
+
+
+@register
+class CDRUF2Processor(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 = CDRUF2Config
+        self.normalizer = normalizer.NormalizerFactory.create_multi(normalization)
+
+    async def inference(
+        self,
+        batch_input_data: list[np.ndarray],
+        batch_output_data: np.ndarray,
+    ):
+        session = self.session
+        _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]
+        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 = CDRUF2Context(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(
+            output_vector=context.output_vector,
+            patches=context.output_patches,
+        )
+
+        self.normalizer.denormalize(context.output_vector, 2)
+        return EngineIOData(
+            frame_id=input_data.frame_id,
+            frame=self.caster.uncast(context.output_vector),
+        )
+
+    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,
+            )
+        )
+        self.patcher = patcher.Patcher(
+            **patcher_config.build_patcher_params(
+                input_vector=padded_input_image,
+                output_vector=output_image,
+            )
+        )
+        n_patches = len(self.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,
+        )
+
+        # 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]
+        )
+
+        return True
+
+    def _get_live(self) -> bool:
+        return True
+
+    def _get_concurrency(self) -> int:
+        return self._concurrency

{pms_nvidia_processor-2.5.0 → pms_nvidia_processor-2.6.0}/pms_nvidia_processor/processor/col_en/processor.py

@@ -29,7 +29,7 @@ def _post_processing(
 
 
 @register
-class ColEnProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
+class COLENProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
 
     def __init__(
         self,

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

@@ -94,7 +94,8 @@ class PatcherW:
         output_vector: np.ndarray,
         scale: int = 1,
     ):
-        w_vector = np.ones((*output_vector.shape[:-1], 1), np.float32) * (1e-8)
+        # w_vector = np.ones((*output_vector.shape[:-1], 1), np.float32) * (1e-8)
+        w_vector = np.zeros((*output_vector.shape[:-1], 1), np.float32)
         for (ys, xs), patch in zip(self._patch_pos_list, patches, strict=True):
             ys_out = slice(ys.start * scale, ys.stop * scale)
             xs_out = slice(xs.start * scale, xs.stop * scale)

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

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