WGU-Capstone/validation/compare_to_gt.py

import argparse
import os

def dir_path(string):
    if os.path.exists(string):
        return string
    else:
        raise NotADirectoryError(string)

def init_argparse() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="FaceDetection",
        usage="%(prog)s [OPTION]",
        description="Run face localization"
    )
    parser.add_argument(
        "ground_truth", type=dir_path, help="ground truth file"
    )
    parser.add_argument(
        "test_file", type=dir_path, help="file to compare to the ground truth"
    )
    parser.add_argument(
        "out_file", help="the file to write the output to"
    )
    parser.add_argument(
        "-f", "--faces_count_file", help="the file output the number of faces found in each frame"
    )
    return parser

multiplication_factor = 0.05
def get_adjustment_amount(imgSize, x1, y1, x2, y2):

    # find the difference between the left gap and the right gap, divide it by two, and multiply it by the speed scale
    horizontal_adjustment = multiplication_factor * (x1 - (imgSize[0] - x2)) / 2
    vertical_adjustment = multiplication_factor * (y1 - (imgSize[0] - y2)) / 2

    return [horizontal_adjustment, vertical_adjustment]

parser = init_argparse()
args = parser.parse_args()

class FrameBox:
    frame_number: int 
    top_left: tuple 
    bottom_right: tuple

    def __init__(self, frame_number, x1, y1, x2, y2):
        self.frame_number = frame_number
        self.top_left = (x1, y1)
        self.bottom_right = (x2, y2)

ground_truth = []
with open(args.ground_truth, 'r') as gt_file:
    lines = gt_file.readlines()
    for line in lines:
        items = line.split(",")
        ground_truth.append(FrameBox(int(items[0].strip()), int(items[1].strip()), int(items[2].strip()), int(items[3].strip()), int(items[4].strip())))

test_results = []
with open(args.test_file, 'r') as test_file:
    lines = test_file.readlines()
    for line in lines:
        items = line.split(",")
        test_results.append(FrameBox(int(items[0].strip()), int(items[1].strip()), int(items[2].strip()), int(items[3].strip()), int(items[4].strip())))

test_fb: FrameBox
gt_frame: FrameBox
last_frame_num = 0
average_sum = 0
average_count = 0
print("length of test file: " + str(len(test_results)))
for test_fb in test_results:
    # make
    def bring_up_gt():
        if test_fb.frame_number > ground_truth[0].frame_number:
            ground_truth.pop(0)

            if test_fb.frame_number > ground_truth[0].frame_number:
                # we need to include the empty frames too
                if args.faces_count_file:
                    with open(args.faces_count_file, 'a') as out_file:
                        out_file.write(f"{ground_truth[0].frame_number}, 0,\n")

            bring_up_gt()

    bring_up_gt()

    assert(ground_truth[0].frame_number == test_fb.frame_number)
    gt_frame = ground_truth[0]

    gt_adjustment = get_adjustment_amount((1920, 1000), gt_frame.top_left[0], gt_frame.top_left[1], gt_frame.bottom_right[0], gt_frame.bottom_right[1])
    test_adjustment = get_adjustment_amount((1920, 1000), test_fb.top_left[0], test_fb.top_left[1], test_fb.bottom_right[0], test_fb.bottom_right[1])

    if last_frame_num != test_fb.frame_number :
        if average_count > 0:
            with open(args.out_file, 'a') as out_file:
                out_file.write(f"{average_sum},\n")
        if args.faces_count_file:
            with open(args.faces_count_file, 'a') as out_file:
                out_file.write(f"{test_fb.frame_number}, {average_count},\n")
        average_sum = 0
        average_count = 0

    average_count += 1
    # get the average 
    average_sum += abs(( (gt_adjustment[0] - test_adjustment[0]) + (gt_adjustment[1] - test_adjustment[1]) ) / 2)
    
    last_frame_num = test_fb.frame_number