Spaces:
Running
Running
File size: 1,816 Bytes
1fec9bd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
'''
'''
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()
|