""" 图片去背景工具 使用rembg库自动去除图片背景 """ import os import sys import argparse from pathlib import Path import numpy as np import cv2 from PIL import Image, ImageOps # 避免 numba 在某些环境下缓存失败 os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache") from rembg import remove, new_session AOT_MODEL_CACHE: dict[tuple[str, str, int, tuple[int, ...]], tuple[object, object]] = {} def _resolve_aot_paths(aot_root, aot_pretrain): """解析 AOT-GAN 的路径配置。""" aot_root = Path(aot_root).resolve() if aot_pretrain is None: aot_pretrain = None else: aot_pretrain = Path(aot_pretrain) if not aot_pretrain.is_absolute(): # 预训练权重相对路径以当前工作目录为基准 aot_pretrain = (Path.cwd() / aot_pretrain).resolve() return aot_root, aot_pretrain def _parse_aot_rates(rates_str): parts = [p for p in rates_str.split("+") if p] return [int(p) for p in parts] def _get_aot_model(aot_root, aot_pretrain, device="cpu", block_num=8, rates=None): """加载/缓存 AOT-GAN 模型。""" aot_root, aot_pretrain = _resolve_aot_paths(aot_root, aot_pretrain) if not aot_root.exists(): raise FileNotFoundError(f"AOT-GAN目录不存在: {aot_root}") if aot_pretrain is None: raise ValueError("AOT-GAN需要指定预训练权重路径（--aot-pretrain）") if not aot_pretrain.exists(): raise FileNotFoundError(f"AOT-GAN权重不存在: {aot_pretrain}") if rates is None: rates = [1, 2, 4, 8] rates_tuple = tuple(rates) key = (str(aot_pretrain), device, int(block_num), rates_tuple) if key in AOT_MODEL_CACHE: return AOT_MODEL_CACHE[key] src_root = aot_root / "src" if str(src_root) not in sys.path: sys.path.insert(0, str(src_root)) import importlib try: import torch except ImportError as exc: raise ImportError("未找到 PyTorch，请先按README安装依赖。") from exc net = importlib.import_module("model.aotgan") class _Args: pass args = _Args() args.block_num = int(block_num) args.rates = rates model = net.InpaintGenerator(args) state = torch.load(str(aot_pretrain), map_location=device) if isinstance(state, dict): if "state_dict" in state: state = state["state_dict"] elif "model" in state: state = state["model"] elif "generator" in state: state = state["generator"] if isinstance(state, dict) and any(k.startswith("module.") for k in state.keys()): state = {k.replace("module.", "", 1): v for k, v in state.items()} model.load_state_dict(state, strict=True) model.to(device) model.eval() AOT_MODEL_CACHE[key] = (model, device) return AOT_MODEL_CACHE[key] def _mask_bbox(mask): ys, xs = np.where(mask > 0) if len(xs) == 0: return None return int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max()) def _expand_bbox(bbox, pad, width, height): if bbox is None: return None x0, y0, x1, y1 = bbox pad = max(0, int(pad)) x0 = max(0, x0 - pad) y0 = max(0, y0 - pad) x1 = min(width - 1, x1 + pad) y1 = min(height - 1, y1 + pad) return x0, y0, x1, y1 def _expand_bbox_min_size(bbox, pad, width, height, min_size): """在给定 pad 基础上，确保裁剪区域至少为 min_size。""" expanded = _expand_bbox(bbox, pad, width, height) if expanded is None: return None x0, y0, x1, y1 = expanded roi_w = x1 - x0 + 1 roi_h = y1 - y0 + 1 need_w = max(0, int(min_size) - roi_w) need_h = max(0, int(min_size) - roi_h) if need_w == 0 and need_h == 0: return expanded extra = max((need_w + 1) // 2, (need_h + 1) // 2) return _expand_bbox(bbox, pad + extra, width, height) def _build_noise_prefill(img, mask_t, strength): """为mask区域生成噪声预填充，img取值范围为[-1, 1]。""" import torch img01 = (img + 1.0) / 2.0 unmasked = 1.0 - mask_t denom = unmasked.sum() if denom.item() < 1.0: mean = torch.full((1, 3, 1, 1), 0.5, device=img.device) std = torch.full((1, 3, 1, 1), 0.2, device=img.device) else: mean = (img01 * unmasked).sum(dim=(0, 2, 3), keepdim=True) / denom var = ((img01 - mean) ** 2 * unmasked).sum(dim=(0, 2, 3), keepdim=True) / denom std = torch.sqrt(var + 1e-6) noise = mean + std * float(strength) * torch.randn_like(img01) noise = noise.clamp(0.0, 1.0) return noise * 2.0 - 1.0 def _inpaint_with_aot_core( bgr, mask, aot_root, aot_pretrain, device="cpu", block_num=8, rates=None, noise_prefill=False, noise_strength=1.0, ): """使用 AOT-GAN 进行修补，返回 BGR 图像（与输入同尺寸）。""" if mask.dtype != np.uint8: mask = mask.astype(np.uint8) if mask.max() <= 1: mask = mask * 255 if mask.max() == 0: return bgr h, w = mask.shape grid = 4 h2 = (h // grid) * grid w2 = (w // grid) * grid if h2 == 0 or w2 == 0: return bgr rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) img_crop = rgb[:h2, :w2, :] mask_crop = mask[:h2, :w2] import torch img = torch.from_numpy(img_crop).permute(2, 0, 1).float() / 255.0 img = img * 2.0 - 1.0 mask_t = torch.from_numpy((mask_crop > 0).astype(np.float32)).unsqueeze(0) img = img.unsqueeze(0) mask_t = mask_t.unsqueeze(0) if device != "cpu": img = img.to(device) mask_t = mask_t.to(device) model, _ = _get_aot_model( aot_root=aot_root, aot_pretrain=aot_pretrain, device=device, block_num=block_num, rates=rates, ) with torch.no_grad(): if noise_prefill: noise = _build_noise_prefill(img, mask_t, noise_strength) image_masked = img * (1 - mask_t) + noise * mask_t else: image_masked = img * (1 - mask_t) + mask_t pred = model(image_masked, mask_t) comp = pred * mask_t + img * (1 - mask_t) comp = comp[0].clamp(-1.0, 1.0) comp = (comp + 1.0) / 2.0 * 255.0 comp = comp.permute(1, 2, 0).byte().cpu().numpy() result_bgr = cv2.cvtColor(comp, cv2.COLOR_RGB2BGR) if h2 == h and w2 == w: return result_bgr output_full = bgr.copy() output_full[:h2, :w2, :] = result_bgr return output_full def _inpaint_with_aot( bgr, mask, aot_root, aot_pretrain, device="cpu", block_num=8, rates=None, crop=False, crop_pad=0, max_size=0, noise_prefill=False, noise_strength=1.0, ): """使用 AOT-GAN 进行修补，支持裁剪与限幅以加速。""" min_side = 32 if mask.dtype != np.uint8: mask = mask.astype(np.uint8) if mask.max() <= 1: mask = mask * 255 if mask.max() == 0: return bgr h, w = mask.shape x0 = y0 = 0 x1 = w - 1 y1 = h - 1 if crop: bbox = _mask_bbox(mask) if bbox is None: return bgr expanded = _expand_bbox_min_size(bbox, crop_pad, w, h, min_side) if expanded is None: return bgr x0, y0, x1, y1 = expanded roi_w = x1 - x0 + 1 roi_h = y1 - y0 + 1 if roi_w < min_side or roi_h < min_side: # 裁剪区域过小会导致 AOT 失败，回退到全图修补 crop = False x0 = y0 = 0 x1 = w - 1 y1 = h - 1 bgr_roi = bgr[y0 : y1 + 1, x0 : x1 + 1] mask_roi = mask[y0 : y1 + 1, x0 : x1 + 1] roi_h, roi_w = bgr_roi.shape[:2] scale = 1.0 max_size = int(max_size) if max_size else 0 if max_size > 0: max_dim = max(roi_h, roi_w) if max_dim > max_size: scale = max_size / float(max_dim) new_w = max(1, int(round(roi_w * scale))) new_h = max(1, int(round(roi_h * scale))) if min(new_w, new_h) < min_side: scale = 1.0 new_w = roi_w new_h = roi_h interp = cv2.INTER_AREA if scale < 1.0 else cv2.INTER_LINEAR if scale != 1.0: bgr_roi = cv2.resize(bgr_roi, (new_w, new_h), interpolation=interp) mask_roi = cv2.resize(mask_roi, (new_w, new_h), interpolation=cv2.INTER_NEAREST) filled_roi = _inpaint_with_aot_core( bgr_roi, mask_roi, aot_root=aot_root, aot_pretrain=aot_pretrain, device=device, block_num=block_num, rates=rates, noise_prefill=noise_prefill, noise_strength=noise_strength, ) if scale != 1.0: filled_roi = cv2.resize( filled_roi, (roi_w, roi_h), interpolation=cv2.INTER_LINEAR, ) if not crop: return filled_roi output_full = bgr.copy() output_full[y0 : y1 + 1, x0 : x1 + 1] = filled_roi return output_full # 支持HEIC格式 try: from pillow_heif import register_heif_opener register_heif_opener() HEIC_SUPPORTED = True except ImportError: HEIC_SUPPORTED = False def remove_background(input_path, output_path, session=None, **kwargs): """ 去除图片背景 Args: input_path: 输入图片路径 output_path: 输出图片路径 session: rembg会话对象（可选） **kwargs: 其他参数，如alpha_matting相关参数 """ print(f"正在处理: {input_path}") # 读取输入图片 input_image = ImageOps.exif_transpose(Image.open(input_path)) # 去除背景 output_image = remove(input_image, session=session, **kwargs) # 保存输出图片 output_image.save(output_path) print(f"已保存: {output_path}") def _pil_to_bgr(image): """将PIL图片转换为OpenCV BGR格式""" rgb = image.convert("RGB") arr = np.array(rgb) return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR) def _alpha_to_mask(alpha_channel, threshold=10): """将alpha通道转换为二值mask（0或255）""" mask = (alpha_channel > threshold).astype(np.uint8) * 255 return mask def _ensure_rgba_size(image, target_size): if image.mode != "RGBA": image = image.convert("RGBA") if image.size != target_size: image = image.resize(target_size, resample=Image.NEAREST) return image def _prepare_mask(mask, mask_dilate=3, mask_blur=3, edge_grow=0): """生成硬边mask与处理后mask（用于填补/过渡）""" if mask_dilate and mask_dilate > 0: kernel = np.ones((mask_dilate, mask_dilate), np.uint8) mask = cv2.dilate(mask, kernel, iterations=1) if edge_grow and edge_grow > 0: kernel = np.ones((3, 3), np.uint8) mask = cv2.dilate(mask, kernel, iterations=int(edge_grow)) mask_hard = mask.copy() mask_used = mask_hard if mask_blur and mask_blur > 0: k = mask_blur if mask_blur % 2 == 1 else mask_blur + 1 mask_used = cv2.GaussianBlur(mask_hard, (k, k), 0) mask_used = (mask_used > 0).astype(np.uint8) * 255 return mask_hard, mask_used def remove_subject_and_inpaint( input_path, output_path, session=None, mask_dilate=3, mask_blur=3, mask_threshold=10, edge_grow=0, aot_root="AOT-GAN-for-Inpainting", aot_pretrain=None, aot_device="cpu", aot_block_num=8, aot_rates="1+2+4+8", aot_crop=False, aot_crop_pad=0, aot_max_size=0, aot_noise_prefill=False, aot_noise_strength=1.0, black_subject=False, black_threshold=50, gray_subject=False, gray_saturation_threshold=30, gray_value_threshold=200, feather=False, feather_radius=5, save_mask=False, mask_output_path=None, **kwargs, ): """ 去掉主体并补全背景 Args: input_path: 输入图片路径 output_path: 输出图片路径 session: rembg会话对象（可选） aot_root: AOT-GAN目录 aot_pretrain: AOT-GAN权重文件路径 aot_device: AOT-GAN设备 aot_block_num: AOT-GAN AOTBlock 数量 aot_rates: AOT-GAN AOTBlock 膨胀率 aot_crop: AOT仅裁剪mask区域进行修补 aot_crop_pad: AOT裁剪边缘留白像素 aot_max_size: AOT输入最大边限制（0为不限制） aot_noise_prefill: AOT使用随机噪声预填充 aot_noise_strength: 噪声强度系数 mask_dilate: mask膨胀大小 mask_blur: mask模糊大小（奇数） mask_threshold: alpha阈值 save_mask: 是否保存mask mask_output_path: mask保存路径 **kwargs: rembg参数 """ print(f"正在处理(去主体补背景): {input_path}") input_image = ImageOps.exif_transpose(Image.open(input_path)) # 使用rembg获取主体mask output_image = remove(input_image, session=session, **kwargs) output_image = _ensure_rgba_size(output_image, input_image.size) alpha = np.array(output_image.getchannel("A")) mask = _alpha_to_mask(alpha, threshold=mask_threshold) bgr = _pil_to_bgr(input_image) if black_subject: gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) black_mask = (gray <= black_threshold).astype(np.uint8) * 255 mask = cv2.bitwise_or(mask, black_mask) if gray_subject: hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) s = hsv[:, :, 1] v = hsv[:, :, 2] gray_mask = ((s <= gray_saturation_threshold) & (v <= gray_value_threshold)).astype(np.uint8) * 255 mask = cv2.bitwise_or(mask, gray_mask) mask_hard, mask_used = _prepare_mask( mask, mask_dilate=mask_dilate, mask_blur=mask_blur, edge_grow=edge_grow, ) rates = _parse_aot_rates(aot_rates) filled = _inpaint_with_aot( bgr, mask_used, aot_root=aot_root, aot_pretrain=aot_pretrain, device=aot_device, block_num=aot_block_num, rates=rates, crop=aot_crop, crop_pad=aot_crop_pad, max_size=aot_max_size, noise_prefill=aot_noise_prefill, noise_strength=aot_noise_strength, ) if feather and feather_radius > 0: # 仅在mask外侧做过渡，避免把原主体带回 mask_bin = (mask_hard > 0).astype(np.uint8) * 255 dist_out = cv2.distanceTransform(255 - mask_bin, cv2.DIST_L2, 3) alpha = np.ones_like(dist_out, dtype=np.float32) outside = mask_bin == 0 alpha[outside] = np.clip(1.0 - (dist_out[outside] / float(feather_radius)), 0.0, 1.0) alpha = alpha[:, :, None] blended = (alpha * filled + (1.0 - alpha) * bgr).astype(np.uint8) result = cv2.cvtColor(blended, cv2.COLOR_BGR2RGB) else: result = cv2.cvtColor(filled, cv2.COLOR_BGR2RGB) Image.fromarray(result).save(output_path) if save_mask: if mask_output_path is None: mask_output_path = str(Path(output_path).with_suffix("")) + "_mask.png" Image.fromarray(mask_used).save(mask_output_path) print(f"已保存mask: {mask_output_path}") print(f"已保存: {output_path}") def process_images_folder( input_folder, output_folder, model_name="u2net", alpha_matting=False, alpha_matting_foreground_threshold=240, alpha_matting_background_threshold=10, alpha_matting_erode_size=10, post_process_mask=False, remove_subject=False, mask_dilate=3, mask_blur=3, mask_threshold=10, edge_grow=0, aot_root="AOT-GAN-for-Inpainting", aot_pretrain=None, aot_device="cpu", aot_block_num=8, aot_rates="1+2+4+8", aot_crop=False, aot_crop_pad=0, aot_max_size=0, aot_noise_prefill=False, aot_noise_strength=1.0, black_subject=False, black_threshold=50, gray_subject=False, gray_saturation_threshold=30, gray_value_threshold=200, feather=False, feather_radius=5, save_mask=False, ): """ 批量处理文件夹中的所有图片 Args: input_folder: 输入文件夹路径 output_folder: 输出文件夹路径 model_name: 模型名称，可选值: - u2net (默认): 通用模型 - u2netp: 轻量版u2net - u2net_human_seg: 人物分割 - silueta: 精简版u2net (43MB) - isnet-general-use: 新的通用模型 - isnet-anime: 动漫角色高精度分割 - birefnet-general: 通用模型 - birefnet-portrait: 人像模型 alpha_matting: 是否启用alpha matting后处理（改善边缘质量） alpha_matting_foreground_threshold: 前景阈值 (0-255)，值越大保留越多前景 alpha_matting_background_threshold: 背景阈值 (0-255)，值越大去除越多背景 alpha_matting_erode_size: 侵蚀大小，用于平滑边缘 post_process_mask: 是否启用mask后处理 """ # 创建输出文件夹 Path(output_folder).mkdir(parents=True, exist_ok=True) # 创建会话（重用会话可以提高性能） print(f"使用模型: {model_name}") session = new_session(model_name) # 支持的图片格式 image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'} if HEIC_SUPPORTED: image_extensions.update({'.heic', '.heif'}) else: print("提示: 未安装pillow-heif，HEIC格式不可用。安装方法: pip install pillow-heif") # 获取所有图片文件 input_path = Path(input_folder) image_files = [f for f in input_path.iterdir() if f.is_file() and f.suffix.lower() in image_extensions] if not image_files: print(f"在 {input_folder} 中没有找到图片文件") return print(f"找到 {len(image_files)} 张图片，开始处理...") print(f"Alpha Matting: {'启用' if alpha_matting else '禁用'}") if alpha_matting: print(f" - 前景阈值: {alpha_matting_foreground_threshold}") print(f" - 背景阈值: {alpha_matting_background_threshold}") print(f" - 侵蚀大小: {alpha_matting_erode_size}") print(f"Mask后处理: {'启用' if post_process_mask else '禁用'}") print(f"去主体补背景: {'启用' if remove_subject else '禁用'}") if remove_subject: print(f" - AOT目录: {aot_root}") print(f" - AOT权重: {aot_pretrain}") print(f" - AOT设备: {aot_device}") print(f" - AOT块数: {aot_block_num}") print(f" - AOT膨胀率: {aot_rates}") print(f" - AOT裁剪: {'是' if aot_crop else '否'}") if aot_crop: print(f" - AOT裁剪边界: {aot_crop_pad}") print(f" - AOT最大边: {aot_max_size}") print(f" - AOT噪声预填充: {'是' if aot_noise_prefill else '否'}") if aot_noise_prefill: print(f" - AOT噪声强度: {aot_noise_strength}") print(f" - mask膨胀: {mask_dilate}") print(f" - mask模糊: {mask_blur}") print(f" - mask阈值: {mask_threshold}") print(f" - 边缘扩张: {edge_grow}") print(f" - 黑色内容作为主体: {'是' if black_subject else '否'}") if black_subject: print(f" - 黑色阈值: {black_threshold}") print(f" - 灰阶内容作为主体: {'是' if gray_subject else '否'}") if gray_subject: print(f" - 灰阶饱和度阈值: {gray_saturation_threshold}") print(f" - 灰阶亮度阈值: {gray_value_threshold}") print(f" - 边缘过渡: {'是' if feather else '否'}") if feather: print(f" - 过渡半径: {feather_radius}") print(f" - 保存mask: {'是' if save_mask else '否'}") print("-" * 50) # 处理每张图片 for i, image_file in enumerate(image_files, 1): try: if remove_subject: suffix = image_file.suffix.lower() if suffix in {".jpg", ".jpeg", ".png", ".bmp", ".webp"}: output_filename = image_file.stem + "_bgfill" + suffix else: output_filename = image_file.stem + "_bgfill.jpg" else: # 去背景默认使用PNG格式以支持透明背景 output_filename = image_file.stem + "_nobg.png" output_path = Path(output_folder) / output_filename print(f"[{i}/{len(image_files)}] ", end="") if remove_subject: mask_output_path = None if save_mask: mask_output_path = str(Path(output_folder) / (image_file.stem + "_mask.png")) remove_subject_and_inpaint( str(image_file), str(output_path), session=session, alpha_matting=alpha_matting, alpha_matting_foreground_threshold=alpha_matting_foreground_threshold, alpha_matting_background_threshold=alpha_matting_background_threshold, alpha_matting_erode_size=alpha_matting_erode_size, post_process_mask=post_process_mask, mask_dilate=mask_dilate, mask_blur=mask_blur, mask_threshold=mask_threshold, edge_grow=edge_grow, aot_root=aot_root, aot_pretrain=aot_pretrain, aot_device=aot_device, aot_block_num=aot_block_num, aot_rates=aot_rates, aot_crop=aot_crop, aot_crop_pad=aot_crop_pad, aot_max_size=aot_max_size, aot_noise_prefill=aot_noise_prefill, aot_noise_strength=aot_noise_strength, black_subject=black_subject, black_threshold=black_threshold, gray_subject=gray_subject, gray_saturation_threshold=gray_saturation_threshold, gray_value_threshold=gray_value_threshold, feather=feather, feather_radius=feather_radius, save_mask=save_mask, mask_output_path=mask_output_path, ) else: remove_background( str(image_file), str(output_path), session=session, alpha_matting=alpha_matting, alpha_matting_foreground_threshold=alpha_matting_foreground_threshold, alpha_matting_background_threshold=alpha_matting_background_threshold, alpha_matting_erode_size=alpha_matting_erode_size, post_process_mask=post_process_mask, ) except Exception as e: print(f"处理 {image_file.name} 时出错: {e}") print("-" * 50) print(f"处理完成！结果保存在 {output_folder} 文件夹中") if __name__ == "__main__": parser = argparse.ArgumentParser( description='图片去背景工具 - 使用rembg自动去除图片背景', formatter_class=argparse.RawDescriptionHelpFormatter, epilog=''' 示例用法: # 使用默认参数处理images文件夹 python remove_background.py # 处理单个文件 python remove_background.py input.jpg output.png # 处理指定文件夹 python remove_background.py my_images/ my_output/ # 使用不同模型 python remove_background.py input.jpg output.png -m birefnet-portrait # 自定义alpha matting参数 python remove_background.py input.jpg output.png -ft 260 -bt 12 -es 5 ''' ) # 必需参数 parser.add_argument('input', nargs='?', default='images', help='输入文件或文件夹路径（默认: images）') parser.add_argument('output', nargs='?', default=None, help='输出文件或文件夹路径（可选，默认为output/）') # 模型选择 parser.add_argument('-m', '--model', default='isnet-general-use', choices=['u2net', 'u2netp', 'u2net_human_seg', 'silueta', 'isnet-general-use', 'isnet-anime', 'birefnet-general', 'birefnet-general-lite', 'birefnet-portrait', 'birefnet-dis', 'birefnet-hrsod', 'birefnet-cod', 'birefnet-massive'], help='选择使用的模型 (默认: isnet-general-use)') # Alpha Matting参数 parser.add_argument('-a', '--alpha-matting', '--alpha_matting', action='store_true', help='启用alpha matting后处理（默认: false）') parser.add_argument('-ft', '--foreground-threshold', type=int, default=245, help='前景阈值 (0-255)，值越大保留越多细节 (默认: 245)') parser.add_argument('-bt', '--background-threshold', type=int, default=8, help='背景阈值 (0-255)，值越大去除越多背景 (默认: 8)') parser.add_argument('-es', '--erode-size', type=int, default=2, help='侵蚀大小，用于平滑边缘，值越大越平滑但可能丢失细节 (默认: 2)') # 其他选项 parser.add_argument('-p', '--post-process', '--post_process', action='store_true', help='启用mask后处理（默认: false）') # 去主体补背景参数 parser.add_argument('--remove-subject', '--remove_subject', action='store_true', help='去掉主体并补全背景（默认: false）') parser.add_argument('--aot-root', type=str, default='AOT-GAN-for-Inpainting', help='AOT-GAN目录（默认: AOT-GAN-for-Inpainting）') parser.add_argument('--aot-pretrain', type=str, default=None, help='AOT-GAN预训练权重文件路径（必填，相对路径基于当前目录）') parser.add_argument('--aot-device', type=str, default='cpu', help='AOT-GAN设备（默认: cpu）') parser.add_argument('--aot-block-num', type=int, default=8, help='AOTBlock数量（默认: 8）') parser.add_argument('--aot-rates', type=str, default='1+2+4+8', help='AOTBlock膨胀率（默认: 1+2+4+8）') parser.add_argument('--aot-crop', action='store_true', help='AOT仅对mask区域裁剪修补（默认: false）') parser.add_argument('--aot-crop-pad', type=int, default=0, help='AOT裁剪边缘留白像素（默认: 0）') parser.add_argument('--aot-max-size', type=int, default=0, help='AOT输入最大边限制，0为不限制（默认: 0）') parser.add_argument('--aot-noise-prefill', action='store_true', help='AOT使用随机噪声预填充（默认: false）') parser.add_argument('--aot-noise-strength', type=float, default=1.0, help='AOT噪声强度系数（默认: 1.0）') parser.add_argument('--mask-dilate', type=int, default=3, help='mask膨胀大小（默认: 3）') parser.add_argument('--mask-blur', type=int, default=3, help='mask模糊大小（默认: 3，建议奇数）') parser.add_argument('--mask-threshold', type=int, default=10, help='alpha阈值（默认: 10）') parser.add_argument('--edge-grow', type=int, default=0, help='主体边缘扩张像素（默认: 0）') parser.add_argument('--save-mask', '--save_mask', action='store_true', help='保存mask到output目录（默认: false）') parser.add_argument('--black-subject', '--black_subject', action='store_true', help='将黑色内容也视为主体（默认: false）') parser.add_argument('--black-threshold', type=int, default=50, help='黑色阈值(0-255，灰度越小越黑，默认: 50)') parser.add_argument('--gray-subject', '--gray_subject', action='store_true', help='将灰阶内容也视为主体（默认: false）') parser.add_argument('--gray-saturation-threshold', type=int, default=30, help='灰阶饱和度阈值(0-255，越小越接近灰阶，默认: 30)') parser.add_argument('--gray-value-threshold', type=int, default=200, help='灰阶亮度阈值(0-255，越小越暗，默认: 200)') parser.add_argument('--feather', action='store_true', help='启用边缘过渡融合（默认: false）') parser.add_argument('--feather-radius', type=int, default=5, help='边缘过渡半径（默认: 5）') args = parser.parse_args() print("=" * 50) print("图片去背景工具") print("=" * 50) # 判断输入是文件还是文件夹 input_path = Path(args.input) if not input_path.exists(): print(f"错误: 输入路径不存在: {args.input}") exit(1) # 处理单个文件 if input_path.is_file(): suffix = input_path.suffix.lower() if args.remove_subject: if suffix in {".jpg", ".jpeg", ".png", ".bmp", ".webp"}: output_name = input_path.stem + "_bgfill" + suffix else: output_name = input_path.stem + "_bgfill.jpg" else: output_name = input_path.stem + "_nobg.png" # 确定输出路径 if args.output is None: if args.remove_subject: output_path = input_path.parent / "output" / output_name else: output_path = input_path.parent / 'output' / output_name output_path.parent.mkdir(parents=True, exist_ok=True) else: output_candidate = Path(args.output) if output_candidate.exists() and output_candidate.is_dir(): output_path = output_candidate / output_name output_path.parent.mkdir(parents=True, exist_ok=True) elif output_candidate.suffix == "": output_candidate.mkdir(parents=True, exist_ok=True) output_path = output_candidate / output_name else: output_path = output_candidate output_path.parent.mkdir(parents=True, exist_ok=True) print(f"输入文件: {input_path}") print(f"输出文件: {output_path}") print(f"模型: {args.model}") print(f"Alpha Matting: {'启用' if args.alpha_matting else '禁用'}") if args.alpha_matting: print(f" - 前景阈值: {args.foreground_threshold}") print(f" - 背景阈值: {args.background_threshold}") print(f" - 侵蚀大小: {args.erode_size}") print(f"Mask后处理: {'启用' if args.post_process else '禁用'}") print(f"去主体补背景: {'启用' if args.remove_subject else '禁用'}") if args.remove_subject: print(f" - AOT目录: {args.aot_root}") print(f" - AOT权重: {args.aot_pretrain}") print(f" - AOT设备: {args.aot_device}") print(f" - AOT块数: {args.aot_block_num}") print(f" - AOT膨胀率: {args.aot_rates}") print(f" - AOT裁剪: {'是' if args.aot_crop else '否'}") if args.aot_crop: print(f" - AOT裁剪边界: {args.aot_crop_pad}") print(f" - AOT最大边: {args.aot_max_size}") print(f" - AOT噪声预填充: {'是' if args.aot_noise_prefill else '否'}") if args.aot_noise_prefill: print(f" - AOT噪声强度: {args.aot_noise_strength}") print(f" - mask膨胀: {args.mask_dilate}") print(f" - mask模糊: {args.mask_blur}") print(f" - mask阈值: {args.mask_threshold}") print(f" - 边缘扩张: {args.edge_grow}") print(f" - 黑色内容作为主体: {'是' if args.black_subject else '否'}") if args.black_subject: print(f" - 黑色阈值: {args.black_threshold}") print(f" - 灰阶内容作为主体: {'是' if args.gray_subject else '否'}") if args.gray_subject: print(f" - 灰阶饱和度阈值: {args.gray_saturation_threshold}") print(f" - 灰阶亮度阈值: {args.gray_value_threshold}") print(f" - 边缘过渡: {'是' if args.feather else '否'}") if args.feather: print(f" - 过渡半径: {args.feather_radius}") print(f" - 保存mask: {'是' if args.save_mask else '否'}") print("-" * 50) # 创建会话 session = new_session(args.model) # 处理图片 if args.remove_subject: mask_output_path = None if args.save_mask: mask_output_path = str(output_path.with_suffix("")) + "_mask.png" remove_subject_and_inpaint( str(input_path), str(output_path), session=session, alpha_matting=args.alpha_matting, alpha_matting_foreground_threshold=args.foreground_threshold, alpha_matting_background_threshold=args.background_threshold, alpha_matting_erode_size=args.erode_size, post_process_mask=args.post_process, mask_dilate=args.mask_dilate, mask_blur=args.mask_blur, mask_threshold=args.mask_threshold, edge_grow=args.edge_grow, aot_root=args.aot_root, aot_pretrain=args.aot_pretrain, aot_device=args.aot_device, aot_block_num=args.aot_block_num, aot_rates=args.aot_rates, aot_crop=args.aot_crop, aot_crop_pad=args.aot_crop_pad, aot_max_size=args.aot_max_size, aot_noise_prefill=args.aot_noise_prefill, aot_noise_strength=args.aot_noise_strength, black_subject=args.black_subject, black_threshold=args.black_threshold, gray_subject=args.gray_subject, gray_saturation_threshold=args.gray_saturation_threshold, gray_value_threshold=args.gray_value_threshold, feather=args.feather, feather_radius=args.feather_radius, save_mask=args.save_mask, mask_output_path=mask_output_path, ) else: remove_background( str(input_path), str(output_path), session=session, alpha_matting=args.alpha_matting, alpha_matting_foreground_threshold=args.foreground_threshold, alpha_matting_background_threshold=args.background_threshold, alpha_matting_erode_size=args.erode_size, post_process_mask=args.post_process, ) print("-" * 50) print(f"处理完成！结果保存在: {output_path}") # 处理文件夹 elif input_path.is_dir(): output_folder = args.output if args.output else 'output' process_images_folder( str(input_path), output_folder, model_name=args.model, alpha_matting=args.alpha_matting, alpha_matting_foreground_threshold=args.foreground_threshold, alpha_matting_background_threshold=args.background_threshold, alpha_matting_erode_size=args.erode_size, post_process_mask=args.post_process, remove_subject=args.remove_subject, mask_dilate=args.mask_dilate, mask_blur=args.mask_blur, mask_threshold=args.mask_threshold, edge_grow=args.edge_grow, aot_root=args.aot_root, aot_pretrain=args.aot_pretrain, aot_device=args.aot_device, aot_block_num=args.aot_block_num, aot_rates=args.aot_rates, aot_crop=args.aot_crop, aot_crop_pad=args.aot_crop_pad, aot_max_size=args.aot_max_size, aot_noise_prefill=args.aot_noise_prefill, aot_noise_strength=args.aot_noise_strength, black_subject=args.black_subject, black_threshold=args.black_threshold, gray_subject=args.gray_subject, gray_saturation_threshold=args.gray_saturation_threshold, gray_value_threshold=args.gray_value_threshold, feather=args.feather, feather_radius=args.feather_radius, save_mask=args.save_mask, ) else: print(f"错误: 不支持的输入类型: {args.input}") exit(1)