File size: 4,252 Bytes
1da285f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | 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 = {}
# Batch-by-batch generation, 1 point per image in each batch
# for k in tqdm(range(1, MAX_INTERACT_PER_IMAGE + 1), desc='Randomly generating interact points'):
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
# Batch-by-batch generation, 1 point per image in each batch
# for k in tqdm(range(1, MAX_INTERACT_PER_IMAGE + 1), desc='Generating guided interact points'):
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
# if outside or img_id not in img_pred.keys() or np.sum(img_pred_mask[img_id]) == 0:
# X = random.randint(1, width)
# Y = random.randint(1, height)
# else:
# Xs, Ys = np.where(img_pred_mask[img_id] == 1)
# idx = random.randint(0, len(Xs) - 1)
# X = int(Xs[idx])
# Y = int(Ys[idx])
for _ in range(10000): # To avoid infinite loop
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)
|