多线程检测图像是不是网页截图v3.0

# python data/utils/select_none_web_v2.py
# 多线程检测是不是网页截图

import cv2
import pytesseract
import numpy as np
import os
import shutil
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

# 设置Tesseract路径
pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract'

# 图像预处理和特征提取
def extract_features(image_path):
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 文本检测
    text = pytesseract.image_to_string(gray)
    text_len = len(text)
    
    # 颜色分布特征
    color_hist = cv2.calcHist([gray], [0], None, [256], [0,256])
    color_hist = color_hist.flatten()
    
    # 计算总像素数
    total_pixels = np.prod(img.shape[:2])
    
    return text_len, color_hist, total_pixels

# 筛选规则
def is_webpage_screenshot(image_path, min_text_length=100, max_text_density=0.05):
    text_len, _, total_pixels = extract_features(image_path)
    
    # 文本长度阈值
    if text_len < min_text_length:
        return False
    
    # 颜色密度阈值，假设文本区域颜色分布较为集中
    text_density = text_len / total_pixels
    
    if text_density > max_text_density:
        return False
    
    return True

# 设置路径
source_dir = "data/aigc_0"
dest_dir = "data/aigc_0_non"

# 确保目标目录存在
if not os.path.exists(dest_dir):
    os.makedirs(dest_dir)

# 应用筛选规则到未知图片，并移动非网页截图
image_paths = [os.path.join(source_dir, filename) for filename in os.listdir(source_dir)]

# 自定义线程数
num_threads = 10  # 根据你的系统资源进行调整

# 使用ThreadPoolExecutor进行多线程处理
with ThreadPoolExecutor(max_workers=num_threads) as executor:
    futures = {executor.submit(is_webpage_screenshot, path): path for path in image_paths}
    
    # 使用tqdm显示进度
    for future in tqdm(as_completed(futures), total=len(image_paths), desc="Processing images"):
        path = futures[future]
        try:
            result = future.result()
        except Exception as e:
            print(f"An error occurred while processing {path}: {e}")
        else:
            if not result:
                shutil.move(path, dest_dir)