Spaces:
Runtime error
Runtime error
File size: 5,570 Bytes
1b2a9b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import cv2
import torch
from PIL import Image
import os.path as osp
import numpy as np
from torch.utils import data
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
import torchvision.transforms.functional as TF
from .custom_transform import *
class _Coco164kCuratedFew(data.Dataset):
"""Base class
This contains fields and methods common to all COCO 164k curated few datasets:
(curated) Coco164kFew_Stuff
(curated) Coco164kFew_Stuff_People
(curated) Coco164kFew_Stuff_Animals
(curated) Coco164kFew_Stuff_People_Animals
"""
def __init__(self, root, img_size, crop_size, split = "train2017"):
super(_Coco164kCuratedFew, self).__init__()
# work out name
self.split = split
self.root = root
self.include_things_labels = False # people
self.incl_animal_things = False # animals
version = 6
name = "Coco164kFew_Stuff"
if self.include_things_labels and self.incl_animal_things:
name += "_People_Animals"
elif self.include_things_labels:
name += "_People"
elif self.incl_animal_things:
name += "_Animals"
self.name = (name + "_%d" % version)
print("Specific type of _Coco164kCuratedFew dataset: %s" % self.name)
self._set_files()
self.transform = transforms.Compose([
transforms.Resize(int(img_size)),
transforms.RandomCrop(crop_size)])
N = len(self.files)
# eqv transform
self.random_horizontal_flip = RandomHorizontalTensorFlip(N=N)
self.random_vertical_flip = RandomVerticalFlip(N=N)
self.random_resized_crop = RandomResizedCrop(N=N, res=288)
# photometric transform
self.random_color_brightness = [RandomColorBrightness(x=0.3, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)]
self.random_color_contrast = [RandomColorContrast(x=0.3, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)
self.random_color_saturation = [RandomColorSaturation(x=0.3, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)
self.random_color_hue = [RandomColorHue(x=0.1, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)
self.random_gray_scale = [RandomGrayScale(p=0.2, N=N) for _ in range(2)]
self.random_gaussian_blur = [RandomGaussianBlur(sigma=[.1, 2.], p=0.5, N=N) for _ in range(2)]
self.eqv_list = ['random_crop', 'h_flip']
self.inv_list = ['brightness', 'contrast', 'saturation', 'hue', 'gray', 'blur']
self.transform_tensor = TensorTransform()
def _set_files(self):
# Create data list by parsing the "images" folder
if self.split in ["train2017", "val2017"]:
file_list = osp.join(self.root, "curated", self.split, self.name + ".txt")
file_list = tuple(open(file_list, "r"))
file_list = [id_.rstrip() for id_ in file_list]
self.files = file_list
print("In total {} images.".format(len(self.files)))
else:
raise ValueError("Invalid split name: {}".format(self.split))
def transform_eqv(self, indice, image):
if 'random_crop' in self.eqv_list:
image = self.random_resized_crop(indice, image)
if 'h_flip' in self.eqv_list:
image = self.random_horizontal_flip(indice, image)
if 'v_flip' in self.eqv_list:
image = self.random_vertical_flip(indice, image)
return image
def transform_inv(self, index, image, ver):
"""
Hyperparameters same as MoCo v2.
(https://github.com/facebookresearch/moco/blob/master/main_moco.py)
"""
if 'brightness' in self.inv_list:
image = self.random_color_brightness[ver](index, image)
if 'contrast' in self.inv_list:
image = self.random_color_contrast[ver](index, image)
if 'saturation' in self.inv_list:
image = self.random_color_saturation[ver](index, image)
if 'hue' in self.inv_list:
image = self.random_color_hue[ver](index, image)
if 'gray' in self.inv_list:
image = self.random_gray_scale[ver](index, image)
if 'blur' in self.inv_list:
image = self.random_gaussian_blur[ver](index, image)
return image
def transform_image(self, index, image):
image1 = self.transform_inv(index, image, 0)
image1 = self.transform_tensor(image)
image2 = self.transform_inv(index, image, 1)
#image2 = TF.resize(image2, self.crop_size, Image.BILINEAR)
image2 = self.transform_tensor(image2)
return image1, image2
def __getitem__(self, index):
# same as _Coco164k
# Set paths
image_id = self.files[index]
image_path = osp.join(self.root, "images", self.split, image_id + ".jpg")
# Load an image
ori_img = Image.open(image_path)
ori_img = self.transform(ori_img)
image1, image2 = self.transform_image(index, ori_img)
if image1.shape[0] < 3:
image1 = image1.repeat(3, 1, 1)
if image2.shape[0] < 3:
image2 = image2.repeat(3, 1, 1)
rets = []
rets.append(image1)
rets.append(image2)
rets.append(index)
return rets
def __len__(self):
return len(self.files)
|