Spaces:
Build error
Build error
File size: 1,850 Bytes
1865436 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
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()
|