| import argparse | |
| import json | |
| import os | |
| import tqdm | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description="Convert image list to coco format") | |
| parser.add_argument("--image_root", help="img root", required=True) | |
| parser.add_argument("--json_path", help="output path", required=True) | |
| args = parser.parse_args() | |
| return args | |
| def main(): | |
| args = parse_args() | |
| print(args) | |
| imgs = [] | |
| for root, dirs, files in os.walk(args.image_root): | |
| print(root) | |
| files = [f for f in files if f.endswith((".jpg", ".JPG", ".png", ".PNG", ".jpeg", ".JPEG"))] | |
| for file_name in tqdm.tqdm(files): | |
| path = os.path.join(root, file_name) | |
| rpath = path.replace(args.image_root, "") | |
| json_path = path[:-len(path.split(".")[-1])] + "json" | |
| if os.path.isfile(json_path): | |
| pass | |
| else: | |
| continue | |
| try: | |
| json_data = json.load(open(json_path, "r")) | |
| except Exception as e: | |
| print(json_path) | |
| print(e) | |
| continue | |
| assert json_data["image"]["file_name"] == file_name, json_data | |
| image_id = json_data["image"]["image_id"] | |
| height = json_data["image"]["height"] | |
| width = json_data["image"]["width"] | |
| img = {"file_name": rpath, "height": height, "width": width, "id": image_id,} | |
| imgs.append(img) | |
| if len(imgs) % 1000000 == 0: | |
| save(imgs, args, suffix="_" + str(len(imgs))) | |
| print("#imgs", len(imgs)) | |
| save(imgs, args) | |
| def save(imgs, args, suffix=""): | |
| json_data = dict() | |
| json_data["categories"] = [{"id": 1, "name": "object"}] | |
| json_data["images"] = imgs | |
| json_data["annotations"] = [] | |
| print("json_data", json_data.keys()) | |
| print("categories", len(json_data["categories"])) | |
| print("#images", len(json_data["images"])) | |
| with open(args.json_path + suffix + ".json", "w") as f: | |
| json.dump(json_data, f) | |
| print(args) | |
| if __name__ == "__main__": | |
| main() | |