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

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

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

pms_nvidia_processor-2.4.0/pms_nvidia_processor/__init__.py

@@ -0,0 +1,28 @@
+from .processor.dpir.processor import DPIRProcessor
+from .processor.dru_rbpn_sr_f3.processor import DRURBPNSRF3Processor
+from .processor.dru_rbpn_sr_f5.processor import DRURBPNSRF5Processor
+from .processor.dru_asm_sr_f3.processor import DRUASMSRF3Processor
+from .processor.dru_rbpn_deinter_f3_gloss.processor import (
+    DRURBPNDEINTERF3GLOSSProcessor,
+)
+from .processor.dru_rbpn_deinter_pc_f2.processor import DRURBPNDEINTERPCF2Processor
+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
+
+
+__all__: list[str] = [
+    "DPIRProcessor",
+    "DRURBPNSRF3Processor",
+    "DRURBPNSRF5Processor",
+    "DRUASMSRF3Processor",
+    "DRURBPNDEINTERF3GLOSSProcessor",
+    "DRURBPNDEINTERPCF2Processor",
+    "FISFProcessor",
+    "COLORRESNETPREProcessor",
+    "COLORRESNETPOSTProcessor",
+    "GGProcessor",
+]
+
+__version__ = "2.4.0"

{pms_nvidia_processor-2.2.2/pms_nvidia_processor/processor/dru_asm_sr → pms_nvidia_processor-2.4.0/pms_nvidia_processor/processor/dru_asm_sr_f3}/config.py

@@ -22,6 +22,6 @@ class DRUASMSRF3Config:
     TRT_CONFIG = TRTIOConfig(
         patch_size=PATCH_SIZE,
         upscale_ratio=UPSCALE_RATIO,
-        number_of_input_channels=NUMBER_OF_INPUT_CHANNELS + 1,
+        number_of_input_channels=NUMBER_OF_INPUT_CHANNELS,
         number_of_output_channels=NUMBER_OF_OUTPUT_CHANNELS,
     )

pms_nvidia_processor-2.2.2/pms_nvidia_processor/processor/dru_asm_sr/processor_f3_v2.py → pms_nvidia_processor-2.4.0/pms_nvidia_processor/processor/dru_asm_sr_f3/processor.py

@@ -19,11 +19,11 @@ def _pre_processing(
 ) -> None:
     b = len(batch_input_images)
     # 0.299 * img[..., 2:3, :, :] + 0.587 * img[..., 1:2, :, :] + 0.114 * img[..., 0:1, :, :]
-    e = np.array([0.299, 0.587, 0.114])  # for RGB
-    usm_edge_images = [
-        np.sum(_usm_sharp(img=(img[:, :, 3:6])) * e, axis=-1)
-        for img in batch_input_images
-    ]
+    # e = np.array([0.299, 0.587, 0.114])  # for RGB
+    # usm_edge_images = [
+    #     np.sum(_usm_sharp(img=(img[:, :, 3:6])) * e, axis=-1)
+    #     for img in batch_input_images
+    # ]
 
     for batch_idx in range(b):
         image = batch_input_images[batch_idx]
@@ -31,7 +31,7 @@ def _pre_processing(
         for channel_idx in range(c):
             input_buffer[batch_idx, channel_idx, :h, :w] = image[:, :, channel_idx]
 
-        input_buffer[batch_idx, -1, :h, :w] = usm_edge_images[batch_idx][:h, :w]
+        # input_buffer[batch_idx, -1, :h, :w] = usm_edge_images[batch_idx][:h, :w]
 
 
 def _post_processing(
@@ -62,7 +62,6 @@ class DRUASMSRF3Processor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
             model_path=model_path,
             device=device,
         )
-
         self.config = DRUASMSRF3Config
         self.normalizer = normalizer.NormalizerFactory.create_multi(normalization)
 

pms_nvidia_processor-2.4.0/pms_nvidia_processor/processor/dru_rbpn_deinter_pc_f2/config.py

@@ -0,0 +1,27 @@
+from ..base.config import *
+
+
+class DRURBPNDEINTERPCF2Config:
+    NUMBER_OF_FRAMES = 2
+    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.4.0/pms_nvidia_processor/processor/dru_rbpn_deinter_pc_f2/context.py

@@ -0,0 +1,37 @@
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher, normalizer
+
+
+class DRURBPNDeinterPCF2Context:
+
+    def __init__(
+        self,
+        input_vector: np.ndarray,
+        input_patcher: patcher.Patcher,
+    ) -> None:
+        config = DRURBPNDEINTERPCF2Config
+
+        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.PATCH_SIZE, 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.4.0/pms_nvidia_processor/processor/dru_rbpn_deinter_pc_f2/processor.py

@@ -0,0 +1,161 @@
+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 DRURBPNDEINTERPCF2Processor(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 = DRURBPNDEINTERPCF2Config
+        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 = DRURBPNDeinterPCF2Context(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.4.0/pms_nvidia_processor/processor/dru_rbpn_sr_f3/config.py

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

@@ -0,0 +1,42 @@
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher, normalizer
+
+
+class DRURBPNSRF3Context:
+
+    def __init__(
+        self,
+        input_vector: np.ndarray,
+        input_patcher: patcher.Patcher,
+    ) -> None:
+        config = DRURBPNSRF3Config
+
+        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.4.0/pms_nvidia_processor/processor/dru_rbpn_sr_f5/config.py

@@ -0,0 +1,27 @@
+from ..base.config import *
+
+
+class DRURBPNSRF5Config:
+    NUMBER_OF_FRAMES = 5
+    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 = 32
+
+    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.4.0/pms_nvidia_processor/processor/dru_rbpn_sr_f5/context.py

@@ -0,0 +1,42 @@
+from .config import *
+from ..base.dependency import *
+from ..base.processor import *
+from ...utility import patcher, normalizer
+
+
+class DRURBPNSRF5Context:
+
+    def __init__(
+        self,
+        input_vector: np.ndarray,
+        input_patcher: patcher.Patcher,
+    ) -> None:
+        config = DRURBPNSRF5Config
+
+        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.2.2/pms_nvidia_processor/processor/gg/processor_v2.py → pms_nvidia_processor-2.4.0/pms_nvidia_processor/processor/gg/processor.py

@@ -7,14 +7,16 @@ from ...utility import patcher, normalizer, batch, caster
 
 def _pre_processing(
     batch_input_images: list[np.ndarray],
-    input_buffer: np.ndarray,
+    input_buffer_0: np.ndarray,
+    input_buffer_2: 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]
+            input_buffer_0[batch_idx, channel_idx, :h, :w] = image[:, :, channel_idx]
+    input_buffer_2[:] = np.random.randn(*input_buffer_2.shape).astype(np.float32)[:]
 
 
 def _post_processing(
@@ -50,7 +52,7 @@ class GGProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
             model_path=model_path,
             device=device,
         )
-
+        np.random.seed(index)
         self.config = GGConfig
         self.grain_level = grain_level
         self.normalizer = normalizer.NormalizerFactory.create_multi(normalization)
@@ -63,7 +65,8 @@ class GGProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
         session = self.session
         _pre_processing(
             batch_input_images=batch_input_data,
-            input_buffer=self.input_buffer_0,
+            input_buffer_0=self.input_buffer_0,
+            input_buffer_2=self.input_buffer_2,
         )
         session.run()
         _post_processing(
@@ -131,6 +134,10 @@ class GGProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
                 [self.batch_size, 1, 1, 1],
                 np.float32,
             ),
+            "grain_noise": (
+                [self.batch_size, *trt_config.input_shape],
+                np.float32,
+            ),
             "output_image": (
                 [self.batch_size, *trt_config.output_shape],
                 np.float32,
@@ -151,6 +158,9 @@ class GGProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
             self.io_shapes["grain_level"][0]
         )
         self.input_buffer_1.fill(self.grain_level)
+        self.input_buffer_2 = self.session._input_bindings[2].host_buffer.reshape(
+            self.io_shapes["grain_noise"][0]
+        )
         self.output_buffer = self.session._output_bindings[0].host_buffer.reshape(
             *self.io_shapes["output_image"][0]
         )

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

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

pms_nvidia_processor-2.2.2/pms_nvidia_processor/__init__.py

@@ -1,13 +0,0 @@
-from .processor.dpir.processor_v2 import DPIRProcessor
-from .processor.dru_rbpn_sr.processor_f3_v2 import DRURBPNSRF3Processor
-from .processor.dru_rbpn_sr.processor_f5_v2 import DRURBPNSRF5Processor
-from .processor.dru_asm_sr.processor_f3_v2 import DRUASMSRF3Processor
-from .processor.dru_rbpn_deinter.processor_f3_gloss_v2 import (
-    DRURBPNDEINTERF3GLOSSProcessor,
-)
-from .processor.fisf.processor_v2 import FISFProcessor
-from .processor.color_resnet.processor_pre_v2 import COLORRESNETPREProcessor
-from .processor.color_resnet.processor_post_v2 import COLORRESNETPOSTProcessor
-from .processor.gg.processor_v2 import GGProcessor
-
-__version__ = "2.2.2"