Filter annotation[.txt] file created using cvat_xml_to_txt.py

filter_annotation.py

import argparse
from collections import namedtuple

ROI = namedtuple("ROI", ["xa", "ya", "xb", "yb"])

arg_parser = argparse.ArgumentParser()
arg_parser.add_argument("--input", type=str, help="Input annotation[.txt] file path", dest="input_file",
                        required=True)
arg_parser.add_argument("--output", type=str, help="Output/Filtered annotation[.txt] file path",
                        dest="output_file", required=True)
arg_parser.add_argument("--allowed_classes", nargs="+", dest="allowed_classes",
                        help="List of allowed classes", required=True)
arg_parser.add_argument("--selectroi", action="store_true", help="Select Region of Interest", dest="select_roi")
args = arg_parser.parse_args()

count = 0

print("Allowed Classes", args.allowed_classes)

roi = None

with open(args.input_file, "r") as f:
    with open(args.output_file, "w") as out_f:
        lines = f.readlines()
        
        if args.select_roi:
            import cv2
            
            filepath, xa, ya, xb, yb, label = lines[0].split(",")
            im = cv2.imread(filepath)
            
            roi = cv2.selectROI(im)
            
            roi = ROI(xa=roi[0],
                      ya=roi[1],
                      xb=roi[0] + roi[2],
                      yb=roi[1] + roi[3])

        for line in lines:
            filepath, xa, ya, xb, yb, label = line.split(",")            
            
            if roi is not None:
                xa, ya, xb, yb = int(xa), int(ya), int(xb), int(yb)

                if  ya < roi.ya or yb > roi.yb or xa < roi.xa or xb > roi.xb:
                    continue
            
            if label.strip() not in args.allowed_classes:
                continue
            
            count += 1
            out_f.write(f"{filepath},{xa},{ya},{xb},{yb},{label}")

print(f"Count # {count}")