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

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

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

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

@@ -29,4 +29,4 @@ __all__: list[str] = [
     "CDRUF2Processor",
 ]
 
-__version__ = "2.7.0"
+__version__ = "2.8.0"

{pms_nvidia_processor-2.7.0 → pms_nvidia_processor-2.8.0}/pms_nvidia_processor/processor/col_en/config.py

@@ -6,11 +6,11 @@ class ColEnConfig:
     NUMBER_OF_INPUT_CHANNELS: int = 3 * NUMBER_OF_FRAMES
     NUMBER_OF_OUTPUT_CHANNELS: int = 3
     UPSCALE_RATIO: int = 1
-    PATCH_SIZE = 1280
-    MAX_BATCH_SIZE = 1
+    PATCH_SIZE = 512
+    MAX_BATCH_SIZE = 8
     MIN_BATCH_SIZE = 1
-    OPT_BATCH_SIZE = 1
-    INPUT_OVERLAB_LENGTH = 640  # PATCH_SIZE // 2
+    OPT_BATCH_SIZE = 4
+    INPUT_OVERLAB_LENGTH = 256  # PATCH_SIZE // 2
 
     PATCHER_CONFIG = PatcherWIOConfig(
         patch_size=PATCH_SIZE,

{pms_nvidia_processor-2.7.0 → pms_nvidia_processor-2.8.0}/pms_nvidia_processor/processor/col_en/processor.py

@@ -48,6 +48,7 @@ class COLENProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
         )
         self.config = ColEnConfig
         self.normalizer = normalizer.NormalizerFactory.create_multi(normalization)
+        self.patcher = None
 
     async def inference(
         self,
@@ -67,6 +68,52 @@ class COLENProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
             output_image=batch_output_data,
         )
 
+    async def __run(self, input_data: EngineIOData) -> EngineIOData:
+        tact = {}
+        max_batch_size = self.io_shapes["input"][0][0]
+
+        # 여기서 patching
+        tact["0-preprocessing"] = time.time()
+        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)
+        ]
+        tact["0-preprocessing"] = time.time() - tact["0-preprocessing"]
+        tact["1-create context"] = time.time()
+        context = ColEnContext(input_vector, self.patcher)
+        tact["1-create context"] = time.time() - tact["1-create context"]
+        tact["2-batch inference"] = time.time()
+        # 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,
+            )
+        tact["2-batch inference"] = time.time() - tact["2-batch inference"]
+        tact["3-merge"] = time.time()
+        self.patcher.merge(
+            patches=context.output_patches,
+            output_vector=context.output_vector,
+        )
+        tact["3-merge"] = time.time() - tact["3-merge"]
+        tact["4-crop"] = time.time()
+        context.output_vector = patcher.crop_vector(
+            context.output_vector, context.input_pathcer._input_overlap_length
+        )
+        tact["4-crop"] = time.time() - tact["4-crop"]
+        tact["5-denormalize"] = time.time()
+        self.normalizer.denormalize(context.output_vector, 2)
+        tact["5-denormalize"] = time.time() - tact["5-denormalize"]
+        print("TACT: " + " ".join([f"{k}: {v*1000:.0f}ms" for k, v in tact.items()]))
+        return EngineIOData(
+            frame_id=input_data.frame_id,
+            frame=self.caster.uncast(context.output_vector),
+        )
+
     async def _run(self, input_data: EngineIOData) -> EngineIOData:
         max_batch_size = self.io_shapes["input"][0][0]
 
@@ -123,13 +170,13 @@ class COLENProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
                 self.config.NUMBER_OF_OUTPUT_CHANNELS,
             )
         )
-        self.patcher = patcher.PatcherW(
+        temp_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))
+        n_patches = len(temp_patcher.slice(input_vector=padded_input_image))
 
         # set io shape
         self.batch_size = min(n_patches, self.config.MAX_BATCH_SIZE)
@@ -149,6 +196,13 @@ class COLENProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
             required_batch_size=self.batch_size,
             io_shape=self.io_shapes,
         )
+        self.patcher = patcher.PatcherW(
+            **patcher_config.build_patcher_params(
+                input_vector=padded_input_image,
+                output_vector=output_image,
+            )
+        )
+        self.session._trt_proxy._context.push()
 
         # set io buffer
         self.input_buffer = self.session._input_bindings[0].host_buffer.reshape(
@@ -166,3 +220,10 @@ class COLENProcessor(NVIDIAProcessorBase[EngineIOData, EngineIOData]):
 
     def _get_concurrency(self) -> int:
         return self._concurrency
+
+    # def __del__(self):
+    #     if self._session:
+    #         self._session._trt_proxy._context.pop()
+    #         del self._session
+    #     if self.patcher:
+    #         del self.patcher

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

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