MCPcopy Create free account
hub / github.com/SooLab/CGFormer / __getitem__

Method __getitem__

utils/dataset.py:117–162  ·  view source on GitHub ↗
(self, index)

Source from the content-addressed store, hash-verified

115 return self.length
116
117 def __getitem__(self, index):
118 # Delay loading LMDB data until after initialization: https://github.com/chainer/chainermn/issues/129
119 if self.env is None:
120 self._init_db()
121 env = self.env
122 with env.begin(write=False) as txn:
123 byteflow = txn.get(self.keys[index])
124 ref = loads_pyarrow(byteflow)
125 # img
126 ori_img = cv2.imdecode(np.frombuffer(ref['img'], np.uint8),
127 cv2.IMREAD_COLOR)
128 img = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB)
129 img_size = img.shape[:2]
130 # mask
131 seg_id = ref['seg_id']
132 mask_dir = os.path.join(self.mask_dir, str(seg_id) + '.png')
133 # sentences
134 idx = np.random.choice(ref['num_sents'])
135 sents = ref['sents']
136 # transform
137 # mask transform
138 mask = cv2.imdecode(np.frombuffer(ref['mask'], np.uint8),
139 cv2.IMREAD_GRAYSCALE)
140 mask = mask / 255.
141 if self.mode == 'train':
142 sent = sents[idx]
143 # sentence -> vector
144 img, mask, sent = self.convert(img, mask, sent, inference=False)
145 word_vec, pad_mask = tokenize(sent, self.word_length, True)
146 return img, word_vec, mask, pad_mask
147 elif self.mode == 'val':
148 # sentence -> vector
149 sent = sents[-1]
150 word_vec, pad_mask = tokenize(sent, self.word_length, True)
151 img, mask, sent = self.convert(img, mask, sent, inference=False)
152 return img, word_vec, mask, pad_mask
153 else:
154 # sentence -> vector
155 word_vecs = []
156 pad_masks = []
157 for sent in sents:
158 word_vec, pad_mask = tokenize(sent, self.word_length, True)
159 word_vecs.append(word_vec)
160 pad_masks.append(pad_mask)
161 img, mask, sent = self.convert(img, mask, sent, inference=True)
162 return ori_img, img, word_vecs, mask, pad_masks, seg_id, sents,
163
164 def convert(self, img, mask, sent, inference=False):
165 img = Image.fromarray(np.uint8(img))

Callers

nothing calls this directly

Calls 4

_init_dbMethod · 0.95
convertMethod · 0.95
loads_pyarrowFunction · 0.70
tokenizeFunction · 0.70

Tested by

no test coverage detected