Issue 95: Added PixelFormat class, replaced usages of color formats

juturna/meta/_constants.py

@@ -29,4 +29,4 @@ def get_constant_var(var_name: str):
 JUTURNA_THREAD_JOIN_TIMEOUT = get_constant_var('JUTURNA_THREAD_JOIN_TIMEOUT')
 JUTURNA_MAX_QUEUE_SIZE = get_constant_var('JUTURNA_MAX_QUEUE_SIZE')
 JUTURNA_ENV_VAR_PREFIX = get_constant_var('JUTURNA_ENV_VAR_PREFIX')
-JUTURNA_TELEMETRY_BATCH_SIZE = get_constant_var('JUTURNA_TELEMETRY_BATCH_SIZE')
+JUTURNA_TELEMETRY_BATCH_SIZE = get_constant_var('JUTURNA_TELEMETRY_BATCH_SIZE')

juturna/names/__init__.py

@@ -2,6 +2,7 @@
 from juturna.names._component_status import ComponentStatus
 from juturna.names._pipeline_status import PipelineStatus
 from juturna.names._service_status import ServiceStatus
+from juturna.names._pixel_format import PixelFormat
 
 
 __all__ = ['ComponentStatus', 'PipelineStatus', 'ServiceStatus']

juturna/names/_pixel_format.py

@@ -0,0 +1,27 @@
+class PixelFormat:
+    @staticmethod
+    def _normalize_format(fmt: str) -> str:
+        return fmt.strip().lower()
+
+    __slots__ = ['RGB', 'BGR', 'RGB24', 'BGR24', 'RGBA', 'ARGB', 'BGRA', 'ABGR', 'YUV420P', 'YUV444P', 'YUYV422', 'GRAY', 'GRAY8', 'PAL8', 'GBRP']
+
+    RGB = 'rgb'
+    BGR = 'bgr'
+
+    RGB24 = 'rgb24'
+    BGR24 = 'bgr24'
+
+    RGBA = 'rgba'
+    ARGB = 'argb'
+    BGRA = 'bgra'
+    ABGR = 'abgr'
+
+    YUV420P = 'yuv420p'
+    YUV444P = 'yuv444p'
+    YUYV422 = 'yuyv422'
+
+    GRAY = 'gray'
+    GRAY8 = 'gray8'
+
+    PAL8 = 'pal8'
+    GBRP = 'gbrp'

juturna/nodes/source/_video_file/video_file.py

@@ -16,8 +16,11 @@
 
 from juturna.components import Node
 from juturna.components import Message
+
 from juturna.payloads import BytesPayload, ImagePayload
 
+from juturna.names import PixelFormat
+
 
 class VideoFile(Node[BytesPayload, ImagePayload]):
     """Read video file and steam it locally"""
@@ -140,7 +143,7 @@ def update(self, message: Message[BytesPayload]):
                     image=full_frame,
                     width=self._width,
                     height=self._height,
-                    pixel_format='rgb24',
+                    pixel_format=PixelFormat.RGB24,
                     timestamp=time.time(),
                 ),
             )

juturna/nodes/source/_video_rtp/video_rtp.py

@@ -12,6 +12,8 @@
 
 from juturna.payloads import BytesPayload, ImagePayload
 
+from juturna.names import PixelFormat
+
 
 class VideoRTP(Node[BytesPayload, ImagePayload]):
     """Source node for video streaming"""
@@ -140,7 +142,7 @@ def update(self, message: Message[BytesPayload]):
                     image=full_frame,
                     width=full_frame.shape[0],
                     height=full_frame.shape[1],
-                    pixel_format='rgb24',
+                    pixel_format=PixelFormat.RGB24,
                     timestamp=time.time(),
                 ),
             )

juturna/nodes/source/_video_rtp_av/video_rtp_av.py

@@ -17,6 +17,7 @@
 from juturna.components import Message
 from juturna.components import _resource_broker as rb
 from juturna.payloads import BytesPayload, ImagePayload
+from juturna.names import PixelFormat
 
 
 class VideoRtpAv(Node[BytesPayload, ImagePayload]):
@@ -141,7 +142,7 @@ def _generate_chunks(self):
                     if self._stop_event.is_set():
                         break
 
-                    full_frame = frame.to_ndarray(format='rgb24')
+                    full_frame = frame.to_ndarray(format=PixelFormat.RGB24)
 
                     to_send = Message[ImagePayload](
                         creator=self.name,
@@ -150,7 +151,7 @@ def _generate_chunks(self):
                             image=full_frame,
                             width=full_frame.shape[1],
                             height=full_frame.shape[0],
-                            pixel_format='rgb24',
+                            pixel_format=PixelFormat.RGB24,
                             timestamp=time.time(),
                         ),
                     )

plugins/nodes/proc/_yolo_detector/yolo_detector.py

@@ -17,6 +17,8 @@
 
 from juturna.payloads._payloads import ImagePayload
 
+from juturna.names import PixelFormat
+
 
 class YoloDetector(Node[ImagePayload, ImagePayload]):
     """Node implementation class"""
@@ -89,7 +91,7 @@ def update(self, message: Message[ImagePayload]):
         assert self._model is not None
 
         image = message.payload.image
-        image_format = message.payload.pixel_format
+        image_format = PixelFormat._normalize_format(message.payload.pixel_format)
         meta = dict(message.meta)
 
         normalized_image = None
@@ -103,13 +105,13 @@ def update(self, message: Message[ImagePayload]):
         )
 
         with to_send.timeit(self.name + '_image_preprocessing_numpy'):
-            if image.shape[2] == 4 and image_format == 'RGB':
+            if image.shape[2] == 4 and image_format == PixelFormat.RGB:
                 normalized_image = image[
                     :, :, 2::-1
                 ]  # remove alpha and convert RGB→BGR
             elif image.shape[2] == 4:
                 normalized_image = image[:, :, :3]  # remove alpha only
-            elif image_format == 'RGB':
+            elif image_format == PixelFormat.RGB:
                 normalized_image = image[:, :, ::-1]  # only RGB→BGR
             else:
                 normalized_image = image  # no modification
@@ -129,7 +131,7 @@ def update(self, message: Message[ImagePayload]):
 
         with to_send.timeit(self.name + '_postprocessing'):
             annotated = results[0].plot() if self._plot else image
-            pixel_format = 'BGR' if self._plot else image_format
+            pixel_format = PixelFormat.BGR if self._plot else image_format
 
         to_send.payload = ImagePayload(
             image=annotated,

plugins/nodes/source/_image_loader/image_loader.py

@@ -30,6 +30,8 @@
 from juturna.payloads import ObjectPayload
 from juturna.payloads import ImagePayload
 
+from juturna.names import PixelFormat
+
 
 class ImageLoader(Node[ObjectPayload, ImagePayload]):
     """Node implementation class"""
@@ -126,7 +128,7 @@ def update(self, message: Message[ObjectPayload]):
             image = Image.open(message.payload['src_path'])
 
             if self._convert_rgb:
-                image = image.convert('RGB')
+                image = image.convert(PixelFormat.RGB)
 
             image.load()
         except UnidentifiedImageError: