import argparse import json import math import os import pprint import string from collections import Counter from math import floor import numpy as np from tqdm import tqdm from difflib import SequenceMatcher from detectron2.data.detection_utils import read_image def preprocess_object_labels(data, alias_dict={}): for img in data: for obj in img["objects"]: obj["ids"] = [obj["object_id"]] names = [] for name in obj["names"]: name = name.lower() label = sentence_preprocess(name) if label in alias_dict: label = alias_dict[label] names.append(label) obj["names"] = names def extract_object_token(data, num_tokens, object_list=[], verbose=True): """Builds a set that contains the object names. Filters infrequent tokens.""" token_counter = Counter() for img in data: for obj in img["objects"]: for name in obj["names"]: if len(name): pass else: continue if not object_list or not set([x.lower() for x in name.split(",")]).isdisjoint(object_list): # if not object_list or name in object_list: token_counter.update([name]) tokens = set() # pick top N tokens token_counter_return = {} for token, count in token_counter.most_common(): tokens.add(token) token_counter_return[token] = count if len(tokens) == num_tokens: break if verbose: print(("Keeping %d / %d objects" % (len(tokens), len(token_counter)))) return tokens, token_counter_return def merge_duplicate_boxes(data): def IoU(b1, b2): if b1[2] <= b2[0] or b1[3] <= b2[1] or b1[0] >= b2[2] or b1[1] >= b2[3]: return 0 b1b2 = np.vstack([b1, b2]) minc = np.min(b1b2, 0) maxc = np.max(b1b2, 0) union_area = (maxc[2] - minc[0]) * (maxc[3] - minc[1]) int_area = (minc[2] - maxc[0]) * (minc[3] - maxc[1]) return float(int_area) / float(union_area) def to_x1y1x2y2(obj): x1 = obj["x"] y1 = obj["y"] x2 = obj["x"] + obj["w"] y2 = obj["y"] + obj["h"] return np.array([x1, y1, x2, y2], dtype=np.int32) def inside(b1, b2): return b1[0] >= b2[0] and b1[1] >= b2[1] and b1[2] <= b2[2] and b1[3] <= b2[3] def overlap(obj1, obj2): b1 = to_x1y1x2y2(obj1) b2 = to_x1y1x2y2(obj2) iou = IoU(b1, b2) if all(b1 == b2) or iou > 0.9: # consider as the same box return 1 elif (inside(b1, b2) or inside(b2, b1)) and obj1["names"][0] == obj2["names"][ 0 ]: # same object inside the other return 2 elif iou > 0.6 and obj1["names"][0] == obj2["names"][0]: # multiple overlapping same object return 3 else: return 0 # no overlap num_merged = {1: 0, 2: 0, 3: 0} print("merging boxes..") for img in data: # mark objects to be merged and save their ids objs = img["objects"] num_obj = len(objs) for i in range(num_obj): if "M_TYPE" in objs[i]: # has been merged continue merged_objs = [] # circular refs, but fine for j in range(i + 1, num_obj): if "M_TYPE" in objs[j]: # has been merged continue overlap_type = overlap(objs[i], objs[j]) if overlap_type > 0: objs[j]["M_TYPE"] = overlap_type merged_objs.append(objs[j]) objs[i]["mobjs"] = merged_objs # merge boxes filtered_objs = [] merged_num_obj = 0 for obj in objs: if "M_TYPE" not in obj: ids = [obj["object_id"]] dims = [to_x1y1x2y2(obj)] prominent_type = 1 for mo in obj["mobjs"]: ids.append(mo["object_id"]) obj["names"].extend(mo["names"]) dims.append(to_x1y1x2y2(mo)) if mo["M_TYPE"] > prominent_type: prominent_type = mo["M_TYPE"] merged_num_obj += len(ids) obj["ids"] = ids mdims = np.zeros(4) if prominent_type > 1: # use extreme mdims[:2] = np.min(np.vstack(dims)[:, :2], 0) mdims[2:] = np.max(np.vstack(dims)[:, 2:], 0) else: # use mean mdims = np.mean(np.vstack(dims), 0) obj["x"] = int(mdims[0]) obj["y"] = int(mdims[1]) obj["w"] = int(mdims[2] - mdims[0]) obj["h"] = int(mdims[3] - mdims[1]) num_merged[prominent_type] += len(obj["mobjs"]) obj["mobjs"] = None obj["names"] = list(set(obj["names"])) # remove duplicates filtered_objs.append(obj) else: assert "mobjs" not in obj img["objects"] = filtered_objs assert merged_num_obj == num_obj print("# merged boxes per merging type:") print(num_merged) def build_token_dict(vocab): """build bi-directional mapping between index and token""" token_to_idx, idx_to_token = {}, {} next_idx = 1 vocab_sorted = sorted(list(vocab)) # make sure it's the same order everytime for token in vocab_sorted: token_to_idx[token] = next_idx idx_to_token[next_idx] = token next_idx = next_idx + 1 return token_to_idx, idx_to_token def sentence_preprocess(phrase): """preprocess a sentence: lowercase, clean up weird chars, remove punctuation""" replacements = { "½": "half", "—": "-", "™": "", "¢": "cent", "ç": "c", "û": "u", "é": "e", "°": " degree", "è": "e", "…": "", } # phrase = phrase.encode('utf-8') phrase = phrase.lstrip(" ").rstrip(" ") for k, v in replacements.items(): phrase = phrase.replace(k, v) # return str(phrase).lower().translate(None, string.punctuation).decode('utf-8', 'ignore') return str(phrase).lower().translate(str.maketrans("", "", string.punctuation)) def make_alias_dict_new(dict_file, dict_file2): # return {} alias_list = [] for line in open(dict_file, "r"): alias = [alia.strip("\n").strip("\r") for alia in line.strip("\n").strip("\r").split(",")] alias_list.append(alias) for line in open(dict_file2, "r"): alias = [alia.strip("\n").strip("\r") for alia in line.strip("\n").strip("\r").split(",")] alias_list.append(alias) alias_list_merged = [] merged_id = [] for i in range(len(alias_list)): if i in merged_id: continue merged_id.append(i) a = alias_list[i] a_set = set(a) if len(a) <= 1: continue while True: find_overlap = False for j in range(len(alias_list)): if j in merged_id: continue b = alias_list[j] if not a_set.isdisjoint(b): # print(i, j, a, b) a.extend(b) a_set = set(a) merged_id.append(j) find_overlap = True if not find_overlap: break # a = list(set(a)) if len(a) > 1: # alias_list_merged.append(list(set(a))) alias_list_merged.append(a) out_dict = {} for alias in alias_list_merged: # name = alias[0] # for alia in alias[1:]: # match = SequenceMatcher(None, name, alia).find_longest_match() # name = name[match.a:match.a + match.size] name = ",".join(alias) for alia in alias: out_dict[alia] = name print("merged token", out_dict) print("merge token", len(list(out_dict.keys())), "to", len(set(list(out_dict.values())))) return out_dict def make_alias_dict(dict_file): """create an alias dictionary from a file""" out_dict = {} vocab = [] for line in open(dict_file, "r"): alias = line.strip("\n").strip("\r").split(",") alias_target = alias[0] if alias[0] not in out_dict else out_dict[alias[0]] for a in alias: out_dict[a] = alias_target # use the first term as the aliasing target vocab.append(alias_target) print("merge token", len(list(out_dict.keys())), "to", len(set(list(out_dict.values())))) return out_dict, vocab def make_list(list_file): """create a blacklist list from a file""" return [l.strip("\n").strip("\r") for line in open(list_file) for l in line.strip("\n").strip("\r").split(",")] return [line.strip("\n").strip("\r") for line in open(list_file)] def filter_object_boxes(data, image_data, area_frac_thresh): """ filter boxes by a box area-image area ratio threshold """ thresh_count = 0 all_count = 0 for i, img in enumerate(data): filtered_obj = [] area = float(image_data[i]["height"] * image_data[i]["width"]) for obj in img["objects"]: if float(obj["h"] * obj["w"]) > area * area_frac_thresh: filtered_obj.append(obj) thresh_count += 1 all_count += 1 img["objects"] = filtered_obj print("box threshod: keeping %i/%i boxes" % (thresh_count, all_count)) def filter_by_idx(data, valid_list): return [data[i] for i in valid_list] def object(args): print("start") pprint.pprint(args) if args.apply_exif: print("-" * 60) print("We will apply exif orientation...") print("-" * 60) base_dir = args.path object_alias_path = os.path.join(base_dir, "annotations", "object_alias.txt") object_alias_path2 = os.path.join(os.path.dirname(os.path.realpath(__file__)), "VG/1600-400-20/objects_vocab.txt") object_data_path = os.path.join(base_dir, "annotations", "objects.json") image_data_path = os.path.join(base_dir, "annotations", "image_data.json") obj_alias_dict = {} print("using object alias from %s" % (object_alias_path)) print("using object alias from %s" % (object_alias_path2)) # obj_alias_dict, obj_vocab_list = make_alias_dict(object_alias_path) obj_alias_dict = make_alias_dict_new(object_alias_path, object_alias_path2) object_list = [] if len(args.object_list_path) > 0: print("using object list from %s" % (args.object_list_path)) object_list = make_list(args.object_list_path) object_list = [x.lower() for x in object_list] # assert len(object_list) >= args.num_objects print("number of objects", len(object_list)) exclude_object_list = [] if len(args.exclude_object_list_path) > 0: print("using exclude object list from %s" % (args.exclude_object_list_path)) exclude_object_list = make_list(args.exclude_object_list_path) exclude_object_list = [x.lower() for x in exclude_object_list] print("number of exclude objects", len(exclude_object_list)) # read in the annotation data print("loading json files..") print("using object data from %s" % (object_data_path)) object_data = json.load(open(object_data_path)) print("using image data from %s" % (image_data_path)) image_data = json.load(open(image_data_path)) assert len(object_data) == len(image_data) num_im = len(image_data) # sanity check for i in range(num_im): assert object_data[i]["image_id"] == image_data[i]["image_id"] print("processing %i images" % num_im) # preprocess label data preprocess_object_labels(object_data, alias_dict=obj_alias_dict) if args.min_box_area_frac > 0: # filter out invalid small boxes print("threshold bounding box by %f area fraction" % args.min_box_area_frac) filter_object_boxes(object_data, image_data, args.min_box_area_frac) # filter by box dimensions # merge_duplicate_boxes(object_data) # build vocabulary object_tokens, object_token_counter = extract_object_token(object_data, args.num_objects, object_list) label_to_idx, idx_to_label = build_token_dict(object_tokens) print("object list missing:", list(set(set([x for token in object_tokens for x in token.split(",")]) - set(object_list)).intersection(set(object_list)))) print("object list merged:") for tokens in object_tokens: inter = set(tokens.split(",")).intersection(set(object_list)) if len(list(inter)) > 1: print(tokens) exclude_object_tokens = [] if len(args.exclude_object_list_path) > 0: for token in tqdm(object_tokens): if not set([x.lower() for x in token.split(",")]).isdisjoint(exclude_object_list): # if token in exclude_object_list: exclude_object_tokens.append(token) print("exclude_object_tokens", exclude_object_tokens) print("exclude_object_tokens", len(exclude_object_tokens)) # print out vocabulary print("objects: ") print(list(object_token_counter.items())[:100]) print(list(object_token_counter.items())[-100:]) # Convert image mnetadata print("converting image info ...") images = [] image_ids = [] for i in tqdm(range(num_im), mininterval=0.5): path = image_data[i]["url"] path = os.path.normpath(path) path_split = path.split(os.sep) image_id = image_data[i]["image_id"] coco_id = image_data[i]["coco_id"] flickr_id = image_data[i]["flickr_id"] height = image_data[i]["height"] width = image_data[i]["width"] assert image_data[i]["image_id"] == object_data[i]["image_id"] has_obj = False for obj in object_data[i]["objects"]: names = obj["names"] assert len(names) == 1 name = names[0] if name not in object_tokens: continue if name in exclude_object_tokens: continue has_obj = True break if not has_obj: continue img = {} img["id"] = image_id img["file_name"] = os.path.join(path_split[-2], path_split[-1]) img["height"] = height img["width"] = width if args.apply_exif: filename = os.path.join(base_dir, img["file_name"]) image = read_image(filename, format="BGR") if image.shape[1] != img["width"] or image.shape[0] != img["height"]: print("before exif correction: ", img) img["width"], img["height"] = image.shape[1], image.shape[0] print("after exif correction: ", img) images.append(img) image_ids.append(image_id) # build train/val/test splits print("build train/val/test splits") num_im = len(images) num_im_train = max(int(num_im * 0.7), num_im - 5000) print("build train split") images_train = images[:num_im_train] image_ids_train = image_ids[:num_im_train] print("build val split") images_val = images[num_im_train:] image_ids_val = image_ids[num_im_train:] # Convert instance annotations print("converting annotations info ...") annotations = [] annotations_train = [] annotations_val = [] label_to_synset = {obj: [] for obj in object_tokens} image_count = {obj: 0 for obj in object_tokens} instance_count = {obj: 0 for obj in object_tokens} ann_id = 1 for i in tqdm(range(num_im), mininterval=0.5): image_id = image_data[i]["image_id"] if image_id not in image_ids: continue assert image_data[i]["image_id"] == object_data[i]["image_id"] names = [] for obj in object_data[i]["objects"]: name = obj["names"] assert len(name) == 1 name = name[0] if name not in object_tokens: continue if name in exclude_object_tokens: continue names.append(name) synsets = obj["synsets"] object_id = obj["object_id"] # merged_object_ids = obj['merged_object_ids'] x = obj["x"] y = obj["y"] h = obj["h"] w = obj["w"] ann = {} ann["id"] = ann_id ann_id += 1 ann["image_id"] = image_id ann["category_id"] = label_to_idx[name] ann["phrase"] = name.split(",")[0].strip("\n").strip("\r").strip() ann["isobject"] = 1 # ann["bbox"] = [x, y, x + w, y + h] ann["bbox"] = [x, y, w, h] ann["area"] = h * w ann["iscrowd"] = False annotations.append(ann) if image_id in image_ids_train: annotations_train.append(ann) elif image_id in image_ids_val: annotations_val.append(ann) else: assert 0 # assert len(synsets) <= 1 # if len(synsets) > 0: # synset = synsets[0] # if synset not in label_to_synset[name]: # label_to_synset[name].append(synset) for synset in synsets: if synset not in label_to_synset[name]: label_to_synset[name].append(synset) instance_count[name] += 1 for name in list(set(names)): image_count[name] += 1 print("all images: ", len(images)) print("train images: ", len(images_train)) print("val images: ", len(images_val)) print("all annotations: ", len(annotations)) print("train annotations: ", len(annotations_train)) print("val annotations: ", len(annotations_val)) oi_train = {} oi_val = {} oi_all = {} # Add basic dataset info print("adding basic dataset info") oi_train["info"] = {} oi_val["info"] = {} oi_all["info"] = {} # Add license information print("adding basic license info") oi_train["licenses"] = [] oi_val["licenses"] = [] oi_all["licenses"] = [] # Convert category information print("converting category info") categories = [] for i, name in idx_to_label.items(): cat = {} cat["id"] = i cat["name"] = name # cat["synsets"] = label_to_synset[name] alias = name.split(",") name = alias[0].strip("\n").strip("\r").strip() # for alia in alias[1:]: # match = SequenceMatcher(None, name, alia).find_longest_match() # name = name[match.a:match.a + match.size] cat["name"] = name categories.append(cat) oi_train["categories"] = categories oi_val["categories"] = categories oi_all["categories"] = categories # Convert image mnetadata print("converting image info ...") oi_train["images"] = images_train oi_val["images"] = images_val oi_all["images"] = images # Convert instance annotations print("converting annotations ...") oi_train["annotations"] = annotations_train oi_val["annotations"] = annotations_val oi_all["annotations"] = annotations # Write annotations into .json file filename = os.path.join(base_dir, "annotations/", f"visualgenome_{len(categories)}_box_train.json") if exclude_object_list: filename = os.path.join(base_dir, "annotations/", f"visualgenome_{len(categories)}minus{len(exclude_object_list)}_box_train.json") print("writing output to {}".format(filename)) json.dump(oi_train, open(filename, "w")) filename = os.path.join(base_dir, "annotations/", f"visualgenome_{len(categories)}_box_val.json") if exclude_object_list: filename = os.path.join(base_dir, "annotations/", f"visualgenome_{len(categories)}minus{len(exclude_object_list)}_box_val.json") print("writing output to {}".format(filename)) json.dump(oi_val, open(filename, "w")) filename = os.path.join(base_dir, "annotations/", f"visualgenome_{len(categories)}_box.json") if exclude_object_list: filename = os.path.join(base_dir, "annotations/", f"visualgenome_{len(categories)}minus{len(exclude_object_list)}_box.json") print("writing output to {}".format(filename)) json.dump(oi_all, open(filename, "w")) filename = os.path.join(base_dir, "annotations/", f"visualgenome_{len(categories)}_box_categories.json") if exclude_object_list: filename = os.path.join(base_dir, "annotations/", f"visualgenome_{len(categories)}minus{len(exclude_object_list)}_box_categories.json") print("writing output to {}".format(filename)) json.dump(categories, open(filename, "w")) print("Done") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("-p", "--path", dest="path", help="path to visual genome data", type=str) parser.add_argument( "--apply-exif", dest="apply_exif", action="store_true", help="apply the exif orientation correctly", ) parser.add_argument("--object_list_path", default="VG/object_list.txt", type=str) parser.add_argument("--exclude_object_list_path", default="", type=str) parser.add_argument( "--num_objects", default=150, type=int, help="set to 0 to disable filtering" ) parser.add_argument("--min_box_area_frac", default=0.002, type=float) args = parser.parse_args() object(args)