Linhz/ViMNer
1
1import torch
2import logging
3import os
4
5logger = logging.getLogger(__name__)
6from torchvision import transforms
7from PIL import Image
8
9
10class SBInputExample(object):
11 """A single training/test example for simple sequence classification."""
12
13 def __init__(self, guid, text_a, text_b, img_id, label=None, auxlabel=None):
14 """Constructs a InputExample.
15
16 Args:
17 guid: Unique id for the example.
18 text_a: string. The untokenized text of the first sequence. For single
19 sequence tasks, only this sequence must be specified.
20 text_b: (Optional) string. The untokenized text of the second sequence.
21 Only must be specified for sequence pair tasks.
22 label: (Optional) string. The label of the example. This should be
23 specified for train and dev examples, but not for test examples.
24 """
25 self.guid = guid
26 self.text_a = text_a
27 self.text_b = text_b
28 self.img_id = img_id
29 self.label = label
30 # Please note that the auxlabel is not used in SB
31 # it is just kept in order not to modify the original code
32 self.auxlabel = auxlabel
33
34
35class SBInputFeatures(object):
36 """A single set of features of data"""
37
38 def __init__(self, input_ids, input_mask, added_input_mask, segment_ids, img_feat, label_id, auxlabel_id):
39 self.input_ids = input_ids
40 self.input_mask = input_mask
41 self.added_input_mask = added_input_mask
42 self.segment_ids = segment_ids
43 self.img_feat = img_feat
44 self.label_id = label_id
45 self.auxlabel_id = auxlabel_id
46
47
48def sbreadfile(filename):
49 '''
50 Đọc dữ liệu từ tệp và trả về dưới dạng danh sách các cặp từ và nhãn, cùng với danh sách hình ảnh và nhãn phụ.
51 '''
52 print("Chuẩn bị dữ liệu cho ", filename)
53 f = open(filename, encoding='utf8')
54 data = []
55 imgs = []
56 auxlabels = []
57 sentence = []
58 label = []
59 auxlabel = []
60 imgid = ''
61
62 for line in f:
63 line = line.strip() # Loại bỏ các dấu cách thừa ở đầu và cuối dòng
64 if line.startswith('IMGID:'):
65 imgid = line.split('IMGID:')[1] + '.jpg'
66 continue
67 if line == '':
68 if len(sentence) > 0:
69 data.append((sentence, label))
70 imgs.append(imgid)
71 auxlabels.append(auxlabel)
72 sentence = []
73 label = []
74 auxlabel = []
75 imgid = ''
76 continue
77 splits = line.split('\t')
78 if len(splits) == 2: # Đảm bảo dòng có ít nhất một từ và một nhãn
79 word, cur_label = splits
80 sentence.append(word)
81 label.append(cur_label)
82 auxlabel.append(cur_label[0]) # Lấy ký tự đầu tiên của nhãn làm nhãn phụ
83
84 if len(sentence) > 0: # Xử lý dữ liệu cuối cùng trong tệp
85 data.append((sentence, label))
86 imgs.append(imgid)
87 auxlabels.append(auxlabel)
88
89 print("Số lượng mẫu: " + str(len(data)))
90 print("Số lượng hình ảnh: " + str(len(imgs)))
91 return data, imgs, auxlabels
92
93
94# def sbreadfile(filename): #code gốc
95# '''
96# read file
97# return format :
98# [ ['EU', 'B-ORG'], ['rejects', 'O'], ['German', 'B-MISC'], ['call', 'O'], ['to', 'O'], ['boycott', 'O'], ['British', 'B-MISC'], ['lamb', 'O'], ['.', 'O'] ]
99# '''
100# print("prepare data for ",filename)
101# f = open(filename,encoding='utf8')
102# data = []
103# imgs = []
104# auxlabels = []
105# sentence = []
106# label = []
107# auxlabel = []
108# imgid = ''
109# a = 0
110# for line in f:
111# if line.startswith('IMGID:'):
112# imgid = line.strip().split('IMGID:')[1] + '.jpg'
113# continue
114# if line[0] == "\n":
115# if len(sentence) > 0:
116# data.append((sentence, label))
117# imgs.append(imgid)
118# auxlabels.append(auxlabel)
119# sentence = []
120# label = []
121# imgid = ''
122# auxlabel = []
123# continue
124# splits = line.split('\t')
125# sentence.append(splits[0])
126# cur_label = splits[-1][:-1]
127# # if cur_label == 'B-OTHER':
128# # cur_label = 'B-MISC'
129# # elif cur_label == 'I-OTHER':
130# # cur_label = 'I-MISC'
131# label.append(cur_label)
132# auxlabel.append(cur_label[0])
133
134# if len(sentence) > 0:
135# data.append((sentence, label))
136# imgs.append(imgid)
137# auxlabels.append(auxlabel)
138# sentence = []
139# label = []
140# auxlabel = []
141
142# print("The number of samples: " + str(len(data)))
143# print("The number of images: " + str(len(imgs)))
144# return data, imgs, auxlabels
145
146class DataProcessor(object):
147 """Base class for data converters for sequence classification data sets."""
148
149 def get_train_examples(self, data_dir):
150 """Gets a collection of `InputExample`s for the train set."""
151 raise NotImplementedError()
152
153 def get_dev_examples(self, data_dir):
154 """Gets a collection of `InputExample`s for the dev set."""
155 raise NotImplementedError()
156
157 def get_labels(self):
158 """Gets the list of labels for this data set."""
159 raise NotImplementedError()
160
161 @classmethod
162 def _read_sbtsv(cls, input_file, quotechar=None):
163 """Reads a tab separated value file."""
164 return sbreadfile(input_file)
165
166
167class MNERProcessor_2021(DataProcessor):
168 """Processor for the CoNLL-2003 data set."""
169
170 def get_train_examples(self, data_dir):
171 """See base class."""
172 data, imgs, auxlabels = self._read_sbtsv(os.path.join(data_dir, "train.txt"))
173 return self._create_examples(data, imgs, auxlabels, "train")
174
175 def get_dev_examples(self, data_dir):
176 """See base class."""
177 data, imgs, auxlabels = self._read_sbtsv(os.path.join(data_dir, "dev.txt"))
178 return self._create_examples(data, imgs, auxlabels, "dev")
179
180 def get_test_examples(self, data_dir):
181 """See base class."""
182 data, imgs, auxlabels = self._read_sbtsv(os.path.join(data_dir, "test.txt"))
183 return self._create_examples(data, imgs, auxlabels, "test")
184
185 def get_labels(self):
186 return [
187 "O", # 1
188 "I-PRODUCT-AWARD", # 2
189 "B-MISCELLANEOUS", # 3
190 "B-QUANTITY-NUM", # 4
191 "B-ORGANIZATION-SPORTS", # 5
192 "B-DATETIME", # 6
193 "I-ADDRESS", # 7
194 "I-PERSON", # 8
195 "I-EVENT-SPORT", # 9
196 "B-ADDRESS", # 10
197 "B-EVENT-NATURAL", # 11
198 "I-LOCATION-GPE", # 12
199 "B-EVENT-GAMESHOW", # 13
200 "B-DATETIME-TIMERANGE", # 14
201 "I-QUANTITY-NUM", # 15
202 "I-QUANTITY-AGE", # 16
203 "B-EVENT-CUL", # 17
204 "I-QUANTITY-TEM", # 18
205 "I-PRODUCT-LEGAL", # 19
206 "I-LOCATION-STRUC", # 20
207 "I-ORGANIZATION", # 21
208 "B-PHONENUMBER", # 22
209 "B-IP", # 23
210 "B-QUANTITY-AGE", # 24
211 "I-DATETIME-TIME", # 25
212 "I-DATETIME", # 26
213 "B-ORGANIZATION-MED", # 27
214 "B-DATETIME-SET", # 28
215 "I-EVENT-CUL", # 29
216 "B-QUANTITY-DIM", # 30
217 "I-QUANTITY-DIM", # 31
218 "B-EVENT", # 32
219 "B-DATETIME-DATERANGE", # 33
220 "I-EVENT-GAMESHOW", # 34
221 "B-PRODUCT-AWARD", # 35
222 "B-LOCATION-STRUC", # 36
223 "B-LOCATION", # 37
224 "B-PRODUCT", # 38
225 "I-MISCELLANEOUS", # 39
226 "B-SKILL", # 40
227 "I-QUANTITY-ORD", # 41
228 "I-ORGANIZATION-STOCK", # 42
229 "I-LOCATION-GEO", # 43
230 "B-PERSON", # 44
231 "B-PRODUCT-COM", # 45
232 "B-PRODUCT-LEGAL", # 46
233 "I-LOCATION", # 47
234 "B-QUANTITY-TEM", # 48
235 "I-PRODUCT", # 49
236 "B-QUANTITY-CUR", # 50
237 "I-QUANTITY-CUR", # 51
238 "B-LOCATION-GPE", # 52
239 "I-PHONENUMBER", # 53
240 "I-ORGANIZATION-MED", # 54
241 "I-EVENT-NATURAL", # 55
242 "I-EMAIL", # 56
243 "B-ORGANIZATION", # 57
244 "B-URL", # 58
245 "I-DATETIME-TIMERANGE", # 59
246 "I-QUANTITY", # 60
247 "I-IP", # 61
248 "B-EVENT-SPORT", # 62
249 "B-PERSONTYPE", # 63
250 "B-QUANTITY-PER", # 64
251 "I-QUANTITY-PER", # 65
252 "I-PRODUCT-COM", # 66
253 "I-DATETIME-DURATION", # 67
254 "B-LOCATION-GPE-GEO", # 68
255 "B-QUANTITY-ORD", # 69
256 "I-EVENT", # 70
257 "B-DATETIME-TIME", # 71
258 "B-QUANTITY", # 72
259 "I-DATETIME-SET", # 73
260 "I-LOCATION-GPE-GEO", # 74
261 "B-ORGANIZATION-STOCK", # 75
262 "I-ORGANIZATION-SPORTS", # 76
263 "I-SKILL", # 77
264 "I-URL", # 78
265 "B-DATETIME-DURATION", # 79
266 "I-DATETIME-DATE", # 80
267 "I-PERSONTYPE", # 81
268 "B-DATETIME-DATE", # 82
269 "I-DATETIME-DATERANGE", # 83
270 "B-LOCATION-GEO", # 84
271 "B-EMAIL", # 85
272 "X", # 86
273 "<s>", # 87
274 "</s>" # 88
275 ]
276
277 # vlsp2016
278
279
280 # vlsp2018
281 # return [
282 # "O","I-ORGANIZATION",
283 # "B-ORGANIZATION",
284 # "I-LOCATION",
285 # "B-MISCELLANEOUS",
286 # "I-PERSON",
287 # "B-PERSON",
288 # "I-MISCELLANEOUS",
289 # "B-LOCATION",
290 # "X",
291 # "<s>",
292 # "</s>"]
293
294 def get_auxlabels(self):
295 return ["O", "B", "I", "X", "<s>", "</s>"]
296
297 def get_start_label_id(self):
298 label_list = self.get_labels()
299 label_map = {label: i for i, label in enumerate(label_list, 1)}
300 return label_map['<s>']
301
302 def get_stop_label_id(self):
303 label_list = self.get_labels()
304 label_map = {label: i for i, label in enumerate(label_list, 1)}
305 return label_map['</s>']
306
307 def _create_examples(self, lines, imgs, auxlabels, set_type):
308 examples = []
309 for i, (sentence, label) in enumerate(lines):
310 guid = "%s-%s" % (set_type, i)
311 text_a = ' '.join(sentence)
312 text_b = None
313 img_id = imgs[i]
314 label = label
315 auxlabel = auxlabels[i]
316 examples.append(
317 SBInputExample(guid=guid, text_a=text_a, text_b=text_b, img_id=img_id, label=label, auxlabel=auxlabel))
318 return examples
319
320
321def image_process(image_path, transform):
322 image = Image.open(image_path).convert('RGB')
323 image = transform(image)
324 return image
325
326
327def convert_mm_examples_to_features(examples, label_list, auxlabel_list,
328 max_seq_length, tokenizer, crop_size, path_img):
329 label_map = {label: i for i, label in enumerate(label_list, 1)}
330 auxlabel_map = {label: i for i, label in enumerate(auxlabel_list, 1)}
331
332 features = []
333 count = 0
334
335 transform = transforms.Compose([
336 transforms.Resize([256, 256]),
337 transforms.RandomCrop(crop_size), # args.crop_size, by default it is set to be 224
338 transforms.RandomHorizontalFlip(),
339 transforms.ToTensor(),
340 transforms.Normalize((0.485, 0.456, 0.406),
341 (0.229, 0.224, 0.225))])
342
343 for (ex_index, example) in enumerate(examples):
344 textlist = example.text_a.split(' ')
345 labellist = example.label
346 auxlabellist = example.auxlabel
347 tokens = []
348 labels = []
349 auxlabels = []
350 for i, word in enumerate(textlist):
351 token = tokenizer.tokenize(word)
352 tokens.extend(token)
353 label_1 = labellist[i]
354 auxlabel_1 = auxlabellist[i]
355 for m in range(len(token)):
356 if m == 0:
357 labels.append(label_1)
358 auxlabels.append(auxlabel_1)
359 else:
360 labels.append("X")
361 auxlabels.append("X")
362 if len(tokens) >= max_seq_length - 1:
363 tokens = tokens[0:(max_seq_length - 2)]
364 labels = labels[0:(max_seq_length - 2)]
365 auxlabels = auxlabels[0:(max_seq_length - 2)]
366 ntokens = []
367 segment_ids = []
368 label_ids = []
369 auxlabel_ids = []
370 ntokens.append("<s>")
371 segment_ids.append(0)
372 label_ids.append(label_map["<s>"])
373 auxlabel_ids.append(auxlabel_map["<s>"])
374 for i, token in enumerate(tokens):
375 ntokens.append(token)
376 segment_ids.append(0)
377 label_ids.append(label_map[labels[i]])
378 auxlabel_ids.append(auxlabel_map[auxlabels[i]])
379 ntokens.append("</s>")
380 segment_ids.append(0)
381 label_ids.append(label_map["</s>"])
382 auxlabel_ids.append(auxlabel_map["</s>"])
383 input_ids = tokenizer.convert_tokens_to_ids(ntokens)
384 input_mask = [1] * len(input_ids)
385 added_input_mask = [1] * (len(input_ids) + 49) # 1 or 49 is for encoding regional image representations
386
387 while len(input_ids) < max_seq_length:
388 input_ids.append(0)
389 input_mask.append(0)
390 added_input_mask.append(0)
391 segment_ids.append(0)
392 label_ids.append(0)
393 auxlabel_ids.append(0)
394
395 assert len(input_ids) == max_seq_length
396 assert len(input_mask) == max_seq_length
397 assert len(segment_ids) == max_seq_length
398 assert len(label_ids) == max_seq_length
399 assert len(auxlabel_ids) == max_seq_length
400
401 image_name = example.img_id
402 image_path = os.path.join(path_img, image_name)
403
404 if not os.path.exists(image_path):
405 if 'NaN' not in image_path:
406 print(image_path)
407 try:
408 image = image_process(image_path, transform)
409 except:
410 count += 1
411 image_path_fail = os.path.join(path_img, 'background.jpg')
412 image = image_process(image_path_fail, transform)
413
414 else:
415 if ex_index < 2:
416 logger.info("*** Example ***")
417 logger.info("guid: %s" % (example.guid))
418 logger.info("tokens: %s" % " ".join(
419 [str(x) for x in tokens]))
420 logger.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
421 logger.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
422 logger.info(
423 "segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
424 logger.info("label: %s" % " ".join([str(x) for x in label_ids]))
425 logger.info("auxlabel: %s" % " ".join([str(x) for x in auxlabel_ids]))
426
427 features.append(
428 SBInputFeatures(input_ids=input_ids, input_mask=input_mask, added_input_mask=added_input_mask,
429 segment_ids=segment_ids, img_feat=image, label_id=label_ids, auxlabel_id=auxlabel_ids))
430
431 print('the number of problematic samples: ' + str(count))
432 return features
433
434
435# if __name__ == "__main__":
436# processor = MNERProcessor_2016()
437# label_list = processor.get_labels()
438# auxlabel_list = processor.get_auxlabels()
439# num_labels = len(label_list) + 1 # label 0 corresponds to padding, label in label_list starts from 1
440#
441# start_label_id = processor.get_start_label_id()
442# stop_label_id = processor.get_stop_label_id()
443#
444# data_dir = r'sample_data'
445# train_examples = processor.get_train_examples(data_dir)
446# print(train_examples[0].img_id)