abupy.DLBu.ABuDLImgStd 源代码

# -*- encoding:utf-8 -*-
"""
    深度学习工具模块,为caffe等工具库标准化图片格式
"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import glob
import imghdr
import os

import PIL.Image
from PIL import ImageFile

# noinspection PyUnresolvedReferences
from ..CoreBu.ABuFixes import map

from ..UtilBu import ABuFileUtil

__all__ = ['std_img_from_root_dir',
           'covert_to_jpeg',
           'find_img_by_ext',
           'change_to_real_type'
           ]

__author__ = '阿布'
__weixin__ = 'abu_quant'

# 为了im.convert('RGB')的异常错误,需要设置ImageFile.LOAD_TRUNCATED_IMAGES = True
ImageFile.LOAD_TRUNCATED_IMAGES = True


[文档]def std_img_from_root_dir(root_dir, a_ext): """ 找到root_dir目录下所有.ext的图片,都转换其真实类型后, 将不是jpeg的全部转换为jpeg,即标准化图像格式 :param root_dir: str对象,文件夹路径信息 :param a_ext: 图像文件后缀,eg. png,jpg """ # root_dir目录下及一级子目录下所有.ext后缀的文件全路径返回 img_list = find_img_by_ext(a_ext, root_dir) # img_list中所有的转换为其真实类型 all_type = change_to_real_type(img_list) for ext in all_type: # 迭代change_to_real_type返回的集合 if ext != 'jpeg': # 标准化所有输入图片,将不是jpeg的都转换为jpeg if ext is None: ext = a_ext # 找到非标准类型子序列sub_img_list sub_img_list = find_img_by_ext(ext, root_dir) # 一个个进行转换 _ = list(map(lambda img: covert_to_jpeg(img), sub_img_list)) # 然后再次转换其真实后缀名 change_to_real_type(sub_img_list)
[文档]def covert_to_jpeg(org_img, dst_img=None): """ 将输入img转换为RGB的jpeg格式图像 :param org_img: 原始图像路径,str对象 :param dst_img: 转换后的图像路径输出路径,str对象,默认=None将覆盖org_img路径 """ im = PIL.Image.open(org_img) if dst_img is None: # 不设置输出路径,即在输入文件上直接覆盖 dst_img = org_img im.convert('RGB').save(dst_img, 'JPEG')
[文档]def find_img_by_ext(ext, root_dir): """ 将root_dir目录下及一级子目录下所有.ext后缀的文件全路径返回, 注意只遍历根目录及一级子目录 :param ext: str对象 图像文件后缀,eg. png,jpg :param root_dir: str对象,文件夹路径信息 :return: list序列对象 """ # 遍历一级子目录 dirs = [root_dir + name for name in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, name))] # 形成根目录及一级子目录 dirs.append(root_dir) img_list = list() for dr in dirs: # 遍历所有目录,通过glob模糊查询所有.ext文件 sub_list = glob.glob('{}/*.{}'.format(dr, ext)) img_list.extend(sub_list) return img_list
[文档]def change_to_real_type(img_list): """ 将img的后缀名转换为其真实类型 eg. a.png 如果 a.png实际上是jpeg,则将后缀修改-> a.jpeg :param img_list: :return: """ # 类型记录集合 record_type = set() for img in img_list: if not ABuFileUtil.file_exist(img): # 过滤实际不存在的文件 continue # 使用imghdr识别图像真实类型 real_type = imghdr.what(img) # 将img_list有的类型做记录,add到集合中 record_type.add(real_type) if real_type is None: # 识别不出来,可能就是残图,过滤 continue # 修改成真实的类型 real_name = img[:img.rfind('.')] + '.' + real_type os.rename(img, real_name) return record_type