package com.chuangzhou.vivid2D.test;

import com.chuangzhou.vivid2D.ai.face_parsing.VividModelWrapper;

import java.io.PrintStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Paths;
import java.util.Set;

/**
 * 测试人脸解析模型
 */
public class AITest {
    public static void main(String[] args) throws Exception {
        System.setOut(new PrintStream(System.out, true, StandardCharsets.UTF_8));
        System.setErr(new PrintStream(System.err, true, StandardCharsets.UTF_8));
        VividModelWrapper wrapper = VividModelWrapper.load(Paths.get("C:\\models\\bisenet_face_parsing.pt"));

        // 使用 BiSeNet 人脸解析模型的 18 个非背景标签
        Set<String> faceLabels = Set.of(
                "skin", "nose", "eye_left", "eye_right", "eyebrow_left",
                "eyebrow_right", "ear_left", "ear_right", "mouth", "lip_upper",
                "lip_lower", "hair", "hat", "earring", "necklace", "clothes",
                "facial_hair", "neck"
        );

        wrapper.segmentAndSave(
                Paths.get("C:\\Users\\Administrator\\Desktop\\b_f4881214f0d18b6cf848b6736f554821.png").toFile(),
                faceLabels,
                Paths.get("C:\\models\\out")
        );
    }
}