TA_EC/dataset/recover.py

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)