본문 바로가기

Python

Fast JPEG image loading by TurboJPEG

반응형
from turbojpeg import TurboJPEG

# define class
class TurboJpegLoader():
    def __init__(self, suffix, resize=None, crop=None):
        super(TurboJpegLoader, self).__init__()

        self.suffix = suffix
        self.resize = resize
        self.crop = crop
        self.jpeg = TurboJPEG()

    def __call__(self, path):

        if (path.exists() is not True):
            sys.exit(f'[Error] {str(path)} does not exists!!')

        in_file = open(str(path), 'rb')
        bgr_array = self.jpeg.decode(in_file.read())
        rgb_array = cv2.cvtColor(bgr_array, cv2.COLOR_BGR2RGB)

        if (self.resize is not None):
            rgb_array = cv2.resize(rgb_array, dsize=(self.resize[0], self.resize[1]), interpolation=cv2.INTER_LINEAR)

        if (self.crop is not None):
            sc, ec = self.crop[0], self.crop[2]
            sr, er = self.crop[1], self.crop[3]
            rgb_array = rgb_array[sr:er, sc:ec, :]

        return rgb_array
        
# create class object        
img_load = TurboJpegLoader(suffix='jpg',
                           resize=self.img_aug_params['resize_dims'],
                           crop=self.img_aug_params['crop'] )
# use loader                           
img = self.img_load(Path(self.nusc.get_sample_data_path(cam_token)))
img = Image.fromarray(img)