src/ss/main.py

from signboard_detect import inference_signboard
import os
import numpy as np
import argparse
import tqdm
import cv2


def get_parser():
    parser = argparse.ArgumentParser(description="Signboard Detection")

    parser.add_argument("--input",
                        type=str,
                        default="./images",
                        help="A list of space separated input images")
    parser.add_argument("--output",
                        type=str,
                        default="./output/output_signboard",
                        help="A list of array of segmentation")
    parser.add_argument("--checkpoint",
                        type=str,
                        default="./checkpoints/ss/ss.ckpt",
                        help="File path to best model checkpoint")

    args = parser.parse_args()
    return args


def handle(args):
    if args.input:
        if os.path.isdir(args.input):
            args.input = [os.path.join(args.input, fname)
                          for fname in os.listdir(args.input)]

    for path in tqdm.tqdm(args.input):
        print(path)
        img = cv2.imread(path)
        dimensions = img.shape
        hei, wid = dimensions[0], dimensions[1]
        print(hei, wid)
        segment_array = inference_signboard(path, args.checkpoint).astype(int)
        h, w = segment_array.shape
        print(h, w)
        if hei == h and wid == w:
            segment_array = segment_array
        else:
            segment_array = cv2.rotate(
                segment_array, cv2.ROTATE_90_CLOCKWISE)
        txt_name = str(path.split("/")[-1].split(".")[0]) + '.txt'
        if args.output:
            output_path = os.path.join(args.output, txt_name)

            np.savetxt(output_path, segment_array)


def main():
    args = get_parser()
    handle(args)


if __name__ == "__main__":
    main()