pms-nvidia-processor 2.8.0__tar.gz → 2.9.0__tar.gz

{pms_nvidia_processor-2.8.0 → pms_nvidia_processor-2.9.0}/PKG-INFO

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

{pms_nvidia_processor-2.8.0 → pms_nvidia_processor-2.9.0}/pms_nvidia_processor/__init__.py

@@ -12,6 +12,7 @@ from .processor.color_resnet.processor_post import COLORRESNETPOSTProcessor
 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
 
 
 __all__: list[str] = [
@@ -27,6 +28,7 @@ __all__: list[str] = [
     "GGProcessor",
     "COLENProcessor",
     "CDRUF2Processor",
+    "DOF3Processor",
 ]
 
-__version__ = "2.8.0"
+__version__ = "2.9.0"

pms_nvidia_processor-2.9.0/pms_nvidia_processor/processor/do_f3/config.py

@@ -0,0 +1,27 @@
+from ..base.config import *
+
+
+class DOF3Config:
+    NUMBER_OF_FRAMES = 3
+    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 = 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.9.0/pms_nvidia_processor/processor/do_f3/context.py

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

pms_nvidia_processor-2.9.0/pms_nvidia_processor/processor/do_f3/processor.py

@@ -0,0 +1,165 @@
+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 DOF3Processor(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 = DOF3Config
+        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
+        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 = DOF3Context(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.input_buffer.fill(1.0 / 255.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.8.0 → pms_nvidia_processor-2.9.0}/pms_nvidia_processor/utility/builder.py

@@ -20,6 +20,7 @@ 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
+from ..processor.do_f3.config import DOF3Config
 
 __MODEL_CONFIG_MAP = {
     "DPIR": DPIRConfig,
@@ -34,6 +35,8 @@ __MODEL_CONFIG_MAP = {
     "GG": GGConfig,
     "COL-EN": ColEnConfig,
     "CDRU-F2": CDRUF2Config,
+    "DO-L-F3": DOF3Config,
+    "DO-H-F3": DOF3Config,
 }
 
 

{pms_nvidia_processor-2.8.0 → pms_nvidia_processor-2.9.0}/pyproject.toml

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