File size: 3,208 Bytes
6d35aff | 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 | # Description: Filter out the bounding boxes based on the following criteria:
# 1. Huge bounding boxes: Filter out the bounding boxes whose area is larger than a certain threshold.
# 2. Non-maximum suppression (NMS): Filter out the bounding boxes that have a high IoU with other bounding boxes.
# 3. Category: Filter out the bounding boxes whose category is in the FILTER_OUT_CATEGORY list.
#
# Usage: python filter_box.py -i <input_file> -o <output_file> [-n] [-hu] [-c]
# -i, --input: Path to the input JSON file.
# -o, --output: Path to the output JSON file.
# -n, --nms: Perform non-maximum suppression.
# -hu, --huge: Filter out huge bounding boxes.
# -c, --category: Filter out bounding boxes based on category.
import json
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', type=str, required=True)
parser.add_argument('-o', '--output', type=str, required=True)
parser.add_argument('-n', '--nms', action='store_true')
parser.add_argument('-hu', '--huge', action='store_true')
parser.add_argument('-c', '--category', action='store_true')
IMG_SIZE = 960*540
HUGE_THRESHOLD = 0.9
NMS_THRESHOLD = 0.7
FILTER_OUT_CATEGORY = ['controller', 'controllers', 'hand', 'hands']
def IoU(box1, box2):
x1, y1, w1, h1 = box1
x2, y2, w2, h2 = box2
inter_x1 = max(x1, x2)
inter_y1 = max(y1, y2)
inter_x2 = min(x1 + w1, x2 + w2)
inter_y2 = min(y1 + h1, y2 + h2)
inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
union_area = w1 * h1 + w2 * h2 - inter_area
iou = inter_area / union_area if union_area > 0 else 0
return iou
def filter_huge(preds):
filtered_preds = []
for pred in preds:
x, y, w, h = pred['bbox']
if w * h / IMG_SIZE < HUGE_THRESHOLD:
filtered_preds.append(pred)
return filtered_preds
def nms(preds):
sorted_preds = sorted(preds, key=lambda x: x['score'], reverse=True)
nms_preds = []
for pred in sorted_preds:
if len(nms_preds) == 0:
nms_preds.append(pred)
continue
for nms_pred in nms_preds:
if IoU(pred['bbox'], nms_pred['bbox']) > NMS_THRESHOLD:
break
else:
nms_preds.append(pred)
return nms_preds
def filter_category(preds):
filtered_preds = []
for pred in preds:
if pred['category_id'].lower() not in FILTER_OUT_CATEGORY:
filtered_preds.append(pred)
return filtered_preds
def main(args):
with open(args.input, 'r') as f:
pred_data = json.load(f)
img_pred = {}
for pred in pred_data:
img_id = pred['image_id']
if img_id not in img_pred:
img_pred[img_id] = []
img_pred[img_id].append(pred)
pred_data = []
for img_id in img_pred:
preds = img_pred[img_id]
if args.huge:
preds = filter_huge(preds)
if args.nms:
preds = nms(preds)
if args.category:
preds = filter_category(preds)
pred_data.extend(preds)
with open(args.output, 'w') as f:
json.dump(pred_data, f)
if __name__ == '__main__':
args = parser.parse_args()
main(args)
|