How to make Dataloader for Image colorisation

vision
1

I am making a model for Image Colorization, and for this model i want to make DataLoader having gray images as the input and and its value to be LAB images.
I am having a folder containing lots of colored images.
I searched for it on the google, but was not able to get it.
Can someone please provide a good source, from where i can learn this.

Thank you

2

Pillow does this https://pillow.readthedocs.io/en/3.1.x/reference/Image.html. It seems pretty straight forward:

gray = img.convert("LA")
lab = img.convert("LAB")

Also OpenCV seems to be able to convert RGB images to both Lab and gray scale so I suppose you can take your pick :slight_smile: .

import cv2
image = cv2.imread('image.png')
gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
lab_image = cv2.cvtColor(image, cv2.COLOR_BGR2Lab)
1 Like
3

I made this class for dataloader. it would be very helpful if you colud please check that i have made it correct way(optimized way) or not.
I have renamed every image with indeces (1 to n).jpg of my dataset.

class DATALODER(Dataset):
	def __init__(self, root_dir, transform=None):
		self.root_dir = root_dir
		self.transform = transform

	def __len__(self):
		return 9294

	def __getitem__(self, idx):
		img_name = os.path.join(self.root_dir,
		str(idx) + ".jpg")
		image = io.imread(img_name)
		gray_img = rgb2gray(image)
		lab_img = rgb2lab(image)[:,:,1:]

		sample = self.transform(sample)

		return sample


class ToTensor(object):
	def __call__(self, sample):
		gray_img, lab_img = sample
		tensor_gray = torch.tensor(gray_img, device=cuda).float()
		tensor_label = torch.tensor(lab_img.transpose((2,0,1)), device=cuda).float()

		tensor_label = tensor_label.view(1,2, tensor_gray.size()[0], tensor_gray.size()[1])
		tensor_gray = tensor_gray.view(1, 1, tensor_gray.size()[0], tensor_gray.size()[1])

		return (tensor_gray,tensor_label)
4

I had a quick look and I don’t think you need to write your own ToTensor class since one already exists in pytorch (https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.ToTensor). I would just use that one to transform into tensor and then rearrange the axes with transpose to get the shape you want. In your version, you load the tensor onto the Cuda device before forming the batch but actually, I think it’s better to do all the preprocessing on the CPU and then move to the device you want.

Here is a snippet showing what I think the result should look like

from torchvision import transforms

class my_loader(Dataset):
def __init__(self, root_dir, transform=None):
	self.root_dir = root_dir
	self.transform = transforms.ToTensor()

def __len__(self):
	return 9294

def __getitem__(self, idx):
	img_name = os.path.join(self.root_dir, str(idx) + ".jpg")
	image = io.imread(img_name)
	gray_img = rgb2gray(image)
	lab_img = rgb2lab(image)[:,:,1:]

	gray_img = self.transform(gray_img)
            lab_img = self.transform(lab_img).transpose((2,0,1)

	return sample