src/ss/ss.py

from src.ss.signboard_detect import inference_signboard
import os
import numpy as np
import cv2


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

            np.savetxt(output_path, segment_array)

    return output_path, segment_array