def __init__(self):
"""
Change the tensor as the gray scale
The function will turn the BCHW tensor into B1HW gray-scaled tensor
"""
INFO("Applied << %15s >>" % self.__class__.__name__)
INFO("* Notice: the rank format of input tensor should be 'BCHW'")
def print(self):
"""
Print the information for each image domain
"""
INFO()
for domain in range(len(self.root)):
INFO("domain index: %d \timage number: %d" %
(domain, len(self.files[domain])))
INFO()
def __init__(self, direction=BHW2BHWC):
"""
Stack the gray-scale image for 3 times to become RGB image
If the input is already RGB image, this function do nothing
Arg: direction - The stack direction you want to conduct
"""
INFO("Applied << %15s >>" % self.__class__.__name__)
INFO("* Notice: the rank format of input tensor should be 'BHWC'")
self.direction = direction
def __init__(self, output_size):
"""
Resize the tensor to the desired size
This function only support for nearest-neighbor interpolation
Since this mechanism can also deal with categorical data
Arg: output_size - The tuple (H, W)
"""
self.output_size = output_size
INFO("Applied << %15s >>" % self.__class__.__name__)
INFO("* Notice: the rank format of input tensor should be 'BHWC'")
def __init__(self, p=0.5):
"""
Flip the tensor toward horizontal direction randomly
Arg: p - The random probability to filp the tensor
"""
INFO("Applied << %15s >>" % self.__class__.__name__)
INFO("* Notice: the rank format of input tensor should be 'BCHW'")
if p < 0.0 or p > 1.0:
raise Exception(
"The parameter 'p' should in (0, 1], but get {}".format(p))
self.p = p
def print(self):
"""
Print the information for each image domain
"""
INFO()
for domain in range(len(self.root)):
total_frame = 0
for video in self.files[domain]:
total_frame += len(video)
INFO("domain index: %d \tvideo number: %d\tframe total: %d" %
(domain, len(self.root[domain]), total_frame))
INFO()
def __init__(self, direction = BHWC2BCHW):
"""
Transfer the rank of tensor into target one
Arg: direction - The direction you want to do the transpose
"""
self.direction = direction
if self.direction == BHWC2BCHW:
INFO("Applied << %15s >>, The rank format is BCHW" % self.__class__.__name__)
elif self.direction == BCHW2BHWC:
INFO("Applied << %15s >>, The rank format is BHWC" % self.__class__.__name__)
else:
raise Exception("Unknown direction symbol: {}".format(self.direction))
def __init__(self, mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]):
"""
Unnormalize the tensor with given mean and standard deviation
* Notice: If you didn't give mean and std, the function will assume that the original distribution locates in [-1, 1]
Args:
mean - The mean of the result tensor
std - The standard deviation
"""
self.mean = mean
self.std = std
INFO("Applied << %15s >>" % self.__class__.__name__)
INFO("* Notice: the rank format of input tensor should be 'BCHW'")
if self.mean == [0.485, 0.456, 0.406] and self.std == [0.229, 0.224, 0.225]:
INFO("* Notice: The function assume that the normalize method is the same as VGG preprocessing")
def __init__(self, mean = [127.5, 127.5, 127.5], std = [127.5, 127.5, 127.5]):
"""
Unnormalize the tensor with given mean and standard deviation
* Notice: If you didn't give mean and std, the function will assume that the original distribution locates in [-1, 1]
Args:
mean - The mean of the result tensor
std - The standard deviation
"""
self.mean = mean
self.std = std
INFO("Applied << %15s >>" % self.__class__.__name__)
INFO("* Notice: the rank format of input tensor should be 'BCHW'")
if self.mean == [127.5, 127.5, 127.5] and self.std == [127.5, 127.5, 127.5]:
INFO("* Notice: The function assume that the input range is [-1, 1]")
def getCategoricalMapping(loader=None,
path='torchvision_sunner_categories_pallete.json'):
"""
This function can statistic the different category with color
And return the list of the mapping OrderedDict object
Arg: loader - The ImageLoader object
path - The path of pallete file
Ret: The list of OrderDict object (palletes object)
"""
INFO("Applied << %15s >>" % getCategoricalMapping.__name__)
INFO("* Notice: the rank format of input tensor should be 'BHWC'")
INFO("* Notice: The range of tensor should be in [0, 255]")
if os.path.exists(path):
palletes = load_pallete(path)
else:
INFO(">> Load from scratch, please wait...")
# Get the number of folder
folder_num = 0
for img_list in loader:
folder_num = len(img_list)
break
# Initialize the pallete list
palletes = [OrderedDict()] * folder_num
color_sets = [set()] * folder_num
# Work
for img_list in tqdm(loader):
for folder_idx in range(folder_num):
img = img_list[folder_idx]
if torch.max(img) > 255 or torch.min(img) < 0:
raise Exception(
'tensor value out of range...\t range is [' +
str(torch.min(img)) + ' ~ ' + str(torch.max(img)))
img = img.cpu().data.numpy().astype(np.uint8)
img = np.reshape(img, [-1, 3])
color_sets[folder_idx] |= set([tuple(_) for _ in img])
# Merge the color
for i in range(folder_num):
for color in color_sets[i]:
if color not in palletes[i].keys():
palletes[i][color] = len(palletes[i])
save_pallete(palletes, path)
return palletes
def getFiles(self):
"""
Construct the files object for the assigned root
We accept the user to mix folder with image
This function can extract whole image in the folder
However, unlike the setting in ImageDataset, we store the video result in root obj.
Also, the 'images' name will be store in files obj
The following list the progress of this function:
1. check if we need to decode again
2. decode if needed
3. form the files obj
"""
if not self.files:
# Check if the decode process should be conducted again
should_decode = not os.path.exists(self.decode_root)
if not should_decode:
for domain_idx, domain in enumerate(self.root):
for video in domain:
if not os.path.exists(
os.path.join(self.decode_root, str(domain_idx),
self.to_folder(video))):
should_decode = True
break
# Decode the video if needed
if should_decode:
INFO("Decode from scratch...")
if os.path.exists(self.decode_root):
subprocess.call(['rm', '-rf', self.decode_root])
os.mkdir(self.decode_root)
self.decodeVideo()
else:
INFO("Skip the decode process!")
# Form the files object
self.files = {}
for domain_idx, domain in enumerate(os.listdir(self.decode_root)):
self.files[domain_idx] = []
for video in os.listdir(os.path.join(self.decode_root,
domain)):
self.files[domain_idx] += [
sorted(
glob(
os.path.join(self.decode_root, domain, video,
"*")))
]
def save(self, remain_file_name, split_ratio, split_file_name = ".split.pkl", save_type = 'image'):
"""
Save the information into record file
Arg: remain_file_name - The path of record file which store the information of remain data
split_ratio - Float. The proportion to split the data. Usually used to split the testing data
split_file_name - The path of record file which store the information of split data
save_type - Str. The type of the record file you want to save
"""
if self.save_file:
if not os.path.exists(remain_file_name):
with open(remain_file_name, 'wb') as f:
pickle.dump({
'type': save_type,
'root': self.root,
'files': self.files
}, f)
if split_ratio:
INFO("Split the dataset, and save as {}".format(split_file_name))
with open(split_file_name, 'wb') as f:
pickle.dump({
'type': save_type,
'root': self.root,
'files': self.split_files
}, f)
def __init__(self, pallete=None, direction=COLOR2INDEX, index_default=0):
"""
Transform the tensor into the particular format
We support for 3 different kinds of format:
1. one hot image
2. index image
3. color
Arg: pallete - The pallete object (default is None)
direction - The direction you want to change
index_default - The default index if the color cannot be found in the pallete
"""
self.pallete = pallete
self.direction = direction
self.index_default = index_default
INFO("Applied << %15s >> , direction: %s" %
(self.__class__.__name__, self.direction))
INFO("* Notice: The range of tensor should be in [-1, 1]")
INFO("* Notice: the rank format of input tensor should be 'BCHW'")
def __init__(self):
"""
Change the tensor into torch.Tensor type
However, if the input is PIL image, then the original ToTensor will be used
For the range of output tensor:
1. [0~255] => [0~1] if the image is PIL object
2. otherwise the value range doesn't change
"""
INFO("Applied << %15s >>" % self.__class__.__name__)
self.official_op_obj = transforms.ToTensor()
def __init__(self,
root=None,
file_name='.remain.pkl',
sample_method=UNDER_SAMPLING,
transforms=None,
split_ratio=0.0,
save_file=False):
"""
The constructor of ImageDataset
Arg: root - The list object. The image set
file_name - The str. The name of record file.
sample_method - sunnerData.UNDER_SAMPLING or sunnerData.OVER_SAMPLING. Use down sampling or over sampling to deal with data unbalance problem.
(default is sunnerData.OVER_SAMPLING)
transform - transform.Compose object. You can declare some pre-process toward the image
split_ratio - Float. The proportion to split the data. Usually used to split the testing data
save_file - Bool. If storing the record file or not. Default is False
"""
super().__init__()
# Record the parameter
self.root = root
self.file_name = file_name
self.sample_method = sample_method
self.transforms = transforms
self.split_ratio = split_ratio
self.save_file = save_file
self.img_num = -1
INFO()
# Substitude the contain of record file if the record file is exist
if os.path.exists(file_name) and self.loadFromFile(file_name):
self.getImgNum()
elif not os.path.exists(file_name) and root is None:
raise Exception(
"Record file {} not found. You should assign 'root' parameter!"
.format(file_name))
else:
# Extend the images of folder into domain list
self.getFiles()
# Change root obj as the index format
self.root = range(len(self.root))
# Adjust the image number
self.getImgNum()
# Split the files if split_ratio is more than 0.0
self.split()
# Save the split information
self.save()
# Print the domain information
self.print()
def loadFromFile(self, file_name, check_type = 'image'):
"""
Load the root and files information from .pkl record file
This function will return False if the record file format is invalid
Arg: file_name - The name of record file
check_type - Str. The type of the record file you want to check
Ret: If the loading procedure are successful or not
"""
with open(file_name, 'rb') as f:
obj = pickle.load(f)
self.type = obj['type']
if self.type == check_type:
INFO("Load from file: {}".format(file_name))
self.root = obj['root']
self.files = obj['files']
return True
else:
INFO("Record file type: {}\tFail to load...".format(self.type))
INFO("Form the contain from scratch...")
return False
def __init__(self, mean = [127.5, 127.5, 127.5], std = [127.5, 127.5, 127.5]):
"""
Normalize the tensor with given mean and standard deviation
* Notice: If you didn't give mean and std, the result will locate in [-1, 1]
Args:
mean - The mean of the result tensor
std - The standard deviation
"""
self.mean = mean
self.std = std
INFO("Applied << %15s >>" % self.__class__.__name__)
INFO("* Notice: the rank format of input tensor should be 'BCHW'")
INFO("*****************************************************************")
INFO("* Notice: You should must call 'ToFloat' before normalization")
INFO("*****************************************************************")
if self.mean == [127.5, 127.5, 127.5] and self.std == [127.5, 127.5, 127.5]:
INFO("* Notice: The result will locate in [-1, 1]")
def __init__(self, mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]):
"""
Normalize the tensor with given mean and standard deviation
We recommand you to set mean as [0.5, 0.5, 0.5], and std as [0.5, 0.5, 0.5]
Then the range will locate in [-1, 1]
* Notice: If you didn't give mean and std, then we will follow the preprocessing of VGG
However, The range is NOT located in [-1, 1]
Args:
mean - The mean of the result tensor
std - The standard deviation
"""
self.mean = mean
self.std = std
INFO("Applied << %15s >>" % self.__class__.__name__)
INFO("* Notice: the rank format of input tensor should be 'BCHW'")
INFO("*****************************************************************")
INFO("* Notice: You should must call 'ToFloat' before normalization")
INFO("*****************************************************************")
if self.mean == [0.485, 0.456, 0.406] and self.std == [0.229, 0.224, 0.225]:
INFO("* Notice: The result will NOT locate in [-1, 1]")
def __init__(self):
"""
Change the tensor into torch.FloatTensor
"""
INFO("Applied << %15s >>" % self.__class__.__name__)
def __init__(self,
root=None,
file_name='.remain.pkl',
T=10,
sample_method=UNDER_SAMPLING,
transforms=None,
split_ratio=0.0,
decode_root='./.decode',
save_file=False):
"""
The constructor of VideoDataset
Arg: root - The list object. The image set
file_name - Str. The name of record file.
T - Int. The maximun length of small video sequence
sample_method - sunnerData.UNDER_SAMPLING or sunnerData.OVER_SAMPLING. Use down sampling or over sampling to deal with data unbalance problem.
(default is sunnerData.OVER_SAMPLING)
transform - transform.Compose object. You can declare some pre-process toward the image
split_ratio - Float. The proportion to split the data. Usually used to split the testing data
decode_root - Str. The path to store the ffmpeg decode result.
save_file - Bool. If storing the record file or not. Default is False
"""
super().__init__()
# Record the parameter
self.root = root
self.file_name = file_name
self.T = T
self.sample_method = sample_method
self.transforms = transforms
self.split_ratio = split_ratio
self.decode_root = decode_root
self.video_num = -1
self.split_root = None
INFO()
# Substitude the contain of record file if the record file is exist
if not os.path.exists(file_name) and root is None:
raise Exception(
"Record file {} not found. You should assign 'root' parameter!"
.format(file_name))
elif os.path.exists(file_name):
INFO("Load from file: {}".format(file_name))
self.loadFromFile(file_name)
# Extend the images of folder into domain list
self.extendFolder()
# Split the image
self.split()
# Form the files obj
self.getFiles()
# Adjust the image number
self.getVideoNum()
# Save the split information
self.save()
# Print the domain information
self.print()