-
Notifications
You must be signed in to change notification settings - Fork 6
Expand file tree
/
Copy pathIRDSTDataLoader.py
More file actions
111 lines (88 loc) · 3.79 KB
/
IRDSTDataLoader.py
File metadata and controls
111 lines (88 loc) · 3.79 KB
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
import os
import torch
from PIL import Image
from torch.utils.data import Dataset
from numpy import *
import numpy as np
import scipy.io as scio
import cv2
from match_images import matching
#load image
class IRDST_TrainSetLoader(Dataset):
def __init__(self, root, fullSupervision=True, align=True):
txtpath = root + 'img_idx/train_IRDST-simulation.txt'
txt = np.loadtxt(txtpath, dtype=bytes).astype(str)
self.imgs_arr = txt
self.root = root
self.fullSupervision = fullSupervision
self.align = align
self.train_mean = 106.8523
self.train_std = 56.9243
def __getitem__(self, index):
imgfolder = 'images'
img_path = os.path.join(self.root, imgfolder, self.imgs_arr[index] + '.bmp')
seq = self.imgs_arr[index].split('/')[0]
frame = int(self.imgs_arr[index].split('/')[1])
img_ori = cv2.imread(img_path)
if np.dim(img_ori) == 3:
img_ori = img_ori[:,:,0]
img = np.expand_dims(img_ori.astype(np.float32), axis=0)
for i in range(1,5):
img_hispath = os.path.join(self.root, imgfolder, seq, str(max(frame-i, 1)) + '.bmp')
img_his = cv2.imread(img_hispath)
if np.ndim(img_his) == 3:
img_his= img_his[:,:,0]
if self.align:
img_his = matching(img_his, img_ori)
img_his = np.expand_dims(img_his.astype(np.float32), axis=0)
img= np.concatenate((img_his, img), axis=0)
# Read Label/Tgt
label_path = img_path.replace(imgfolder, 'masks')
label = cv2.imread(label_path).astype(np.float32) / 255.0
# Mix preprocess
img = (img - self.train_mean)/self.train_std
img = torch.unsqueeze(torch.from_numpy(img), 0)
label = torch.unsqueeze(torch.from_numpy(label), 0)
[_, m, n] = np.shape(label)
return img, label, m, n
def __len__(self):
return len(self.imgs_arr)
class IRDST_TestSetLoader(Dataset):
def __init__(self, root, fullSupervision=True, align=True):
txtpath = root + 'img_idx/test_IRDST-simulation.txt'
txt = np.loadtxt(txtpath, dtype=bytes).astype(str)
self.imgs_arr = txt
self.root = root
self.fullSupervision = fullSupervision
self.align = align
self.train_mean = 106.8523
self.train_std = 56.9243
def __getitem__(self, index):
imgfolder = 'images'
img_path = os.path.join(self.root, imgfolder, self.imgs_arr[index] + '.bmp')
seq = self.imgs_arr[index].split('/')[0]
frame = int(self.imgs_arr[index].split('/')[1])
img_ori = cv2.imread(img_path)
if np.dim(img_ori) == 3:
img_ori = img_ori[:,:,0]
img = np.expand_dims(img_ori.astype(np.float32), axis=0)
for i in range(1,5):
img_hispath = os.path.join(self.root, imgfolder, seq, str(max(frame-i, 1)) + '.bmp')
img_his = cv2.imread(img_hispath)
if np.ndim(img_his) == 3:
img_his = img_his[:,:,0]
if self.align:
img_his = matching(img_his, img_ori)
img_his = np.expand_dims(img_his.astype(np.float32), axis=0)
img= np.concatenate((img_his, img), axis=0)
# Read Label/Tgt
label_path = img_path.replace(imgfolder, 'masks')
label = cv2.imread(label_path).astype(np.float32) / 255.0
# Mix preprocess
img = (img - self.train_mean)/self.train_std
img = torch.unsqueeze(torch.from_numpy(img), 0)
label = torch.unsqueeze(torch.from_numpy(label), 0)
[_, m, n] = np.shape(label)
return img, label, m, n
def __len__(self):
return len(self.imgs_arr)