from seg.hairseg_single_model import Evaluator
import os
import cv2
import numpy as np
if __name__ == "__main__":

    data_path = "/home/liyang/project/matting/合格/origin"
    dst_path = "/home/liyang/project/matting/合格/origin_seg_res1102"
    if not os.path.exists(dst_path):
        os.mkdir(dst_path)
    seg_model = Evaluator(gpu_id=0, output_img_size=512, nclass=3)
    for imgs in os.listdir(data_path):
        if imgs.endswith(".txt"):
            continue

        img_path = os.path.join(data_path, imgs)
        img = cv2.imread(img_path)
        if img_path.endswith('.png'):
            kpts_1k = np.loadtxt(img_path.replace('.png', '_landmark1k.txt'))
        elif img_path.endswith('.jpg'):
            kpts_1k = np.loadtxt(img_path.replace('.jpg', '_landmark1k.txt'))
        output_img_size = 512
        mask = seg_model.eval(img, kpts_1k)
        cv2.imwrite(os.path.join(dst_path, imgs), mask)