utils/image.py

import os
import cv2
from src.ss.ss import handle_ss
from src.sts.demo.sts import handle_sts
from src.ir.ir import handle_ir


def create_folder(dir_path):
    if not os.path.isdir(dir_path):
        os.mkdir(dir_path)
    return


def parse(image, filename):
    dir_path = os.path.basename(filename).split('.')[0]
    dir_path = os.path.join("./outputs", dir_path)
    create_folder(dir_path)
    filepath = os.path.join(dir_path, filename)
    image.save(filepath)
    segment_path, segment_array = handle_ss(filepath, dir_path)
    output_path_box, output_path_text, output_path_visual, dict_box_sign_out, dict_rec_sign_out = handle_sts(
        filepath, segment_path, dir_path)
    predicted = handle_ir(filepath, dict_rec_sign_out, dir_path)
    print(predicted)
    return