import os import os.path as op import json # import logging import base64 import yaml import errno import io import math from PIL import Image, ImageDraw from maskrcnn_benchmark.structures.bounding_box import BoxList from .box_label_loader import LabelLoader def load_linelist_file(linelist_file): if linelist_file is not None: line_list = [] with open(linelist_file, 'r') as fp: for i in fp: line_list.append(int(i.strip())) return line_list def img_from_base64(imagestring): try: img = Image.open(io.BytesIO(base64.b64decode(imagestring))) return img.convert('RGB') except ValueError: return None def load_from_yaml_file(yaml_file): with open(yaml_file, 'r') as fp: return yaml.load(fp, Loader=yaml.CLoader) def find_file_path_in_yaml(fname, root): if fname is not None: if op.isfile(fname): return fname elif op.isfile(op.join(root, fname)): return op.join(root, fname) else: raise FileNotFoundError( errno.ENOENT, os.strerror(errno.ENOENT), op.join(root, fname) ) def create_lineidx(filein, idxout): idxout_tmp = idxout + '.tmp' with open(filein, 'r') as tsvin, open(idxout_tmp, 'w') as tsvout: fsize = os.fstat(tsvin.fileno()).st_size fpos = 0 while fpos != fsize: tsvout.write(str(fpos) + "\n") tsvin.readline() fpos = tsvin.tell() os.rename(idxout_tmp, idxout) def read_to_character(fp, c): result = [] while True: s = fp.read(32) assert s != '' if c in s: result.append(s[: s.index(c)]) break else: result.append(s) return ''.join(result) class TSVFile(object): def __init__(self, tsv_file, generate_lineidx=False): self.tsv_file = tsv_file self.lineidx = op.splitext(tsv_file)[0] + '.lineidx' self._fp = None self._lineidx = None # the process always keeps the process which opens the file. # If the pid is not equal to the currrent pid, we will re-open the file. self.pid = None # generate lineidx if not exist if not op.isfile(self.lineidx) and generate_lineidx: create_lineidx(self.tsv_file, self.lineidx) def __del__(self): if self._fp: self._fp.close() def __str__(self): return "TSVFile(tsv_file='{}')".format(self.tsv_file) def __repr__(self): return str(self) def num_rows(self): self._ensure_lineidx_loaded() return len(self._lineidx) def seek(self, idx): self._ensure_tsv_opened() self._ensure_lineidx_loaded() try: pos = self._lineidx[idx] except: # logging.info('{}-{}'.format(self.tsv_file, idx)) raise self._fp.seek(pos) return [s.strip() for s in self._fp.readline().split('\t')] def seek_first_column(self, idx): self._ensure_tsv_opened() self._ensure_lineidx_loaded() pos = self._lineidx[idx] self._fp.seek(pos) return read_to_character(self._fp, '\t') def get_key(self, idx): return self.seek_first_column(idx) def __getitem__(self, index): return self.seek(index) def __len__(self): return self.num_rows() def _ensure_lineidx_loaded(self): if self._lineidx is None: # logging.info('loading lineidx: {}'.format(self.lineidx)) with open(self.lineidx, 'r') as fp: self._lineidx = [int(i.strip()) for i in fp.readlines()] def _ensure_tsv_opened(self): if self._fp is None: self._fp = open(self.tsv_file, 'r') self.pid = os.getpid() if self.pid != os.getpid(): # logging.info('re-open {} because the process id changed'.format(self.tsv_file)) self._fp = open(self.tsv_file, 'r') self.pid = os.getpid() class CompositeTSVFile(): def __init__(self, file_list, seq_file, root='.'): if isinstance(file_list, str): self.file_list = load_list_file(file_list) else: assert isinstance(file_list, list) self.file_list = file_list self.seq_file = seq_file self.root = root self.initialized = False self.initialize() def get_key(self, index): idx_source, idx_row = self.seq[index] k = self.tsvs[idx_source].get_key(idx_row) return '_'.join([self.file_list[idx_source], k]) def num_rows(self): return len(self.seq) def __getitem__(self, index): idx_source, idx_row = self.seq[index] return self.tsvs[idx_source].seek(idx_row) def __len__(self): return len(self.seq) def initialize(self): ''' this function has to be called in init function if cache_policy is enabled. Thus, let's always call it in init funciton to make it simple. ''' if self.initialized: return self.seq = [] with open(self.seq_file, 'r') as fp: for line in fp: parts = line.strip().split('\t') self.seq.append([int(parts[0]), int(parts[1])]) self.tsvs = [TSVFile(op.join(self.root, f)) for f in self.file_list] self.initialized = True def load_list_file(fname): with open(fname, 'r') as fp: lines = fp.readlines() result = [line.strip() for line in lines] if len(result) > 0 and result[-1] == '': result = result[:-1] return result class TSVDataset(object): def __init__(self, img_file, label_file=None, hw_file=None, linelist_file=None, imageid2idx_file=None): """Constructor. Args: img_file: Image file with image key and base64 encoded image str. label_file: An optional label file with image key and label information. A label_file is required for training and optional for testing. hw_file: An optional file with image key and image height/width info. linelist_file: An optional file with a list of line indexes to load samples. It is useful to select a subset of samples or duplicate samples. """ self.img_file = img_file self.label_file = label_file self.hw_file = hw_file self.linelist_file = linelist_file self.img_tsv = TSVFile(img_file) self.label_tsv = None if label_file is None else TSVFile(label_file, generate_lineidx=True) self.hw_tsv = None if hw_file is None else TSVFile(hw_file) self.line_list = load_linelist_file(linelist_file) self.imageid2idx = None if imageid2idx_file is not None: self.imageid2idx = json.load(open(imageid2idx_file, 'r')) self.transforms = None def __len__(self): if self.line_list is None: if self.imageid2idx is not None: assert self.label_tsv is not None, "label_tsv is None!!!" return self.label_tsv.num_rows() return self.img_tsv.num_rows() else: return len(self.line_list) def __getitem__(self, idx): img = self.get_image(idx) img_size = img.size # w, h annotations = self.get_annotations(idx) # print(idx, annotations) target = self.get_target_from_annotations(annotations, img_size, idx) img, target = self.apply_transforms(img, target) if self.transforms is None: return img, target, idx, 1.0 else: new_img_size = img.shape[1:] scale = math.sqrt(float(new_img_size[0] * new_img_size[1]) / float(img_size[0] * img_size[1])) return img, target, idx, scale def get_line_no(self, idx): return idx if self.line_list is None else self.line_list[idx] def get_image(self, idx): line_no = self.get_line_no(idx) if self.imageid2idx is not None: assert self.label_tsv is not None, "label_tsv is None!!!" row = self.label_tsv.seek(line_no) annotations = json.loads(row[1]) imageid = annotations["img_id"] line_no = self.imageid2idx[imageid] row = self.img_tsv.seek(line_no) # use -1 to support old format with multiple columns. img = img_from_base64(row[-1]) return img def get_annotations(self, idx): line_no = self.get_line_no(idx) if self.label_tsv is not None: row = self.label_tsv.seek(line_no) annotations = json.loads(row[1]) return annotations else: return [] def get_target_from_annotations(self, annotations, img_size, idx): # This function will be overwritten by each dataset to # decode the labels to specific formats for each task. return annotations def apply_transforms(self, image, target=None): # This function will be overwritten by each dataset to # apply transforms to image and targets. return image, target def get_img_info(self, idx): if self.imageid2idx is not None: assert self.label_tsv is not None, "label_tsv is None!!!" line_no = self.get_line_no(idx) row = self.label_tsv.seek(line_no) annotations = json.loads(row[1]) return {"height": int(annotations["img_w"]), "width": int(annotations["img_w"])} if self.hw_tsv is not None: line_no = self.get_line_no(idx) row = self.hw_tsv.seek(line_no) try: # json string format with "height" and "width" being the keys data = json.loads(row[1]) if type(data) == list: return data[0] elif type(data) == dict: return data except ValueError: # list of strings representing height and width in order hw_str = row[1].split(' ') hw_dict = {"height": int(hw_str[0]), "width": int(hw_str[1])} return hw_dict def get_img_key(self, idx): line_no = self.get_line_no(idx) # based on the overhead of reading each row. if self.imageid2idx is not None: assert self.label_tsv is not None, "label_tsv is None!!!" row = self.label_tsv.seek(line_no) annotations = json.loads(row[1]) return annotations["img_id"] if self.hw_tsv: return self.hw_tsv.seek(line_no)[0] elif self.label_tsv: return self.label_tsv.seek(line_no)[0] else: return self.img_tsv.seek(line_no)[0] class TSVYamlDataset(TSVDataset): """ TSVDataset taking a Yaml file for easy function call """ def __init__(self, yaml_file, root=None, replace_clean_label=False): print("Reading {}".format(yaml_file)) self.cfg = load_from_yaml_file(yaml_file) if root: self.root = root else: self.root = op.dirname(yaml_file) img_file = find_file_path_in_yaml(self.cfg['img'], self.root) label_file = find_file_path_in_yaml(self.cfg.get('label', None), self.root) hw_file = find_file_path_in_yaml(self.cfg.get('hw', None), self.root) linelist_file = find_file_path_in_yaml(self.cfg.get('linelist', None), self.root) imageid2idx_file = find_file_path_in_yaml(self.cfg.get('imageid2idx', None), self.root) if replace_clean_label: assert ("raw_label" in label_file) label_file = label_file.replace("raw_label", "clean_label") super(TSVYamlDataset, self).__init__( img_file, label_file, hw_file, linelist_file, imageid2idx_file) class ODTSVDataset(TSVYamlDataset): """ Generic TSV dataset format for Object Detection. """ def __init__(self, yaml_file, extra_fields=(), transforms=None, is_load_label=True, **kwargs): if yaml_file is None: return super(ODTSVDataset, self).__init__(yaml_file) self.transforms = transforms self.is_load_label = is_load_label self.attribute_on = False # self.attribute_on = kwargs['args'].MODEL.ATTRIBUTE_ON if "args" in kwargs else False if self.is_load_label: # construct maps jsondict_file = find_file_path_in_yaml( self.cfg.get("labelmap", None), self.root ) if jsondict_file is None: jsondict_file = find_file_path_in_yaml( self.cfg.get("jsondict", None), self.root ) if "json" in jsondict_file: jsondict = json.load(open(jsondict_file, 'r')) if "label_to_idx" not in jsondict: jsondict = {'label_to_idx': jsondict} elif "tsv" in jsondict_file: label_to_idx = {} counter = 1 with open(jsondict_file) as f: for line in f: label_to_idx[line.strip()] = counter counter += 1 jsondict = {'label_to_idx': label_to_idx} else: assert (0) self.labelmap = {} self.class_to_ind = jsondict['label_to_idx'] self.class_to_ind['__background__'] = 0 self.ind_to_class = {v: k for k, v in self.class_to_ind.items()} self.labelmap['class_to_ind'] = self.class_to_ind if self.attribute_on: self.attribute_to_ind = jsondict['attribute_to_idx'] self.attribute_to_ind['__no_attribute__'] = 0 self.ind_to_attribute = {v: k for k, v in self.attribute_to_ind.items()} self.labelmap['attribute_to_ind'] = self.attribute_to_ind self.label_loader = LabelLoader( labelmap=self.labelmap, extra_fields=extra_fields, ) def get_target_from_annotations(self, annotations, img_size, idx): if isinstance(annotations, list): annotations = {"objects": annotations} if self.is_load_label: return self.label_loader(annotations['objects'], img_size) def apply_transforms(self, img, target=None): if self.transforms is not None: img, target = self.transforms(img, target) return img, target