<reponame>EleutherAGI/experiment-manager
import imageio
import numpy as np
import os


class DoubleImage(object):
    def __init__(self, storage):
        self.storage = storage

    def get(self, raw: []) -> []:
        self.storage.download_file(raw[0], 'temp.jpg')
        im1 = imageio.imread('temp.jpg')
        os.remove('temp.jpg')
        self.storage.download_file(raw[1], 'temp.jpg')
        im2 = imageio.imread('temp.jpg')
        os.remove('temp.jpg')
        return im1, im2

    def put(self, val: []) -> []:
        imageio.imwrite('temp.jpg', val[0])
        key2 = self.storage.upload_file('temp.jpg')
        os.remove('temp.jpg')
        imageio.imwrite('temp.jpg', val[1])
        key1 = self.storage.upload_file('temp.jpg')
        os.remove('temp.jpg')
        return [key1, key2]
