prepare_csv.py

import os
import csv
import argparse
import pandas as pd
import numpy as np
from tqdm import tqdm


parser = argparse.ArgumentParser()
parser.add_argument('--label_path', type=str,
                    help='Path of the label folder generated by detect.py')
parser.add_argument('--test_path', type=str,
                    help='Path of the test images folder')

opt = parser.parse_args()


label_dir = sorted(os.listdir(opt.label_path))
test_dir = sorted(os.listdir(opt.test_path))


label = []
box = []


for i in tqdm(test_dir):
    file = i[:-4] + '.txt'
    # print(file)
    if file in label_dir:
        label.append([i, 1])
    else:
        label.append([i, 0])

for idx, i in enumerate(label_dir):
    file = i[:-4] + '.png'

    with open(os.path.join(opt.label_path, i), "r") as f:
        for i in f.readlines():
            box.append([file] + np.asarray(i.strip().split(' ')
                       [1:]).astype(np.float).tolist())


with open("Final_DefectBoxes_No_Free_Lunch.csv", "w") as csvfile:
    csvwriter = csv.writer(csvfile)

    csvwriter.writerow(
        [
            "image_id", "X", "Y", "W", "H"
        ]
    )

    csvwriter.writerows(box)

with open("Final_DefectTypes_No_Free_Lunch.csv", "w") as csvfile:
    csvwriter = csv.writer(csvfile)

    csvwriter.writerow(
        [
            "image_id", "defect_flag"
        ]
    )

    csvwriter.writerows(label)