mypkg/convert.py

import io
import cv2
from PIL import Image as pilImage
import numpy as np
import numpy.typing as npt
from base64 import b64encode

__all__ = ["Convert"]

class Convert:
    @staticmethod
    def cv_to_pil(img: np.ndarray) -> pilImage.Image:
        result = img.copy()
        if result.ndim == 2: # grayscale
            pass
        else:
            assert result.ndim == 3
            if result.shape[2] == 3: # rgb
                result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
            elif result.shape[2] == 4: # rgba
                result = cv2.cvtColor(result, cv2.COLOR_RGBA2BGRA)
            else:
                raise ValueError
        return pilImage.fromarray(result)
    
    @staticmethod
    def pil_to_cv(img: pilImage.Image | npt.NDArray[np.uint8]) -> np.ndarray:
        if type(img) is pilImage.Image:
            result = np.array(img, dtype=np.uint8)
        elif type(img) is np.ndarray:
            pass
        else:
            raise TypeError
        
        if result.ndim == 2: # grayscale
            pass
        else:
            assert result.ndim == 3
            if result.shape[2] == 3: # rgb
                result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
            elif result.shape[2] == 4: # rgba
                result = cv2.cvtColor(result, cv2.COLOR_RGBA2BGRA)
            else:
                raise ValueError
        return result
    
    @staticmethod
    def cv_to_base64(img: np.ndarray) -> str:
        return Convert.pil_to_base64(Convert.cv_to_pil(img))

    @staticmethod
    def pil_to_base64(img: pilImage.Image | npt.NDArray[np.uint8]) -> str:
        if type(img) is pilImage.Image:
            pass
        elif type(img) is np.ndarray:
            img = pilImage.fromarray(img.astype('uint8'))
        else:
            raise TypeError
        
        # https://stackoverflow.com/a/59583262/13797085
        file_object = io.BytesIO()
        img.save(file_object, 'PNG')
        base64img = "data:image/png;base64," \
            + b64encode(file_object.getvalue()).decode('ascii')
        return base64img