201 lines
7.1 KiB
Python
201 lines
7.1 KiB
Python
from concurrent.futures import ThreadPoolExecutor
|
||
import random
|
||
import cv2
|
||
import numpy as np
|
||
import os
|
||
|
||
# 假设我们有一个函数 `remove_background` 使用某种方法去背景，返回前景掩码和前景图像
|
||
def remove_background(image_path):
|
||
# 加载图像
|
||
source_image = cv2.imread(image_path)
|
||
# 这里可以使用一个预训练的模型去背景，比如 U2-Net。为了简化，假设我们得到一个二值掩码
|
||
# 掩码生成逻辑可以替换为实际的模型推理
|
||
# 转换为灰度图像
|
||
GRAY = cv2.cvtColor(source_image, cv2.COLOR_BGR2GRAY)
|
||
|
||
# 二值化处理
|
||
_, mask_threshold = cv2.threshold(GRAY, 0, 1, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
|
||
|
||
# 定义结构元素
|
||
element = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
|
||
element1 = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
|
||
|
||
# 膨胀和腐蚀操作
|
||
mask_dilate = cv2.dilate(mask_threshold, element)
|
||
mask_erode = cv2.erode(mask_dilate, element1)
|
||
|
||
# 计算非零像素数量
|
||
count2 = cv2.countNonZero(mask_erode)
|
||
|
||
# 查找轮廓
|
||
contours, hierarchy = cv2.findContours(mask_erode, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
|
||
|
||
# 过滤轮廓
|
||
contours = [c for c in contours if cv2.contourArea(c) >= count2 * 0.3]
|
||
|
||
# 绘制轮廓
|
||
mask = np.zeros_like(mask_erode)
|
||
cv2.drawContours(mask, contours, -1, 1, -1)
|
||
|
||
# 将掩码转换为3通道
|
||
mask_cvtColor = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
|
||
|
||
# 应用掩码
|
||
source_image_multiply = cv2.multiply(source_image, mask_cvtColor)
|
||
|
||
# 转换为HSV颜色空间
|
||
imgHSV = cv2.cvtColor(source_image_multiply, cv2.COLOR_BGR2HSV)
|
||
|
||
# 定义HSV范围
|
||
scalarL = np.array([0, 46, 46])
|
||
scalarH = np.array([45, 255, 255])
|
||
|
||
# 根据HSV范围生成掩码
|
||
mask_inRange = cv2.inRange(imgHSV, scalarL, scalarH)
|
||
|
||
# 二值化处理
|
||
_, mask_tthreshold = cv2.threshold(mask_inRange, 0, 1, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
|
||
|
||
# 中值滤波
|
||
mask_medianBlur = cv2.medianBlur(mask_tthreshold, 7)
|
||
|
||
# 将掩码转换为3通道
|
||
mask_scvtColor = cv2.cvtColor(mask_medianBlur, cv2.COLOR_GRAY2BGR)
|
||
|
||
# 应用掩码
|
||
source_image = cv2.multiply(source_image, mask_scvtColor)
|
||
|
||
return source_image
|
||
|
||
def synthesize_background(foreground, background):
|
||
|
||
# 创建前景掩膜（非黑色区域）
|
||
gray = cv2.cvtColor(foreground, cv2.COLOR_BGR2GRAY)
|
||
_, mask = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY) # 阈值设为1以保留所有非纯黑像素
|
||
|
||
mask = cv2.GaussianBlur(mask, (5,5), 0) # 高斯模糊柔化边缘
|
||
_, mask = cv2.threshold(mask, 200, 255, cv2.THRESH_BINARY) # 重新二值化
|
||
|
||
# 精准形态学处理
|
||
kernel = np.ones((2,2), np.uint8)
|
||
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=1) # 闭运算填充小孔
|
||
|
||
# 反转掩膜用于获取背景区域
|
||
mask_inv = cv2.bitwise_not(mask)
|
||
|
||
# 提取背景和前景的ROI区域
|
||
background_roi = cv2.bitwise_and(background, background, mask=mask_inv)
|
||
foreground_roi = cv2.bitwise_and(foreground, foreground, mask=mask)
|
||
|
||
# 合成图像
|
||
result = cv2.add(foreground_roi, background_roi)
|
||
|
||
return result
|
||
|
||
def edge_fill2(img):
|
||
(height, width, p) = img.shape
|
||
H = 2384
|
||
W = 1560
|
||
top = bottom=int((W - height) / 2)
|
||
left= right= int((H - width) / 2)
|
||
img_result = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0)
|
||
return img_result
|
||
|
||
|
||
def process_images(input_folder, background_image_path, output_base):
|
||
"""
|
||
递归处理所有子文件夹并保持目录结构
|
||
"""
|
||
# 预处理背景路径（只需执行一次）
|
||
# if os.path.isfile(background_image_path):
|
||
# background_paths = [background_image_path]
|
||
# else:
|
||
# valid_ext = ['.jpg', '.jpeg', '.png', '.bmp', '.webp']
|
||
# background_paths = [
|
||
# os.path.join(background_image_path, f)
|
||
# for f in os.listdir(background_image_path)
|
||
# if os.path.splitext(f)[1].lower() in valid_ext
|
||
# ]
|
||
|
||
# 递归遍历输入目录
|
||
for root, dirs, files in os.walk(input_folder):
|
||
# 计算相对路径
|
||
relative_path = os.path.relpath(root, input_folder)
|
||
|
||
# 创建对应的输出目录
|
||
output_dir = os.path.join(output_base, relative_path)
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 处理当前目录的文件
|
||
for filename in files:
|
||
input_path = os.path.join(root, filename)
|
||
output_path = os.path.join(output_dir, filename)
|
||
|
||
# 跳过非图像文件
|
||
if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
|
||
continue
|
||
|
||
try:
|
||
# 去背景处理
|
||
result = remove_background(input_path)
|
||
|
||
|
||
# result = edge_fill2(result)
|
||
|
||
# 保存结果
|
||
cv2.imwrite(output_path, result)
|
||
print(f"Processed: {input_path} -> {output_path}")
|
||
|
||
except Exception as e:
|
||
print(f"Error processing {input_path}: {str(e)}")
|
||
|
||
def process_single_file(input_path, output_path):
|
||
"""处理单个文件的独立函数"""
|
||
try:
|
||
result = remove_background(input_path)
|
||
# result = edge_fill2(result) # 按需启用
|
||
cv2.imwrite(output_path, result)
|
||
print(f"Processed: {input_path} -> {output_path}")
|
||
except Exception as e:
|
||
print(f"Error processing {input_path}: {str(e)}")
|
||
|
||
def process_imageswithpool(input_folder, background_image_path, output_base):
|
||
"""
|
||
多线程版本的处理函数
|
||
使用ThreadPoolExecutor并行处理文件
|
||
"""
|
||
with ThreadPoolExecutor(max_workers=os.cpu_count()*2) as executor:
|
||
futures = []
|
||
for root, dirs, files in os.walk(input_folder):
|
||
# 创建输出目录（主线程保证目录创建顺序）
|
||
relative_path = os.path.relpath(root, input_folder)
|
||
output_dir = os.path.join(output_base, relative_path)
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 提交任务到线程池
|
||
for filename in files:
|
||
if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
|
||
continue
|
||
|
||
input_path = os.path.join(root, filename)
|
||
output_path = os.path.join(output_dir, filename)
|
||
futures.append(executor.submit(
|
||
process_single_file,
|
||
input_path,
|
||
output_path
|
||
))
|
||
|
||
# 可选：等待所有任务完成并处理异常
|
||
for future in futures:
|
||
try:
|
||
future.result()
|
||
except Exception as e:
|
||
print(f"Unhandled error in thread: {str(e)}")
|
||
|
||
|
||
# 使用示例
|
||
input_directory = 'L:/Grade_datasets/JY_A'
|
||
background_image_path = 'F:/dataset/02.TA_EC/rundata/BACKGROUND/ZY_B'
|
||
output_directory = 'L:/Grade_datasets/MOVE_BACKGROUND'
|
||
|
||
process_imageswithpool(input_directory, background_image_path, output_directory) |