import random

import numpy as np
import skimage.color as sc

import torch

def set_channel(*args, n_channels=3):
    def _set_channel(img):
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        c = img.shape[2]
        if n_channels == 1 and c == 3:
            img = np.expand_dims(sc.rgb2ycbcr(img)[:, :, 0], 2)
        elif n_channels == 3 and c == 1:
            img = np.concatenate([img] * n_channels, 2)

        return img

    return [_set_channel(a) for a in args]

def np2Tensor(*args, rgb_range=255, format='NCHW'):
    def _np2Tensor(img, channel_format):
        assert channel_format in ('NCHW', 'NHWC')
        img = np.ascontiguousarray(img.transpose((2, 0, 1))) if channel_format == ('NCHW') else img
        tensor = torch.from_numpy(img).float()
        tensor.mul_(rgb_range / 255)

        return tensor

    return [_np2Tensor(a, format) for a in args]