import argparse
import json
import os
from pycocotools.coco import COCO


def new_path(old_path):
    # 使用os.path.dirname获取目录部分
    directory = os.path.dirname(old_path)
    # 使用os.path.basename获取文件名部分（包括扩展名）
    filename_with_extension = os.path.basename(old_path)
    # 使用os.path.splitext分割文件名和扩展名
    filename, extension = os.path.splitext(filename_with_extension)
    # 修改文件名部分，添加_filter并重新组合扩展名
    new_filename = filename + '_filter' + extension
    # 如果directory为空（即文件名只包含一个文件名没有目录），则使用'.'作为目录
    if not directory:
        directory = '.'
    # 使用os.path.join组合目录和新的文件名
    return os.path.join(directory, new_filename)


def process_coco(file, annot_name, output=""):
    # 初始化COCO对象
    coco = COCO(file)
    # 获取所有的类别信息
    categories = coco.loadCats(coco.getCatIds())
    # 找到"line"类别的ID
    cat_id = []
    for cat in categories:
        if cat['name'] in annot_name:
            cat_id.append(cat['id'])

    if cat_id is None:
        print("Category not found in annotations.")
    else:
        # 获取所有的标注
        annotations = coco.dataset['annotations']

        # 创建一个新的标注列表，不包含"line"类别的标注
        new_annotations = [ann for ann in annotations if ann['category_id'] not in cat_id]

        # 将新的标注列表写回到一个新的JSON文件中
        if output == "":
            output = new_path(file)

        with open(output, 'w') as f:
            coco.dataset['annotations'] = new_annotations
            json.dump(coco.dataset, f, ensure_ascii=True, indent=2)


def main():
    parser = argparse.ArgumentParser("删除特定标注信息\n")
    parser.add_argument("input", help="输入路径")
    parser.add_argument("type", help="数据标注格式类型，目前支持coco")
    parser.add_argument("--labels", nargs='+', help="需要清除的标签")
    parser.add_argument("--output", help="[可选]输出路径")

    args = parser.parse_args()
    if args.type == 'coco':
        process_coco(args.input, args.labels)


if __name__ == "__main__":
    main()