pms-nvidia-processor 2.4.0__tar.gz → 2.5.1__tar.gz

{pms_nvidia_processor-2.4.0 → pms_nvidia_processor-2.5.1}/PKG-INFO

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

{pms_nvidia_processor-2.4.0 → pms_nvidia_processor-2.5.1}/pms_nvidia_processor/__init__.py

@@ -10,6 +10,7 @@ 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
 
 
 __all__: list[str] = [
@@ -23,6 +24,7 @@ __all__: list[str] = [
     "COLORRESNETPREProcessor",
     "COLORRESNETPOSTProcessor",
     "GGProcessor",
+    "COLENProcessor",
 ]
 
-__version__ = "2.4.0"
+__version__ = "2.5.1"

{pms_nvidia_processor-2.4.0 → pms_nvidia_processor-2.5.1}/pms_nvidia_processor/processor/base/config.py

@@ -67,3 +67,16 @@ class TRTIOConfig:
             self.patch_size * self.upscale_ratio,
             self.patch_size * self.upscale_ratio,
         )
+
+
+@dataclass
+class PatcherWIOConfig(PatcherIOConfig):
+
+    @property
+    def output_shape(self) -> tuple[int, int, int]:
+        output_patch_size = self.patch_size * self.upscale_ratio
+        return (
+            output_patch_size,
+            output_patch_size,
+            self.number_of_output_channels,
+        )

pms_nvidia_processor-2.5.1/pms_nvidia_processor/processor/col_en/config.py

@@ -0,0 +1,27 @@
+from ..base.config import *
+
+
+class ColEnConfig:
+    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 = 1536
+    MAX_BATCH_SIZE = 1
+    MIN_BATCH_SIZE = 1
+    OPT_BATCH_SIZE = 1
+    INPUT_OVERLAB_LENGTH = 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.5.1/pms_nvidia_processor/processor/col_en/context.py

@@ -0,0 +1,38 @@
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher, normalizer
+
+
+class ColEnContext:
+
+    def __init__(
+        self,
+        input_vector: np.ndarray,
+        input_patcher: patcher.PatcherW,
+    ) -> None:
+        config = ColEnConfig()
+
+        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.5.1/pms_nvidia_processor/processor/col_en/processor.py

@@ -0,0 +1,168 @@
+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 COLENProcessor(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 = ColEnConfig
+        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 = ColEnContext(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)
+        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.PatcherW(
+            **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(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.5.1/pms_nvidia_processor/utility/patcher/__init__.py

@@ -0,0 +1,2 @@
+from ._patcher import Patcher, pad_vector
+from ._patcher_w import PatcherW, crop_vector

{pms_nvidia_processor-2.4.0 → pms_nvidia_processor-2.5.1}/pms_nvidia_processor/utility/patcher/_patcher.py

@@ -80,3 +80,34 @@ class Patcher:
             patches=patches,
             overlab_length=self._output_overlap_length,
         )
+
+    def merge_w(
+        self, output_vector: np.ndarray, patches: list[np.ndarray] | np.ndarray
+    ):
+        import cv2
+
+        n = 0
+
+        pos_collection = self._output_pos_collection
+        vector = output_vector
+        weight_vector = np.ones((*output_vector.shape[:-1], 1), np.float32) * (1e-8)
+        overlab_length = self._output_overlap_length
+        print(f"pos_collection.size: {pos_collection.size}")
+        print(f"pos_collection.patch_pos_list: {pos_collection.patch_pos_list}")
+        for pos, patch in zip(pos_collection, patches, strict=True):  # inplace copy
+            h, w, c = patch.shape
+            np.add(
+                vector[pos.y.range, pos.x.range],
+                patch[
+                    overlab_length : overlab_length + pos.y.dp,
+                    overlab_length : overlab_length + pos.x.dp,
+                ],
+                out=vector[pos.y.range, pos.x.range],
+            )
+            weight_vector[pos.y.range, pos.x.range] += 1.0
+            cv2.imwrite(
+                f"/home/hskim/pms-nvidia-processor/.pytest_cache/2025-04-29 color test/{n:04d}.jpg",
+                np.clip(vector * 255, 0, 255).astype(np.uint8),
+            )
+            n += 1
+        # np.divide(vector, weight_vector, out=vector)

pms_nvidia_processor-2.5.1/pms_nvidia_processor/utility/patcher/_patcher_w.py

@@ -0,0 +1,109 @@
+from typing import Literal
+import numpy as np
+
+
+def crop_vector(
+    vector: np.ndarray,
+    overlap_length: int,
+) -> np.ndarray:
+    return vector[overlap_length:-overlap_length, overlap_length:-overlap_length].copy()
+
+
+class PatcherW:
+
+    def __init__(
+        self,
+        input_vector_shape: tuple[int, int, int],
+        input_patch_shape: tuple[int, int, int],
+        input_overlap_length: int,
+        output_vector_shape: tuple[int, int, int],
+        output_patch_shape: tuple[int, int, int],
+        output_overlap_length: int,
+    ) -> None:
+        assert input_overlap_length > -1, "assert input_overlap_length > -1"
+        assert all(
+            [e > 0 for e in input_patch_shape]
+        ), "assert all([e > 0 for e in input_patch_shape])"
+        assert (
+            len(input_patch_shape) == 3
+        ), "assert len(input_patch_shape) == 3"  # only allow image-like vector
+        self._input_vector_shape = input_vector_shape
+        self._input_patch_shape = input_patch_shape
+        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)
+
+    def _get_patch_pos_list(self):
+        input_vector_width = self._input_vector_shape[1]
+        input_vector_height = self._input_vector_shape[0]
+        input_patch_width = self._input_patch_shape[1]
+        input_patch_height = self._input_patch_shape[0]
+        input_overlap_length = self._input_overlap_length
+        patchs: list[tuple[slice, slice]] = []
+        for y0 in range(
+            0, input_vector_height, input_patch_height - input_overlap_length
+        ):
+            y1 = min(y0 + input_patch_height, input_vector_height)
+            for x0 in range(
+                0, input_vector_width, input_patch_width - input_overlap_length
+            ):
+                x1 = min(x0 + input_patch_width, input_vector_width)
+                patchs.append((slice(y0, y1), slice(x0, x1)))
+        return patchs
+
+    def _make_weight_mask(
+        self,
+        shape: tuple[int, int, int],
+        min_val=0.0,
+        max_val=1.0,
+        sigma_scale=0.25,
+    ):
+        """
+        h, w: 패치 높이/너비
+        sigma_scale: 중심에서 분포하는 정도 (0.25이면 가우시안 표준편차가 크기 대비 1/4)
+        """
+        import cv2
+
+        h, w, c = shape
+        y = np.linspace(-1, 1, h)
+        x = np.linspace(-1, 1, w)
+        yy, xx = np.meshgrid(y, x, indexing="ij")
+        d2 = xx**2 + yy**2
+        sigma2 = sigma_scale**2
+        weights = np.exp(-d2 / (2 * sigma2))
+        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)
+
+    def slice(self, input_vector: np.ndarray) -> list[np.ndarray]:
+        patchs = []
+        for ys, xs in self._get_patch_pos_list():
+            patchs.append(input_vector[ys, xs])
+        return patchs
+
+    def merge(
+        self,
+        patches: list[np.ndarray] | np.ndarray,
+        output_vector: np.ndarray,
+        scale: int = 1,
+    ):
+        w_vector = np.ones((*output_vector.shape[:-1], 1), np.float32) * (1e-8)
+        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)
+            target_vector = output_vector[ys_out, xs_out]
+            target_patch = patch[: target_vector.shape[0], : target_vector.shape[1]]
+            target_weight_mask = self._weight_mask[
+                : target_patch.shape[0], : target_patch.shape[1]
+            ]
+            np.multiply(target_patch, target_weight_mask, out=target_patch)
+            np.add(target_vector, target_patch, out=target_vector)
+            w_vector[ys_out, xs_out] += target_weight_mask
+        output_vector /= w_vector

{pms_nvidia_processor-2.4.0 → pms_nvidia_processor-2.5.1}/pyproject.toml

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

pms_nvidia_processor-2.4.0/pms_nvidia_processor/utility/patcher/__init__.py

@@ -1 +0,0 @@
-from ._patcher import Patcher, pad_vector