image_script.py

import os

import cv2
from tqdm import tqdm

from utils.stripe_code_utils import *
from utils.logger import *


def image_to_stripe_image(input_dir: str, output_dir: str, ExistingNotToConvert: bool = False, row: int = 50):
    """
    图片转换为条纹图像
    :param input_dir: 原图片所在文件夹
    :param output_dir: 条纹图像输出文件夹
    :param ExistingNotToConvert: 已存在图片不转换（用于断点续转）
    :param row: 转换条纹图像横向切割行数
    """
    # 遍历输入目录下所有的文件和子目录
    for root, dirs, files in os.walk(input_dir):
        # 遍历所有文件
        for file in tqdm(files):
            # 如果是图片（jpg or png）
            if file.endswith('.jpg') or file.endswith('.png'):
                # 构建完整的文件读取路径
                input_path = os.path.join(root, file)
                # 构建文件输出目录，os.path.relpath函数会返回input_path在input_dir中的相对路径
                output_path = os.path.join(output_dir, os.path.relpath(input_path, input_dir))
                # 判断是否需要转换
                if ExistingNotToConvert and os.path.exists(output_path):
                    continue
                # 图像读取和处理
                image = cv2.imread(input_path)  # 读取图像
                image = flit_average(image, row)  # 对图片做RGB均值切割
                image = flit_adaptive_threshold(image)  # 自适应二值化
                image = flit_average_gray(image, row)  # 对图片做灰度均值切割

                # 创建输出目录文件夹
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
                cv2.imwrite(output_path, image)


def count_images(input_dir: str):
    count = 0
    # 　遍历输入目录下所有的文件和子目录
    for root, dirs, files in os.walk(input_dir):
        # 遍历所有文件
        for file in files:
            # 如果是图片（jpg or png）
            if file.endswith('.jpg') or file.endswith('.png'):
                count += 1
    return count


r"""这是批量生成stripe_image的脚本，不要乱用"""
if __name__ == '__main__':
    logger = setting_logging("generate_stripe_image")
    logger.info("正在转换中，请稍等...")
    image_to_stripe_image(os.path.join('data', 'origin_data'),
                          os.path.join('data', 'stripe_data_100'), ExistingNotToConvert=True, row=100)
    print(count_images('./data/zebra_stripe_data_100'))
    # pass