| import torch
|
| import logging
|
| import os
|
|
|
| logger = logging.getLogger(__name__)
|
| from torchvision import transforms
|
| from PIL import Image
|
|
|
|
|
| class SBInputExample(object):
|
| """A single training/test example for simple sequence classification."""
|
|
|
| def __init__(self, guid, text_a, text_b, img_id, label=None, auxlabel=None):
|
| """Constructs a InputExample.
|
|
|
| Args:
|
| guid: Unique id for the example.
|
| text_a: string. The untokenized text of the first sequence. For single
|
| sequence tasks, only this sequence must be specified.
|
| text_b: (Optional) string. The untokenized text of the second sequence.
|
| Only must be specified for sequence pair tasks.
|
| label: (Optional) string. The label of the example. This should be
|
| specified for train and dev examples, but not for test examples.
|
| """
|
| self.guid = guid
|
| self.text_a = text_a
|
| self.text_b = text_b
|
| self.img_id = img_id
|
|
|
|
|
|
|
|
|
| class SBInputFeatures(object):
|
| """A single set of features of data"""
|
|
|
| def __init__(self, input_ids, input_mask, added_input_mask, segment_ids, img_feat):
|
| self.input_ids = input_ids
|
| self.input_mask = input_mask
|
| self.added_input_mask = added_input_mask
|
| self.segment_ids = segment_ids
|
| self.img_feat = img_feat
|
|
|
|
|
| def sbreadfile(filename):
|
| '''
|
| Đọc dữ liệu từ tệp và trả về dưới dạng danh sách các từ và danh sách hình ảnh.
|
| '''
|
| print("Chuẩn bị dữ liệu từ", filename)
|
| with open(filename, encoding='utf8') as f:
|
| data = []
|
| imgs = []
|
| sentence = []
|
| imgid = ''
|
| for line in f:
|
| line = line.strip()
|
| if line.startswith('IMGID:'):
|
| imgid = line.split('IMGID:')[1] + '.jpg'
|
| continue
|
| if line == '':
|
| if len(sentence) > 0:
|
| data.append(sentence)
|
| imgs.append(imgid)
|
| sentence = []
|
| imgid = ''
|
| continue
|
| word = line.split('\t')[0]
|
| sentence.append(word)
|
|
|
| if len(sentence) > 0:
|
| data.append(sentence)
|
| imgs.append(imgid)
|
|
|
| print("Số lượng mẫu: " + str(len(data)))
|
| print("Số lượng hình ảnh: " + str(len(imgs)))
|
| return data, imgs
|
|
|
|
|
| class DataProcessor(object):
|
| """Base class for data converters for sequence classification data sets."""
|
|
|
| def get_train_examples(self, data_dir):
|
| """Gets a collection of `InputExample`s for the train set."""
|
| raise NotImplementedError()
|
|
|
| def get_dev_examples(self, data_dir):
|
| """Gets a collection of `InputExample`s for the dev set."""
|
| raise NotImplementedError()
|
|
|
| def get_labels(self):
|
| """Gets the list of labels for this data set."""
|
| raise NotImplementedError()
|
|
|
| @classmethod
|
| def _read_sbtsv(cls, input_file, quotechar=None):
|
| """Reads a tab separated value file."""
|
| return sbreadfile(input_file)
|
|
|
|
|
| class MNERProcessor(DataProcessor):
|
| """Processor for the CoNLL-2003 data set."""
|
|
|
| def get_train_examples(self, data_dir):
|
| """See base class."""
|
| data, imgs = self._read_sbtsv(os.path.join(data_dir, "train.txt"))
|
| return self._create_examples(data, imgs, "train")
|
|
|
| def get_dev_examples(self, data_dir):
|
| """See base class."""
|
| data, imgs = self._read_sbtsv(os.path.join(data_dir, "dev.txt"))
|
| return self._create_examples(data, imgs, "dev")
|
|
|
| def get_test_examples(self, data_dir):
|
| """See base class."""
|
| data, imgs = self._read_sbtsv(os.path.join(data_dir, "test.txt"))
|
| return self._create_examples(data, imgs, "test")
|
|
|
| def get_labels(self):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| return [
|
| "I-LOC", "B-MISC",
|
| "I-PER",
|
| "I-ORG",
|
| "B-LOC",
|
| "I-MISC",
|
| "B-ORG",
|
| "O",
|
| "B-PER",
|
| "X",
|
| "<s>",
|
| "</s>"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| def get_auxlabels(self):
|
| return ["O", "B", "I", "X", "<s>", "</s>"]
|
|
|
| def get_start_label_id(self):
|
| label_list = self.get_labels()
|
| label_map = {label: i for i, label in enumerate(label_list, 1)}
|
| return label_map['<s>']
|
|
|
| def get_stop_label_id(self):
|
| label_list = self.get_labels()
|
| label_map = {label: i for i, label in enumerate(label_list, 1)}
|
| return label_map['</s>']
|
|
|
| def _create_examples(self, lines, imgs, set_type):
|
| examples = []
|
| for i, (sentence) in enumerate(lines):
|
| guid = "%s-%s" % (set_type, i)
|
| text_a = ' '.join(sentence)
|
| text_b = None
|
| img_id = imgs[i]
|
| examples.append(
|
| SBInputExample(guid=guid, text_a=text_a, text_b=text_b, img_id=img_id))
|
| return examples
|
|
|
|
|
| def create_examples(lines, imgs, set_type):
|
| examples = []
|
| for i, (sentence) in enumerate(lines):
|
| guid = "%s-%s" % (set_type, i)
|
| text_a = ' '.join(sentence)
|
| text_b = None
|
| img_id = imgs[i]
|
| examples.append(
|
| SBInputExample(guid=guid, text_a=text_a, text_b=text_b, img_id=img_id))
|
| return examples
|
|
|
|
|
| def get_test_examples_predict(data_dir):
|
| """See base class."""
|
| data, imgs = sbreadfile(os.path.join(data_dir, "test.txt"))
|
| return create_examples(data, imgs, "test")
|
|
|
|
|
| def image_process(image_path, transform):
|
| image = Image.open(image_path).convert('RGB')
|
| image = transform(image)
|
| return image
|
|
|
|
|
| def convert_mm_examples_to_features_predict(examples,
|
| max_seq_length, tokenizer, crop_size, path_img):
|
| features = []
|
| count = 0
|
|
|
| transform = transforms.Compose([
|
| transforms.Resize([256, 256]),
|
| transforms.RandomCrop(crop_size),
|
| transforms.RandomHorizontalFlip(),
|
| transforms.ToTensor(),
|
| transforms.Normalize((0.485, 0.456, 0.406),
|
| (0.229, 0.224, 0.225))])
|
|
|
| for (ex_index, example) in enumerate(examples):
|
| textlist = example.text_a.split(' ')
|
| tokens = []
|
| for i, word in enumerate(textlist):
|
| token = tokenizer.tokenize(word)
|
| tokens.extend(token)
|
| if len(tokens) >= max_seq_length - 1:
|
| tokens = tokens[0:(max_seq_length - 2)]
|
| ntokens = []
|
| segment_ids = []
|
| ntokens.append("<s>")
|
| segment_ids.append(0)
|
| for i, token in enumerate(tokens):
|
| ntokens.append(token)
|
| segment_ids.append(0)
|
| ntokens.append("</s>")
|
| segment_ids.append(0)
|
| input_ids = tokenizer.convert_tokens_to_ids(ntokens)
|
| input_mask = [1] * len(input_ids)
|
| added_input_mask = [1] * (len(input_ids) + 49)
|
|
|
| while len(input_ids) < max_seq_length:
|
| input_ids.append(0)
|
| input_mask.append(0)
|
| added_input_mask.append(0)
|
| segment_ids.append(0)
|
|
|
| assert len(input_ids) == max_seq_length
|
| assert len(input_mask) == max_seq_length
|
| assert len(segment_ids) == max_seq_length
|
|
|
| image_name = example.img_id
|
| image_path = os.path.join(path_img, image_name)
|
|
|
| if not os.path.exists(image_path):
|
| if 'NaN' not in image_path:
|
| print(image_path)
|
| try:
|
| image = image_process(image_path, transform)
|
| except:
|
| count += 1
|
| image_path_fail = os.path.join(path_img, 'background.jpg')
|
| image = image_process(image_path_fail, transform)
|
|
|
| else:
|
| if ex_index < 1:
|
| logger.info("*** Example ***")
|
| logger.info("guid: %s" % (example.guid))
|
| logger.info("tokens: %s" % " ".join(
|
| [str(x) for x in tokens]))
|
| logger.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
|
| logger.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
|
| logger.info(
|
| "segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
|
|
|
| features.append(
|
| SBInputFeatures(input_ids=input_ids, input_mask=input_mask, added_input_mask=added_input_mask,
|
| segment_ids=segment_ids, img_feat=image))
|
|
|
| print('the number of problematic samples: ' + str(count))
|
| return features |
