import img_label
from img_nsfw import init_nsfw_pipe, check_nsfw

if __name__ == "__main__":
    img = 'test.jpeg'
    pipe = init_nsfw_pipe()
    nsfw_tags = check_nsfw(img, pipe)

    score_general_threshold = 0.35
    score_character_threshold = 0.85

    img_tags = img_label.label_img(
        image=img,
        model="SwinV2",
        l_score_general_threshold=score_general_threshold,
        l_score_character_threshold=score_character_threshold,
    )['general_res']