def encode()

in petastorm/codecs.py [0:0]

• 20 lines of code
• 6 McCabe index (conditional complexity)

    def encode(self, unischema_field, value):
        """Encodes the image using OpenCV."""
        if unischema_field.numpy_dtype != value.dtype:
            raise ValueError("Unexpected type of {} feature, expected {}, got {}".format(
                unischema_field.name, unischema_field.numpy_dtype, value.dtype
            ))

        if not _is_compliant_shape(value.shape, unischema_field.shape):
            raise ValueError("Unexpected dimensions of {} feature, expected {}, got {}".format(
                unischema_field.name, unischema_field.shape, value.shape
            ))

        if len(value.shape) == 2:
            # Greyscale image
            image_bgr_or_gray = value
        elif len(value.shape) == 3 and value.shape[2] == 3:
            # Convert RGB to BGR
            image_bgr_or_gray = value[:, :, (2, 1, 0)]
        else:
            raise ValueError('Unexpected image dimensions. Supported dimensions are (H, W) or (H, W, 3). '
                             'Got {}'.format(value.shape))

        _, contents = cv2.imencode(self._image_codec,
                                   image_bgr_or_gray,
                                   [int(cv2.IMWRITE_JPEG_QUALITY), self._quality])
        return bytearray(contents)