'''

'''

import os
import cv2
import argparse
import shutil
from pathlib import Path
from tqdm import tqdm

from animeinsseg import AnimeInsSeg

# 设置模型路径
ckpt = r'models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt'
mask_thres = 0.3
instance_thres = 0.3
refine_kwargs = {'refine_method': 'refinenet_isnet'}  # 如果不使用 refinenet，设置为 None
# refine_kwargs = None

# 初始化模型
net = AnimeInsSeg(ckpt, mask_thr=mask_thres, refine_kwargs=refine_kwargs)

def has_instances(image_path):
    # 读取图像
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    # 推理
    instances = net.infer(
        img,
        output_type='numpy',
        pred_score_thr=instance_thres
    )

    # 如果没有检测到对象，返回 False
    if instances.bboxes is None:
        return False
    return True

def copy_images_without_instances(input_path, output_dir):
    input_path = Path(input_path)
    output_dir = Path(output_dir)

    if not output_dir.exists():
        output_dir.mkdir(parents=True)

    image_paths = list(input_path.rglob("*.png")) + list(input_path.rglob("*.jpg"))
    for image_path in tqdm(image_paths, desc="Processing images"):
        if not has_instances(image_path):
            # 拷贝不包含实例的图片到目标文件夹
            shutil.copy(image_path, output_dir / image_path.name)

def main():
    parser = argparse.ArgumentParser(description="Copy images without instances to a target directory")
    parser.add_argument("input_path", type=str, help="Path to the input image or folder")
    parser.add_argument("output_dir", type=str, help="Path to the output directory")
    args = parser.parse_args()

    copy_images_without_instances(args.input_path, args.output_dir)

if __name__ == "__main__":
    main()