diff --git a/lte/modules/image_decode.py b/lte/modules/image_decode.py new file mode 100644 index 0000000..8ae7210 --- /dev/null +++ b/lte/modules/image_decode.py @@ -0,0 +1,24 @@ +""" +Decodes images from JPEG bytes to numpy array. +""" + +import io + +from PIL import Image +import numpy as np + + +def decode(data: "io.BytesIO | bytes") -> np.ndarray: + """ + Decodes a JPEG encoded image and returns it as a numpy array. + + Args: + data: bytes object containing the JPEG encoded image + TODO: When encoding JPEG, make sure the input image only has 3 channels, + if original picture was a PNG (filter out A channel) + Returns: + NDArray with in RGB format. Shape is (Height, Width, 3) + """ + image = Image.open(data, formats=["JPEG"]) + + return np.asarray(image) diff --git a/lte/test_image_encode_decode.py b/lte/test_image_encode_decode.py new file mode 100644 index 0000000..1ddd86e --- /dev/null +++ b/lte/test_image_encode_decode.py @@ -0,0 +1,27 @@ +""" +Test image encode and decode. +TODO: finish image encode and then write this test +""" + +from lte.modules import image_decode + + +def main() -> int: + # TODO: Get test images and Encode + raw_data = 0 + + # Decode + img_array = image_decode.decode(raw_data) + + # Do checks + # eg. check output shape + return 0 + + +if __name__ == "__main__": + result_main = main() + + if result_main < 0: + print(f"ERROR: Status code: {result_main}") + + print("Done!") diff --git a/requirements.txt b/requirements.txt index be7658d..42301b6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,9 @@ pytest opencv-python numpy +# Module lte +Pillow + # Module mavlink dronekit