MCPcopy Create free account
hub / github.com/drinkingcoder/NeuralMarker / __getitem__

Method __getitem__

core/datasets.py:43–126  ·  view source on GitHub ↗
(self, index)

Source from the content-addressed store, hash-verified

41 self.bg_dataset = SequenceDataset(args, 'train')
42
43 def __getitem__(self, index):
44
45 if not self.init_seed:
46 worker_info = torch.utils.data.get_worker_info()
47 if worker_info is not None:
48 torch.manual_seed(worker_info.id)
49 np.random.seed(worker_info.id)
50 random.seed(worker_info.id)
51 self.init_seed = True
52
53 Im1, Im2 = self.entrydata_list[index]
54
55 im1 = cv2.imread(Im1.im_path)
56 im2 = cv2.imread(Im2.im_path)
57
58 if im1.ndim == 2:
59 im1 = np.tile(im1[..., None], (1, 1, 3))
60 im2 = np.tile(im2[..., None], (1, 1, 3))
61 else:
62 im1 = im1[..., :3]
63 im2 = im2[..., :3]
64
65 '''
66 crop the image
67 '''
68 ht, wd = im1.shape[:2]
69 crop_size = self.args.image_size
70
71 assert (crop_size[1] % 8 == 0 and crop_size[0] % 8 == 0)
72
73 im1 = cv2.resize(im1, (crop_size[1], crop_size[0]), interpolation=cv2.INTER_LINEAR)
74 im2 = cv2.resize(im2, (crop_size[1], crop_size[0]), interpolation=cv2.INTER_LINEAR)
75
76 scale = np.diag([crop_size[1] / wd, crop_size[0] / ht, 1.0])
77
78 K1s = scale.dot(Im1.K)
79 K2s = scale.dot(Im2.K)
80
81 F = fundamental_matrix_gen(Im1.Tcw, Im2.Tcw, K1s, K2s)
82
83 im1 = torch.from_numpy(im1).permute(2, 0, 1).float()
84 im2 = torch.from_numpy(im2).permute(2, 0, 1).float()
85
86 '''
87 transformation
88 '''
89
90 im3, backward_map, tnf_type, theta = self.transformation(im2, self.args.tnf_type)
91
92 backward_map = backward_map.permute([2, 0, 1])
93 backward_mask = (backward_map[0, :, :] >= 0) & (backward_map[0, :, :] < wd) & \
94 (backward_map[1, :, :] >= 0) & (backward_map[1, :, :] < ht)
95
96 coords1 = coords_grid(1, ht, wd, device='cpu').contiguous()
97 backward_flow = backward_map[None] - coords1
98
99 out = self.forward_warping(coords1.cuda(), backward_flow.cuda())
100 forward_map = out[0][0] / (out[1][0] + (1e-6))

Callers

nothing calls this directly

Calls 3

fundamental_matrix_genFunction · 0.90
coords_gridFunction · 0.90
random_selectMethod · 0.80

Tested by

no test coverage detected