| import json |
| import random |
| from tqdm import tqdm |
| import argparse |
| import numpy as np |
|
|
| MAX_INTERACT_PER_IMAGE = 60 |
| WIDTH = 960 |
| HEIGHT = 540 |
|
|
|
|
| def in_pred_bbox(img_preds, X, Y): |
| for pred in img_preds: |
| if X - pred['bbox'][0] >= 0 and \ |
| Y - pred['bbox'][1] >= 0 and \ |
| X - pred['bbox'][0] <= pred['bbox'][2] and \ |
| Y - pred['bbox'][1] <= pred['bbox'][3]: |
| return True |
| return False |
|
|
|
|
| def gen_interact_random(annotation_data, width, height): |
| random_inter = {} |
| |
| |
| for k in range(1, MAX_INTERACT_PER_IMAGE + 1): |
| cur_group = [] |
| for img in annotation_data['images']: |
| img_id = img['id'] |
| cur_group.append({ |
| 'X': random.randint(1, width), |
| 'Y': random.randint(1, height), |
| 'img_id': img_id |
| }) |
| random_inter[k] = cur_group |
| return random_inter |
|
|
|
|
| def gen_interact_guided(anno_data, guidance, width, height): |
| interact = {} |
| img_pred = {} |
| for pred in guidance: |
| img_id = pred['image_id'] |
| if img_id not in img_pred: |
| img_pred[img_id] = [] |
| img_pred[img_id].append(pred) |
| img_pred_mask = {} |
| for img_id in img_pred.keys(): |
| img_pred_mask[img_id] = np.zeros((width, height)) |
| if not img_pred[img_id]: |
| raise ValueError(f'No prediction for image id {img_id} in guidance') |
| for pred in img_pred[img_id]: |
| x, y, w, h = pred['bbox'] |
| x = np.floor(x).astype(int) |
| y = np.floor(y).astype(int) |
| w = np.ceil(w).astype(int) |
| h = np.ceil(h).astype(int) |
| img_pred_mask[img_id][x:x+w, y:y+h] = 1 |
| |
| |
| for k in range(1, MAX_INTERACT_PER_IMAGE + 1): |
| outside_prob = k / MAX_INTERACT_PER_IMAGE |
| cur_group = [] |
| for img in anno_data['images']: |
| img_id = img['id'] |
| outside = random.random() < outside_prob |
| |
| |
| |
| |
| |
| |
| |
| |
| for _ in range(10000): |
| X = random.randint(1, width) |
| Y = random.randint(1, height) |
| if outside or img_id not in img_pred.keys() or in_pred_bbox(img_pred[img_id], X, Y): |
| break |
| cur_group.append({ |
| 'X': X, |
| 'Y': Y, |
| 'img_id': img_id |
| }) |
| interact[k] = cur_group |
| return interact |
|
|
|
|
| def main(args): |
| with open(args.annotation, 'r') as f: |
| anno_data = json.load(f) |
| if args.guidance: |
| with open(args.guidance, 'r') as f: |
| guidance = json.load(f) |
| guidance_img = [pred['image_id'] for pred in guidance] |
| anno_data['images'] = [img for img in anno_data['images'] if img['id'] in guidance_img] |
| anno_data['annotations'] = [anno for anno in anno_data['annotations'] if anno['image_id'] in guidance_img] |
| interact_data = gen_interact_guided(anno_data, guidance, WIDTH, HEIGHT) |
| else: |
| interact_data = gen_interact_random(anno_data, WIDTH, HEIGHT) |
| with open(args.output, 'w') as f: |
| json.dump(interact_data, f) |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser(description='Generate interact points') |
| parser.add_argument('-a', '--annotation', type=str, required=True, help='Annotation file path') |
| parser.add_argument('-g', '--guidance', type=str, help='guidance file path') |
| parser.add_argument('-o', '--output', type=str, required=True, help='Output file path') |
| args = parser.parse_args() |
| main(args) |
|
|
