测试

# -*- coding: utf-8 -*-
# @Time : 2021/4/12 15:00
# @Author : Sonny
# @Site : zhzy
# @File : gnt_to_png.py
# @Software: PyCharm
# @Description:
import shutil
from time import sleep
from pathlib import Path
from tqdm import tqdm
from PIL import Image
import numpy as np
import struct
import os
import click
import re


def gnt_to_png_with_label(gnt, out_dir=None, label=None, extra=""):
    if not out_dir:
        out_dir = gnt.replace('.gnt', '-dataSet')
        # out_dir = os.path.join(os.path.dirname(gnt), os.path.splitext(os.path.basename(gnt))[0])
    if not label:
        label = os.path.join(out_dir, 'label.txt')
    # shutil.rmtree(out_dir, ignore_errors=True)
    os.makedirs(out_dir, exist_ok=True)
    with open(gnt, 'rb') as f, open(label, 'a+', encoding='utf-8') as w, tqdm(desc=f'{extra}{os.path.basename(gnt)}', leave=True) as bar: # 读入一份gnt文件
        header_size = 10
        while True:
            header = np.fromfile(f, dtype='uint8', count=header_size)
            if not header.size:
                break
            sample_size = header[0] + (header[1] << 8) + (header[2] << 16) + (header[3] << 24)
            tagcode = header[5] + (header[4] << 8)
            width = header[6] + (header[7] << 8)
            height = header[8] + (header[9] << 8)
            if header_size + width * height != sample_size:
                break
            image = np.fromfile(f, dtype='uint8', count=width * height).reshape((height, width))
            tagcode_unicode = struct.pack('>H', tagcode).decode('gbk').strip(b'\x00'.decode()) # 图片对应的汉字

            im = Image.fromarray(image)
            number = len(os.listdir(out_dir))
            name = f'{number}.png'.rjust(10, '0')
            try:
                im.convert('RGB').save(os.path.join(out_dir, name))
                w.write(name + '\t' + tagcode_unicode + '\n')
            except Exception:
                continue

            bar.update()


@click.command()
@click.option('--folder', help='gnt file path', prompt=True)
def convert_folder(folder):
    foo = Path(folder)
    files = list(foo.rglob('*.gnt'))
    files.sort(key=lambda f: int(re.search(r'^\d+', f.name).group(0)) if re.search(r'^\d+', f.name) else 0)
    num = 1
    for file in files:
        out_dir = file.parent.parent.joinpath('pngs', (re.search(r'^\d+', file.name).group(0) if re.search(r'^\d+', file.name) else "0").rjust(4, '0'))
        gnt_to_png_with_label(str(file), out_dir=str(out_dir), extra=f'{num}/{len(files)} ')
        num += 1


if __name__ == '__main__':

    # gnt_dir = r'E:\Projects\DeepLearning\data-set\Gnt\Gnt1.0TrainPart1'
    #
    # for file in os.listdir(gnt_dir):
    #     gnt = os.path.join(gnt_dir, file)
    #     if gnt.endswith('.gnt'):
    #         gnt_to_png_with_label(gnt)

    convert_folder()