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

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

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

{pms_nvidia_processor-2.2.2 → pms_nvidia_processor-2.3.0}/pms_nvidia_processor/__init__.py

@@ -10,4 +10,4 @@ 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"
+__version__ = "2.3.0"

{pms_nvidia_processor-2.2.2 → pms_nvidia_processor-2.3.0}/pms_nvidia_processor/processor/dru_asm_sr/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-2.3.0}/pms_nvidia_processor/processor/dru_asm_sr/processor_f3_v2.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.2.2 → pms_nvidia_processor-2.3.0}/pms_nvidia_processor/processor/gg/processor_v2.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.3.0}/pyproject.toml

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