Spaces:
Build error
Build error
| 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 | |
| # 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: | |
| 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: # 车牌太大则丢弃 | |
| 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) | |
| # 亮度不好的 | |
| if mean <= IMAGE_MEAN - 10*IMAGE_MEAN_STD or mean >= IMAGE_MEAN + 10*IMAGE_MEAN_STD: | |
| return 0 | |
| # 低对比度的 | |
| 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 # 地址已经存在 | |
| client = LocalCluster(n_workers=jobNum, threads_per_worker=5) # 开启多线程 | |
| for subFolder in ['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() # 登记所有的计算，以便更好地跟踪 | |
| result = toDo.compute() | |
| print('* image cropped: {}. Done ...'.format(sum(result))) | |
| client.close() # 关闭集群 | |
| 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/CCPD2019', '输入图片目录') | |
| add_arg('outputDir', str, 'datasets/CCPD2019_new', '保存图片目录') | |
| args = parser.parse_args() | |
| print_arguments(args) | |
| main(args) | |