File size: 4,299 Bytes
032e687 |
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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
import argparse
import json
import os
from pathlib import Path
from tqdm import tqdm
from detectron2.data.detection_utils import read_image
def parse_args():
"""
Parse input arguments
"""
parser = argparse.ArgumentParser(
description="Convert Objects365 annotations into MS Coco format"
)
parser.add_argument("--root_dir", dest="root_dir", help="path to objects365 data", type=str)
parser.add_argument(
"--apply_exif",
dest="apply_exif",
action="store_true",
help="apply the exif orientation correctly",
)
parser.add_argument(
"--subsets",
type=str,
nargs="+",
default=["val", "train"],
choices=["train", "val", "test", "minival"],
help="subsets to convert",
)
parser.add_argument("--image_info_path", type=str, help="image_info_path")
args = parser.parse_args()
return args
args = parse_args()
root_dir = args.root_dir
if args.apply_exif:
print("-" * 60)
print("We will apply exif orientation...")
print("-" * 60)
if not isinstance(args.subsets, list):
args.subsets = [args.subsets]
for subset in args.subsets:
# Convert annotations
print("converting {} data".format(subset))
# Select correct source files for each subset
if subset == "train":
json_name = "zhiyuan_objv2_train.json"
elif subset == "val":
json_name = "zhiyuan_objv2_val.json"
elif subset == "minival":
json_name = "zhiyuan_objv2_val.json"
# Load original annotations
print("loading original annotations ...")
json_path = os.path.join(root_dir, "annotations", json_name)
json_data = json.load(open(json_path, "r"))
print("loading original annotations ... Done")
print(json_data.keys())
oi = {}
# Add basic dataset info
print("adding basic dataset info")
# Add license information
print("adding basic license info")
oi["licenses"] = json_data["licenses"]
# Convert category information
print("converting category info")
oi["categories"] = json_data["categories"]
# Convert image mnetadata
print("converting image info ...")
images = json_data["images"]
if subset == "minival":
images = images[:5000]
print(f"{len(images)} images get")
rm_image_ids = []
if args.apply_exif:
image_info = {}
with open(args.image_info_path, "r") as f:
for line in f.readlines():
line = line.strip().split()
image_id, file_name, height, width, channel = line
image_id = int(image_id)
height = int(height)
width = int(width)
image_info[image_id] = [file_name, height, width]
print(f"{len(image_info)} image_info get")
new_images = []
for img in tqdm(images):
image_id = img["id"]
if image_id not in image_info.keys():
rm_image_ids.append(image_id)
print("removing", img)
continue
file_name, height, width = image_info[image_id]
assert file_name == img["file_name"]
if width != img["width"] or height != img["height"]:
print("before exif correction: ", img)
img["width"], img["height"] = width, height
print("after exif correction: ", img)
new_images.append(img)
images = new_images
oi["images"] = images
print(f"{len(images)} images keep")
# Convert instance annotations
print("converting annotations ...")
annotations = json_data["annotations"]
print(f"{len(annotations)} annotations get")
annotations = [ann for ann in annotations if ann["image_id"] not in rm_image_ids]
if subset == "minival":
keep_image_ids = [img["id"] for img in images]
annotations = [ann for ann in annotations if ann["image_id"] in keep_image_ids]
oi["annotations"] = annotations
print(f"{len(annotations)} annotations keep")
# Write annotations into .json file
json_path = os.path.join(root_dir, "annotations/", "objects365_{}.json".format(subset))
print("writing output to {}".format(json_path))
json.dump(oi, open(json_path, "w"))
print("Done")
|