import random import cv2 import numpy as np import os from tqdm import tqdm import shutil def correct_qr_code_images(input_folder, output_folder, margin=5): # 确保输出文件夹存在 if not os.path.exists(output_folder): os.makedirs(output_folder) # 读取输入文件夹中的所有图片 images = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.endswith(('.png', '.jpg', '.jpeg'))] total_images = len(images) # 确保WeChatQRCode类和模型文件正确加载 weChatQr = cv2.wechat_qrcode.WeChatQRCode( "wechat_qrcode/detect.prototxt", "wechat_qrcode/detect.caffemodel", "wechat_qrcode/sr.prototxt", "wechat_qrcode/sr.caffemodel" ) # 使用tqdm显示进度条 for i, img_path in enumerate(tqdm(images, desc="Processing images", total=total_images)): img = cv2.imread(img_path) # 灰度转换 gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) kernel = np.ones((1, 3), np.uint8) dilated_img = cv2.dilate(img, kernel, iterations=1) # 使用OTSU阈值算法将灰度图转换为二值图 _, binary_image = cv2.threshold(gray_img, 0, 255, cv2.THRESH_OTSU) # 黑白颜色反转 # binary_img_inv = cv2.bitwise_not(binary_image) # 显示反转后的二值图像 # cv2.imshow('Inverted Binary Image', dilated_img) # 等待按键后关闭窗口 # cv2.waitKey(0) # cv2.destroyAllWindows() res, points = weChatQr.detectAndDecode(dilated_img) if len(res) > 0: point = points[0] width = max(np.linalg.norm(point[0] - point[1]), np.linalg.norm(point[2] - point[3])) height = max(np.linalg.norm(point[1] - point[2]), np.linalg.norm(point[3] - point[0])) # 计算带有边距的新的points_dst points_dst = np.array([[margin, margin], [width - margin, margin], [width - margin, height - margin], [margin, height - margin]], dtype=np.float32) # 生成变换矩阵 matrix = cv2.getPerspectiveTransform(point, points_dst) # 应用透视变换，纠正图像 corrected_image = cv2.warpPerspective(img, matrix, (int(width), int(height))) # 保存裁剪后的二维码图像 output_path = os.path.join(output_folder, f"cropped_qr_code_{i + 1:05d}.jpg") cv2.imwrite(output_path, corrected_image) else: print(f"No QR code found in {img_path}.") def migrate_images(source_folder, target_folder): # 确保目标文件夹存在 if not os.path.exists(target_folder): os.makedirs(target_folder) # 初始化图片计数器 image_counter = 0 # 存储所有图片路径的列表 image_paths = [] # 遍历源文件夹及其所有子文件夹 for root, dirs, files in os.walk(source_folder): for dir in dirs: # 检查是否是名为'output'的文件夹 if dir == 'output': # 拼接完整的output文件夹路径 output_path = os.path.join(root, dir) # 遍历output文件夹中的所有文件 for file in os.listdir(output_path): # 检查是否是图片文件 if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff')): # 将图片路径添加到列表中 image_paths.append(os.path.join(output_path, file)) # 使用tqdm显示进度条 for i, image_path in enumerate(tqdm(image_paths, desc="Migrating images", unit="image")): # 从路径中提取文件名和扩展名 file_name = os.path.basename(image_path) # 构建目标文件的完整路径，使用i作为编号 target_file = os.path.join(target_folder, f"image_{i + 1:04d}{os.path.splitext(file_name)[1]}") # 复制图片到目标文件夹 shutil.copy2(image_path, target_file) image_counter += 1 print(f"Total images migrated: {image_counter}") def split_dataset(dataset_dir, output_dir, train_ratio=0.8): for label in ['real', 'fake']: label_dir = os.path.join(dataset_dir, label) images = os.listdir(label_dir) random.shuffle(images) train_size = int(len(images) * train_ratio) train_images = images[:train_size] val_images = images[train_size:] train_output_dir = os.path.join(output_dir, 'train', label) val_output_dir = os.path.join(output_dir, 'val', label) os.makedirs(train_output_dir, exist_ok=True) os.makedirs(val_output_dir, exist_ok=True) for image in train_images: shutil.copy(os.path.join(label_dir, image), os.path.join(train_output_dir, image)) for image in val_images: shutil.copy(os.path.join(label_dir, image), os.path.join(val_output_dir, image)) # 使用示例 input_folder = '/Users/jasonwong/Downloads/二维码测试1/手机验证码' # 替换为你的输入文件夹路径 output_folder = '/Users/jasonwong/Downloads/二维码测试1/手机验证码/output' # 替换为你的输出文件夹路径 correct_qr_code_images(input_folder, output_folder) # source_folder = '/Users/jasonwong/Downloads/二维码测试1/仿品' # 替换为你的源文件夹路径 # target_folder = '/Users/jasonwong/Downloads/二维码测试1/dataset/fake' # 替换为你的目标文件夹路径 # migrate_images(source_folder, target_folder) # # 使用脚本拆分数据集 # dataset_dir = '/Users/jasonwong/Downloads/二维码测试1/dataset' # 原始数据集路径 # output_dir = '/Users/jasonwong/Downloads/二维码测试1/data' # 拆分后的数据集路径 # split_dataset(dataset_dir, output_dir)