mirror
/
MODNet
mirror of https://github.com/ZHKKKe/MODNet.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

train bug

pull/177/head
actboy 2022-02-11 10:03:00 +08:00
parent f5396b155e
commit 77118b2a4e
1 changed files with 1 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/matting_dataset.py
View File

@ -45,14 +45,6 @@ class MattingDataset(Dataset):
class Rescale(object): class Rescale(object):
"""Rescale the image in a sample to a given size.
Args:
output_size (tuple or int): Desired output size. If tuple, output is
matched to output_size. If int, smaller of image edges is matched
to output_size keeping aspect ratio the same.
"""
def __init__(self, output_size): def __init__(self, output_size):
assert isinstance(output_size, (int, tuple)) assert isinstance(output_size, (int, tuple))
self.output_size = output_size self.output_size = output_size
@ -131,6 +123,7 @@ class Normalize(object):
image = image.type(torch.FloatTensor) image = image.type(torch.FloatTensor)
image = F.normalize(image, self.mean, self.std, self.inplace) image = F.normalize(image, self.mean, self.std, self.inplace)
sample['image'] = image sample['image'] = image
sample['gt_matte'] = sample['gt_matte'] / 255
return sample return sample