forked from fh2019ustc/DocTr
-
Notifications
You must be signed in to change notification settings - Fork 1
/
dataset.py
61 lines (52 loc) · 2.13 KB
/
dataset.py
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
from __future__ import print_function, division
import os
import torch
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import hdf5storage as h5
from PIL import Image
import cv2
class GeometricDataset(Dataset):
def __init__(self, root_dir, segmentation_model, transform=None):
mapping_paths = []
image_paths = []
paths = os.path.join(root_dir, "image")
paths = os.listdir(paths)
for path in paths:
name = path.split(".")[0]
mapping_name = f"{name}.mat"
image_path = os.path.join(root_dir, "image", path)
mapping_path = os.path.join(root_dir, "mapping", mapping_name)
assert os.path.isfile(mapping_path) , f"no mapping file {mapping_name}"
image_paths.append(image_path)
mapping_paths.append(mapping_path)
self.root_dir = root_dir
self.mapping_dir = mapping_paths
self.image_dir = image_paths
self.transform = transform
self.segmentation_model = segmentation_model
def __len__(self):
return len(self.image_dir)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
backward_mapping = h5.loadmat(self.mapping_dir[idx])['bm']
backward_mapping = backward_mapping/ np.array([288, 288])
backward_mapping = (backward_mapping-0.5)*2
# bm0 = cv2.resize(backward_mapping[:, :, 0], (288, 288)) # x flow
# bm1 = cv2.resize(backward_mapping[:, :, 1], (288, 288)) # y flow
# bm0 = cv2.blur(bm0, (3, 3))
# bm1 = cv2.blur(bm1, (3, 3))
# backward_mapping = np.stack([bm0, bm1], axis=2)
backward_mapping = backward_mapping.transpose(2, 0, 1)
image = np.array(Image.open(self.image_dir[idx]))[:, :, :3]/255
print(image.shape)
# image = cv2.resize(image, (288, 288))
image = image.transpose(2, 0, 1)
sample = (image, backward_mapping)
if self.transform:
sample = self.transform(sample)
return sample