149 lines
5.8 KiB
Python
149 lines
5.8 KiB
Python
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) |