Fix in check of chunk type and code restructure for better readability

ifm · May 2, 2024 · f224783 · f224783
1 parent c87a677
commit f224783
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 16 deletions.
diff --git a/source/device/image_client.py b/source/device/image_client.py
@@ -1,12 +1,10 @@
 from __future__ import (absolute_import, division, print_function)
 from builtins import *
 from source.pcic import O2x5xxPCICDevice
-from source.static.formats import serialization_format
+from source.static.formats import serialization_format, ChunkType
 from source.static.configs import images_config
 import struct
 import io
-import enum
-import array
 import matplotlib.pyplot as plt
 import matplotlib.image as mpimg
 import numpy as np
@@ -15,11 +13,6 @@
 SOCKET_TIMEOUT = 10
 
 
-class ChunkType(enum.IntEnum):
-	MONOCHROME_2D_8BIT = 251
-	JPEG_IMAGE = 260
-
-
 class ImageClient(O2x5xxPCICDevice):
 	def __init__(self, address, port, timeout=SOCKET_TIMEOUT):
 		super(ImageClient, self).__init__(address=address, port=port, timeout=timeout)
@@ -93,10 +86,18 @@ def _deserialize_image_chunk(data):
 			# append image
 			image_hex = data[header['HEADER_SIZE']:header['CHUNK_SIZE']]
 			chunk_type = int(header['CHUNK_TYPE'])
+			# check end decode image depending on chunk type
 			if chunk_type == ChunkType.JPEG_IMAGE:
+				# Check that we have received chunk type JPEG_IMAGE
+				# Convert jpeg data to image data
 				image = mpimg.imread(io.BytesIO(image_hex), format='jpg')
+				results[counter].append(image)
 			elif chunk_type == ChunkType.MONOCHROME_2D_8BIT:
-				image = np.reshape(array.array('B', image_hex), (header['IMAGE_HEIGHT'], header['IMAGE_WIDTH']))
+				# Check that we have received chunk type MONOCHROME_2D_8BIT
+				# Read pixel data and reshape to width/height
+				image = np.frombuffer(image_hex, dtype=np.uint8) \
+					.reshape((header["IMAGE_HEIGHT"], header["IMAGE_WIDTH"]))
+				results[counter].append(image)
 			else:
 				image = None
 				print("Unknown image chunk type", header['CHUNK_TYPE'])

diff --git a/source/pcic/client.py b/source/pcic/client.py
@@ -412,13 +412,19 @@ def request_last_image_taken_deserialized(self, image_id=1, datatype='ndarray'):
             # append image
             image_hex = data[header['HEADER_SIZE']:header['CHUNK_SIZE']]
             chunk_type = int(header['CHUNK_TYPE'])
-            if datatype == 'ndarray' and ChunkType.JPEG_IMAGE:
-                image = mpimg.imread(io.BytesIO(image_hex), format='jpg')
-                results[counter].append(image)
-            elif image_id == 2 and chunk_type == ChunkType.MONOCHROME_2D_8BIT:
-                image = np.frombuffer(image_hex, dtype=np.uint8)\
-                    .reshape((header["IMAGE_HEIGHT"], header["IMAGE_WIDTH"]))
-                results[counter].append(image)
+            # check end decode image depending on chunk type
+            if datatype == 'ndarray':
+                if chunk_type == ChunkType.JPEG_IMAGE:
+                    # Check that we have received chunk type JPEG_IMAGE
+                    # Convert jpeg data to image data
+                    image = mpimg.imread(io.BytesIO(image_hex), format='jpg')
+                    results[counter].append(image)
+                elif chunk_type == ChunkType.MONOCHROME_2D_8BIT:
+                    # Check that we have received chunk type MONOCHROME_2D_8BIT
+                    # Read pixel data and reshape to width/height
+                    image = np.frombuffer(image_hex, dtype=np.uint8)\
+                        .reshape((header["IMAGE_HEIGHT"], header["IMAGE_WIDTH"]))
+                    results[counter].append(image)
             elif datatype == 'bytes':
                 results[counter].append(image_hex)
             else: