data_loader.py

from __future__ import print_function, division
import os
import torch
import pandas as pd
import skimage
from skimage import io
import numpy as np
import matplotlib.pyplot as plt
import pdb, random
from torch.utils.data import Dataset, DataLoader
import random, os, cv2

unknown_code = 128


class CompositeData(Dataset):
    """mozi added, modify from class:VideoData"""

    def __init__(self, csv_file, data_config, transform=None):
        self.frames = pd.read_csv(csv_file, sep=';')
        self.transform = transform
        self.resolution = data_config['reso']

    def __len__(self):
        return len(self.frames)

    def __getitem__(self, idx):
        fg = io.imread(self.frames.iloc[idx, 0])
        seg = io.imread(self.frames.iloc[idx, 1])
        img = io.imread(self.frames.iloc[idx, 2])
        back = io.imread(self.frames.iloc[idx, 3])
        back_rnd = io.imread(self.frames.iloc[idx, 4])

        # print('debug', self.frames.iloc[idx, 2])
        if seg.max() <= 1.1:
            seg = seg * 255

        # fr1 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 3]), cv2.COLOR_BGR2GRAY)
        # fr2 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 4]), cv2.COLOR_BGR2GRAY)
        # fr3 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 5]), cv2.COLOR_BGR2GRAY)
        # fr4 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 6]), cv2.COLOR_BGR2GRAY)

        # back_rnd = io.imread(self.frames.iloc[idx, 7])  # target_back

        sz = self.resolution

        if np.random.random_sample() > 0.5:
            img = cv2.flip(img, 1)
            fg = cv2.flip(fg, 1)
            seg = cv2.flip(seg, 1)
            back = cv2.flip(back, 1)
            back_rnd = cv2.flip(back_rnd, 1)
            # fr1 = cv2.flip(fr1, 1);
            # fr2 = cv2.flip(fr2, 1);
            # fr3 = cv2.flip(fr3, 1);
            # fr4 = cv2.flip(fr4, 1)

        # make frames together
        # multi_fr = np.zeros((img.shape[0], img.shape[1], 4))
        # multi_fr[..., 0] = fr1;
        # multi_fr[..., 1] = fr2;
        # multi_fr[..., 2] = fr3;
        # multi_fr[..., 3] = fr4;

        # allow random cropping centered on the segmentation map
        bbox = create_bbox(seg, seg.shape[0], seg.shape[1])
        img = apply_crop(img, bbox, self.resolution)
        fg = apply_crop(fg, bbox, self.resolution)
        seg = apply_crop(seg, bbox, self.resolution)
        back = apply_crop(back, bbox, self.resolution)
        back_rnd = apply_crop(back_rnd, bbox, self.resolution)

        # convert seg to guidance map
        # segg=create_seg_guide(seg,self.resolution)

        sample = {'image': to_tensor(img), 'fg': to_tensor(fg),
                  'seg': to_tensor(create_seg_guide(seg, self.resolution)),
                  'bg': to_tensor(back), 'seg-gt': to_tensor(seg), 'back-rnd': to_tensor(back_rnd)}

        if self.transform:
            sample = self.transform(sample)
        return sample


class RealDataAndAdobe(Dataset):
    """mozi added, modify from class:VideoData"""

    def __init__(self, csv_file, data_config, transform=None):
        self.frames = pd.read_csv(csv_file, sep=';')
        self.transform = transform
        self.resolution = data_config['reso']

    def __len__(self):
        return len(self.frames)

    def __getitem__(self, idx):
        img = io.imread(self.frames.iloc[idx, 0])
        back = io.imread(self.frames.iloc[idx, 1])
        seg = io.imread(self.frames.iloc[idx, 2])
        fg = io.imread(self.frames.iloc[idx, 3])
        back_rnd = io.imread(self.frames.iloc[idx, 4])
        data_type = int(self.frames.iloc[idx, 5])

        # print('debug', self.frames.iloc[idx, 2])
        if seg.max() <= 1.1:
            seg = seg * 255

        sz = self.resolution

        if np.random.random_sample() > 0.5:
            img = cv2.flip(img, 1)
            seg = cv2.flip(seg, 1)
            back = cv2.flip(back, 1)
            fg = cv2.flip(fg, 1)
            back_rnd = cv2.flip(back_rnd, 1)

        # allow random cropping centered on the segmentation map
        try:
            bbox = create_bbox(seg, seg.shape[0], seg.shape[1])
        except:
            print("warning:can't create bbox. ", self.frames.iloc[idx, 2])
            bbox = [0, 0, seg.shape[0], seg.shape[1]]
        if back_rnd.shape[0] <= bbox[0] or back_rnd.shape[1] <= bbox[1]:
            # mozi added，防止back_rnd太小，crop出空。
            back_rnd = cv2.resize(back_rnd, (img.shape[1], img.shape[0]))
            # print('debug resized', back_rnd.shape, bbox)
        img = apply_crop(img, bbox, self.resolution)
        seg = apply_crop(seg, bbox, self.resolution)
        back = apply_crop(back, bbox, self.resolution)
        # print('debug', back_rnd.shape, bbox)
        fg = apply_crop(fg, bbox, self.resolution)
        back_rnd = apply_crop(back_rnd, bbox, self.resolution)

        # convert seg to guidance map
        # segg=create_seg_guide(seg,self.resolution)

        # Perturb Background: random noise addition or gamma change
        if np.random.random_sample() > 0.6:
            sigma = np.random.randint(low=2, high=6)
            mu = np.random.randint(low=0, high=14) - 7
            back_tr = add_noise(back, mu, sigma)
        else:
            back_tr = skimage.exposure.adjust_gamma(back, np.random.normal(1, 0.12))

        sample = {'image': to_tensor(img),
                  'seg': to_tensor(create_seg_guide4rgbd(seg, self.resolution)),
                  'bg': to_tensor(back), 'seg-gt': to_tensor(seg), 'fg': to_tensor(fg), 'back-rnd': to_tensor(back_rnd),
                  'back_tr': to_tensor(back_tr), 'data_type': data_type}

        if self.transform:
            sample = self.transform(sample)
        return sample


class RealDataWoMotion(Dataset):
    """mozi added, modify from class:VideoData"""

    def __init__(self, csv_file, data_config, transform=None):
        self.frames = pd.read_csv(csv_file, sep=';')
        self.transform = transform
        self.resolution = data_config['reso']

    def __len__(self):
        return len(self.frames)

    def __getitem__(self, idx):
        img = io.imread(self.frames.iloc[idx, 0])
        back = io.imread(self.frames.iloc[idx, 1])
        seg = io.imread(self.frames.iloc[idx, 2])
        back_rnd = io.imread(self.frames.iloc[idx, 3])

        # print('debug', self.frames.iloc[idx, 2])
        if seg.max() <= 1.1:
            seg = seg * 255

        # fr1 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 3]), cv2.COLOR_BGR2GRAY)
        # fr2 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 4]), cv2.COLOR_BGR2GRAY)
        # fr3 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 5]), cv2.COLOR_BGR2GRAY)
        # fr4 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 6]), cv2.COLOR_BGR2GRAY)

        # back_rnd = io.imread(self.frames.iloc[idx, 7])  # target_back

        sz = self.resolution

        if np.random.random_sample() > 0.5:
            img = cv2.flip(img, 1)
            seg = cv2.flip(seg, 1)
            back = cv2.flip(back, 1)
            back_rnd = cv2.flip(back_rnd, 1)
            # fr1 = cv2.flip(fr1, 1);
            # fr2 = cv2.flip(fr2, 1);
            # fr3 = cv2.flip(fr3, 1);
            # fr4 = cv2.flip(fr4, 1)

        # make frames together
        # multi_fr = np.zeros((img.shape[0], img.shape[1], 4))
        # multi_fr[..., 0] = fr1;
        # multi_fr[..., 1] = fr2;
        # multi_fr[..., 2] = fr3;
        # multi_fr[..., 3] = fr4;

        # allow random cropping centered on the segmentation map
        try:
            bbox = create_bbox(seg, seg.shape[0], seg.shape[1])
        except:
            print("warning:can't create bbox. ", self.frames.iloc[idx, 2])
            bbox = [0, 0, seg.shape[0], seg.shape[1]]
        if back_rnd.shape[0] <= bbox[0] or back_rnd.shape[1] <= bbox[1]:
            # mozi added，防止back_rnd太小，crop出空。
            back_rnd = cv2.resize(back_rnd, (img.shape[1], img.shape[0]))
            # print('debug resized', back_rnd.shape, bbox)
        img = apply_crop(img, bbox, self.resolution)
        seg = apply_crop(seg, bbox, self.resolution)
        back = apply_crop(back, bbox, self.resolution)
        # print('debug', back_rnd.shape, bbox)
        back_rnd = apply_crop(back_rnd, bbox, self.resolution)

        # convert seg to guidance map
        # segg=create_seg_guide(seg,self.resolution)

        sample = {'image': to_tensor(img),
                  'seg': to_tensor(create_seg_guide4rgbd(seg, self.resolution)),
                  'bg': to_tensor(back), 'seg-gt': to_tensor(seg), 'back-rnd': to_tensor(back_rnd)}

        if self.transform:
            sample = self.transform(sample)
        return sample


class VideoData(Dataset):
    def __init__(self, csv_file, data_config, transform=None):
        self.frames = pd.read_csv(csv_file, sep=';')
        self.transform = transform
        self.resolution = data_config['reso']

    def __len__(self):
        return len(self.frames)

    def __getitem__(self, idx):
        img = io.imread(self.frames.iloc[idx, 0])
        back = io.imread(self.frames.iloc[idx, 1])
        seg = io.imread(self.frames.iloc[idx, 2])

        fr1 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 3]), cv2.COLOR_BGR2GRAY)
        fr2 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 4]), cv2.COLOR_BGR2GRAY)
        fr3 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 5]), cv2.COLOR_BGR2GRAY)
        fr4 = cv2.cvtColor(io.imread(self.frames.iloc[idx, 6]), cv2.COLOR_BGR2GRAY)

        back_rnd = io.imread(self.frames.iloc[idx, 7])  # target_back

        sz = self.resolution

        if np.random.random_sample() > 0.5:
            img = cv2.flip(img, 1)
            seg = cv2.flip(seg, 1)
            back = cv2.flip(back, 1)
            back_rnd = cv2.flip(back_rnd, 1)
            fr1 = cv2.flip(fr1, 1);
            fr2 = cv2.flip(fr2, 1);
            fr3 = cv2.flip(fr3, 1);
            fr4 = cv2.flip(fr4, 1)

        # make frames together
        multi_fr = np.zeros((img.shape[0], img.shape[1], 4))
        multi_fr[..., 0] = fr1;
        multi_fr[..., 1] = fr2;
        multi_fr[..., 2] = fr3;
        multi_fr[..., 3] = fr4;

        # allow random cropping centered on the segmentation map
        bbox = create_bbox(seg, seg.shape[0], seg.shape[1])
        img = apply_crop(img, bbox, self.resolution)
        seg = apply_crop(seg, bbox, self.resolution)
        back = apply_crop(back, bbox, self.resolution)
        back_rnd = apply_crop(back_rnd, bbox, self.resolution)
        multi_fr = apply_crop(multi_fr, bbox, self.resolution)

        # convert seg to guidance map
        # segg=create_seg_guide(seg,self.resolution)

        sample = {'image': to_tensor(img), 'seg': to_tensor(create_seg_guide(seg, self.resolution)),
                  'bg': to_tensor(back), 'multi_fr': to_tensor(multi_fr), 'seg-gt': to_tensor(seg),
                  'back-rnd': to_tensor(back_rnd)}

        if self.transform:
            sample = self.transform(sample)
        return sample


class AdobeDataAffineHR(Dataset):
    def __init__(self, csv_file, data_config, transform=None):
        self.frames = pd.read_csv(csv_file, sep=';')
        self.transform = transform
        self.resolution = data_config['reso']
        self.trimapK = data_config['trimapK']
        self.noise = data_config['noise']

    def __len__(self):
        return len(self.frames)

    def __getitem__(self, idx):
        try:
            # load
            fg = io.imread(self.frames.iloc[idx, 0])
            alpha = io.imread(self.frames.iloc[idx, 1])
            image = io.imread(self.frames.iloc[idx, 2])
            back = io.imread(self.frames.iloc[idx, 3])

            fg = cv2.resize(fg, dsize=(800, 800))
            alpha = cv2.resize(alpha, dsize=(800, 800))
            back = cv2.resize(back, dsize=(800, 800))
            image = cv2.resize(image, dsize=(800, 800))

            sz = self.resolution

            # random flip
            if np.random.random_sample() > 0.5:
                alpha = cv2.flip(alpha, 1)
                fg = cv2.flip(fg, 1)
                back = cv2.flip(back, 1)
                image = cv2.flip(image, 1)

            trimap = generate_trimap(alpha, self.trimapK[0], self.trimapK[1], False)

            # randcom crop+scale
            different_sizes = [(576, 576), (608, 608), (640, 640), (672, 672), (704, 704), (736, 736), (768, 768),
                               (800, 800)]
            crop_size = random.choice(different_sizes)

            x, y = random_choice(trimap, crop_size)

            fg = safe_crop(fg, x, y, crop_size, sz)
            alpha = safe_crop(alpha, x, y, crop_size, sz)
            image = safe_crop(image, x, y, crop_size, sz)
            back = safe_crop(back, x, y, crop_size, sz)
            trimap = safe_crop(trimap, x, y, crop_size, sz)

            # Perturb Background: random noise addition or gamma change
            if self.noise:
                if np.random.random_sample() > 0.6:
                    sigma = np.random.randint(low=2, high=6)
                    mu = np.random.randint(low=0, high=14) - 7
                    back_tr = add_noise(back, mu, sigma)
                else:
                    back_tr = skimage.exposure.adjust_gamma(back, np.random.normal(1, 0.12))

            # Create motion cues: transform foreground and create 4 additional images
            affine_fr = np.zeros((fg.shape[0], fg.shape[1], 4))
            for t in range(0, 4):
                T = np.random.normal(0, 5, (2, 1));
                theta = np.random.normal(0, 7);
                R = np.array([[np.cos(np.deg2rad(theta)), -np.sin(np.deg2rad(theta))],
                              [np.sin(np.deg2rad(theta)), np.cos(np.deg2rad(theta))]])
                sc = np.array([[1 + np.random.normal(0, 0.05), 0], [0, 1]]);
                sh = np.array([[1, np.random.normal(0, 0.05) * (np.random.random_sample() > 0.5)],
                               [np.random.normal(0, 0.05) * (np.random.random_sample() > 0.5), 1]]);
                A = np.concatenate((sc * sh * R, T), axis=1);

                fg_tr = cv2.warpAffine(fg.astype(np.uint8), A, (fg.shape[1], fg.shape[0]), flags=cv2.INTER_LINEAR,
                                       borderMode=cv2.BORDER_REFLECT)
                alpha_tr = cv2.warpAffine(alpha.astype(np.uint8), A, (fg.shape[1], fg.shape[0]),
                                          flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_REFLECT)

                sigma = np.random.randint(low=2, high=6)
                mu = np.random.randint(low=0, high=14) - 7
                back_tr0 = add_noise(back, mu, sigma)

                affine_fr[..., t] = cv2.cvtColor(composite(fg_tr, back_tr0, alpha_tr), cv2.COLOR_BGR2GRAY)

            sample = {'image': to_tensor(image), 'fg': to_tensor(fg), 'alpha': to_tensor(alpha), 'bg': to_tensor(back),
                      'trimap': to_tensor(trimap), 'bg_tr': to_tensor(back_tr),
                      'seg': to_tensor(create_seg(alpha, trimap)), 'multi_fr': to_tensor(affine_fr)}

            if self.transform:
                sample = self.transform(sample)
            return sample
        except Exception as e:
            print("Error loading: " + self.frames.iloc[idx, 0])
            print(e)


# Functions

def create_seg_guide(rcnn, reso):
    kernel_er = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    kernel_dil = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    rcnn = rcnn.astype(np.float32) / 255;
    rcnn[rcnn > 0.2] = 1;
    K = 25

    zero_id = np.nonzero(np.sum(rcnn, axis=1) == 0)
    del_id = zero_id[0][zero_id[0] > 250]
    if len(del_id) > 0:
        del_id = [del_id[0] - 2, del_id[0] - 1, *del_id]
        rcnn = np.delete(rcnn, del_id, 0)
    rcnn = cv2.copyMakeBorder(rcnn, 0, K + len(del_id), 0, 0, cv2.BORDER_REPLICATE)

    rcnn = cv2.erode(rcnn, kernel_er, iterations=np.random.randint(10, 20))
    rcnn = cv2.dilate(rcnn, kernel_dil, iterations=np.random.randint(3, 7))
    k_size_list = [(21, 21), (31, 31), (41, 41)]
    rcnn = cv2.GaussianBlur(rcnn.astype(np.float32), random.choice(k_size_list), 0)
    rcnn = (255 * rcnn).astype(np.uint8)
    rcnn = np.delete(rcnn, range(reso[0], reso[0] + K), 0)

    return rcnn


def create_seg_guide4rgbd(rcnn, reso):
    kernel_er = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    kernel_dil = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    rcnn = rcnn.astype(np.float32) / 255;
    rcnn[rcnn > 0.2] = 1;
    K = 25

    zero_id = np.nonzero(np.sum(rcnn, axis=1) == 0)
    del_id = zero_id[0][zero_id[0] > 250]
    if len(del_id) > 0:
        del_id = [del_id[0] - 2, del_id[0] - 1, *del_id]
        rcnn = np.delete(rcnn, del_id, 0)
    rcnn = cv2.copyMakeBorder(rcnn, 0, K + len(del_id), 0, 0, cv2.BORDER_REPLICATE)

    rcnn = cv2.erode(rcnn, kernel_er, iterations=np.random.randint(5, 20))
    rcnn = cv2.dilate(rcnn, kernel_dil, iterations=np.random.randint(5, 15))
    k_size_list = [(21, 21), (31, 31), (41, 41)]
    rcnn = cv2.GaussianBlur(rcnn.astype(np.float32), random.choice(k_size_list), 0)
    rcnn = (255 * rcnn).astype(np.uint8)
    rcnn = np.delete(rcnn, range(reso[0], reso[0] + K), 0)

    return rcnn


def crop_holes(img, cx, cy, crop_size):
    img[cy:cy + crop_size[0], cx:cx + crop_size[1]] = 0
    return img


def create_seg(alpha, trimap):
    # old
    num_holes = np.random.randint(low=0, high=3)
    crop_size_list = [(15, 15), (25, 25), (35, 35), (45, 45)]
    kernel_er = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    kernel_dil = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    seg = (alpha > 0.5).astype(np.float32)
    # print('Before %.4f max: %.4f' %(seg.sum(),seg.max()))
    # old
    seg = cv2.erode(seg, kernel_er, iterations=np.random.randint(low=10, high=20))
    seg = cv2.dilate(seg, kernel_dil, iterations=np.random.randint(low=15, high=30))
    # print('After %.4f max: %.4f' %(seg.sum(),seg.max()))
    seg = seg.astype(np.float32)
    seg = (255 * seg).astype(np.uint8)
    for i in range(num_holes):
        crop_size = random.choice(crop_size_list)
        cx, cy = random_choice(trimap, crop_size)
        seg = crop_holes(seg, cx, cy, crop_size)
        trimap = crop_holes(trimap, cx, cy, crop_size)
    k_size_list = [(21, 21), (31, 31), (41, 41)]
    seg = cv2.GaussianBlur(seg.astype(np.float32), random.choice(k_size_list), 0)
    return seg.astype(np.uint8)


def apply_crop(img, bbox, reso):
    img_crop = img[bbox[0]:bbox[0] + bbox[2], bbox[1]:bbox[1] + bbox[3], ...];
    img_crop = cv2.resize(img_crop, reso)
    return img_crop


def create_bbox(mask, R, C):
    where = np.array(np.where(mask))
    x1, y1 = np.amin(where, axis=1)
    x2, y2 = np.amax(where, axis=1)

    w = np.maximum(y2 - y1, x2 - x1);
    bd = np.random.uniform(0.1, 0.4)
    x1 = x1 - np.round(bd * w)
    y1 = y1 - np.round(bd * w)
    y2 = y2 + np.round(bd * w)

    if x1 < 0: x1 = 0
    if y1 < 0: y1 = 0
    if y2 >= C: y2 = C
    if x2 >= R: x2 = R - 1

    bbox = np.around([x1, y1, x2 - x1, y2 - y1]).astype('int')

    return bbox


def composite(fg, bg, a):
    fg = fg.astype(np.float32);
    bg = bg.astype(np.float32);
    a = a.astype(np.float32);
    alpha = np.expand_dims(a / 255, axis=2)
    im = alpha * fg + (1 - alpha) * bg
    im = im.astype(np.uint8)
    return im


def add_noise(back, mean, sigma):
    back = back.astype(np.float32)
    row, col, ch = back.shape
    gauss = np.random.normal(mean, sigma, (row, col, ch))
    gauss = gauss.reshape(row, col, ch)
    # gauss = np.repeat(gauss[:, :, np.newaxis], ch, axis=2)
    noisy = back + gauss

    noisy[noisy < 0] = 0;
    noisy[noisy > 255] = 255;

    return noisy.astype(np.uint8)


def safe_crop(mat, x, y, crop_size, img_size, cubic=True):
    img_rows, img_cols = img_size
    crop_height, crop_width = crop_size
    if len(mat.shape) == 2:
        ret = np.zeros((crop_height, crop_width), np.float32)
    else:
        ret = np.zeros((crop_height, crop_width, 3), np.float32)
    crop = mat[y:y + crop_height, x:x + crop_width]
    h, w = crop.shape[:2]
    ret[0:h, 0:w] = crop
    if crop_size != (img_rows, img_cols):
        if cubic:
            ret = cv2.resize(ret, dsize=(img_rows, img_cols))
        else:
            ret = cv2.resize(ret, dsize=(img_rows, img_cols), interpolation=cv2.INTER_NEAREST)
    return ret


def generate_trimap(alpha, K1, K2, train_mode):
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    fg = np.array(np.equal(alpha, 255).astype(np.float32))
    if train_mode:
        K = np.random.randint(K1, K2)
    else:
        K = np.round((K1 + K2) / 2).astype('int')

    fg = cv2.erode(fg, kernel, iterations=K)
    unknown = np.array(np.not_equal(alpha, 0).astype(np.float32))
    unknown = cv2.dilate(unknown, kernel, iterations=2 * K)
    trimap = fg * 255 + (unknown - fg) * 128
    return trimap.astype(np.uint8)


def random_choice(trimap, crop_size=(320, 320)):
    img_height, img_width = trimap.shape[0:2]
    crop_height, crop_width = crop_size

    val_idx = np.zeros((img_height, img_width))
    val_idx[int(crop_height / 2):int(img_height - crop_height / 2),
    int(crop_width / 2):int(img_width - crop_width / 2)] = 1

    y_indices, x_indices = np.where(np.logical_and(trimap == unknown_code, val_idx == 1))
    num_unknowns = len(y_indices)
    x, y = 0, 0
    if num_unknowns > 0:
        ix = np.random.choice(range(num_unknowns))
        center_x = x_indices[ix]
        center_y = y_indices[ix]
        x = max(0, center_x - int(crop_width / 2))
        y = max(0, center_y - int(crop_height / 2))

    # added extra

    return x, y


def to_tensor(pic):
    if len(pic.shape) >= 3:
        img = torch.from_numpy(pic.transpose((2, 0, 1)))
    else:
        img = torch.from_numpy(pic)
        img = img.unsqueeze(0)
    # backward compatibility

    return 2 * (img.float().div(255)) - 1


def select_datatype(adobe_data, real_data, datatype):
    bg_res = torch.ones_like(real_data)
    for i in range(real_data.shape[0]):
        if datatype[i] == 1:
            bg_res[i, ...] = adobe_data[i, ...]
        else:
            bg_res[i, ...] = real_data[i, ...]
    return bg_res