import numpy as np import os import matplotlib.image as mpimage import argparse import functools from utils import add_arguments, print_arguments from dask.distributed import LocalCluster from dask import bag as dbag from dask.diagnostics import ProgressBar from typing import Tuple from PIL import Image import cv2 #-----------------------------------# # 对四个点坐标排序 #-----------------------------------# def order_points(pts): # 一共4个坐标点 rect = np.zeros((4, 2), dtype = "float32") # 按顺序找到对应坐标0123分别是 左上，右上，右下，左下 # 计算左上，右下 s = pts.sum(axis = 1) rect[0] = pts[np.argmin(s)] rect[2] = pts[np.argmax(s)] # 计算右上和左下 diff = np.diff(pts, axis = 1) rect[1] = pts[np.argmin(diff)] rect[3] = pts[np.argmax(diff)] return rect #-----------------------------------# # 透射变换纠正车牌图片 #-----------------------------------# def four_point_transform(image, pts): # 获取输入坐标点 rect = order_points(pts) (tl, tr, br, bl) = rect # 计算输入的w和h值 widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2)) widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2)) maxWidth = max(int(widthA), int(widthB)) heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2)) heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2)) maxHeight = max(int(heightA), int(heightB)) # 变换后对应坐标位置 dst = np.array([ [0, 0], [maxWidth - 1, 0], [maxWidth - 1, maxHeight - 1], [0, maxHeight - 1]], dtype = "float32") # 计算变换矩阵 M = cv2.getPerspectiveTransform(rect, dst) warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight)) # 返回变换后结果 return warped # Dataset statistics that I gathered in development #-----------------------------------# # 用于过滤感知质量较低的不良图片 #-----------------------------------# IMAGE_MEAN = 0.5 IMAGE_MEAN_STD = 0.028 IMG_STD = 0.28 IMG_STD_STD = 0.01 def readImage(fileName: str) -> np.ndarray: image = mpimage.imread(fileName) return image #-----------------------------------# # 从文件名中提取车牌的坐标 #-----------------------------------# def parseLabel(label: str) -> Tuple[np.ndarray, np.ndarray]: annotation = label.split('-')[3].split('_') coor1 = [int(i) for i in annotation[0].split('&')] coor2 = [int(i) for i in annotation[1].split('&')] coor3 = [int(i) for i in annotation[2].split('&')] coor4 = [int(i) for i in annotation[3].split('&')] coor = np.array([coor1, coor2, coor3, coor4]) center = np.mean(coor, axis=0) return coor, center.astype(int) #-----------------------------------# # 根据车牌坐标裁剪出车牌图像 #-----------------------------------# # def cropImage(image: np.ndarray, coor: np.ndarray, center: np.ndarray) -> np.ndarray: # image = four_point_transform(image, coor) # return image def cropImage(image: np.ndarray, coor: np.ndarray, center: np.ndarray) -> np.ndarray: maxW = np.max(coor[:, 0] - center[0]) # max plate width maxH = np.max(coor[:, 1] - center[1]) # max plate height xWanted = [64, 128, 192, 256] yWanted = [32, 64, 96, 128] found = False for w, h in zip(xWanted, yWanted): if maxW < w//2 and maxH < h//2: maxH = h//2 maxW = w//2 found = True break if not found: # plate too large, discard return np.array([]) elif center[1]-maxH < 0 or center[1]+maxH >= image.shape[1] or \ center[0]-maxW < 0 or center[0] + maxW >= image.shape[0]: return np.array([]) else: return image[center[1]-maxH:center[1]+maxH, center[0]-maxW:center[0]+maxW] #-----------------------------------# # 保存车牌图片 #-----------------------------------# def saveImage(image: np.ndarray, fileName: str, outDir: str) -> int: if image.shape[0] == 0: return 0 else: imgShape = image.shape if imgShape[1] == 64: mpimage.imsave(os.path.join(outDir, '64_32', fileName), image) elif imgShape[1] == 128: mpimage.imsave(os.path.join(outDir, '128_64', fileName), image) elif imgShape[1] == 208: mpimage.imsave(os.path.join(outDir, '192_96', fileName), image) else: #resize large images image = Image.fromarray(image).resize((192, 96)) image = np.asarray(image) # back to numpy array mpimage.imsave(os.path.join(outDir, '192_96', fileName), image) return 1 #-----------------------------------# # 包装成一个函数，以便将处理区分到不同目录 #-----------------------------------# def processImage(file: str, inputDir: str, outputDir: str, subFolder: str) -> int: result = parseLabel(file) filePath = os.path.join(inputDir,subFolder, file) image = readImage(filePath) plate = cropImage(image, result[0], result[1]) if plate.shape[0] == 0: return 0 mean = np.mean(plate/255.0) std = np.std(plate/255.0) # bad brightness if mean <= IMAGE_MEAN - 10*IMAGE_MEAN_STD or mean >= IMAGE_MEAN + 10*IMAGE_MEAN_STD: return 0 # low contrast if std <= IMG_STD - 10*IMG_STD_STD: return 0 status = saveImage(plate, file, outputDir) return status def main(argv): jobNum = int(argv.jobNum) outputDir = argv.outputDir inputDir = argv.inputDir try: os.mkdir(outputDir) for shape in ['64_32', '128_64', '192_96']: os.mkdir(os.path.join(outputDir, shape)) except OSError: pass # path already exists client = LocalCluster(n_workers=jobNum, threads_per_worker=5) # IO intensive, more threads print('* number of workers:{}, \n* input dir:{}, \n* output dir:{}\n\n'.format(jobNum, inputDir, outputDir)) for subFolder in ['ccpd_green', 'ccpd_base', 'ccpd_db', 'ccpd_fn', 'ccpd_rotate', 'ccpd_tilt', 'ccpd_weather']: fileList = os.listdir(os.path.join(inputDir, subFolder)) print('* {} images found in {}. Start processing ...'.format(len(fileList), subFolder)) toDo = dbag.from_sequence(fileList, npartitions=jobNum*30).persist() # persist the bag in memory toDo = toDo.map(processImage, inputDir, outputDir, subFolder) pbar = ProgressBar(minimum=2.0) pbar.register() # register all computations for better tracking result = toDo.compute() print('* image cropped: {}. Done ...'.format(sum(result))) client.close() # shut down the cluster if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) add_arg = functools.partial(add_arguments, argparser=parser) add_arg('jobNum', int, 4, '处理图片的线程数') add_arg('inputDir', str, 'datasets/CCPD2020', '输入图片目录') add_arg('outputDir', str, 'datasets/CCPD2020_new', '保存图片目录') args = parser.parse_args() print_arguments(args) main(args)