DocumentAIKit/screenshot_save_by_label/cut_roi_save.py

import argparse
import datetime
from tqdm import tqdm
import json
import os
import cv2
import uuid


def read_jsonfile(self, path):
    with open(path, "r", encoding='utf-8') as f:
        return json.load(f)


def image_cut_save(path, left, upper, right, lower, save_path, pix):
    """
        所截区域图片保存
    :param pix: 将图片height缩放的pix
    :param path: 图片路径
    :param left: 区块左上角位置的像素点离图片左边界的距离
    :param upper：区块左上角位置的像素点离图片上边界的距离
    :param right：区块右下角位置的像素点离图片左边界的距离
    :param lower：区块右下角位置的像素点离图片上边界的距离
     故需满足：lower > upper、right > left
    :param save_path: 所截图片保存位置
    """
    img = cv2.imread(path)  # 打开图像
    cropped = img[upper:lower, left:right]
    if pix != 0:
        height, width = cropped.shape[:2]  # 原始分辨率
        if height != 0 and width != 0:
            while height < pix or width < pix:
                if height <= width:
                    # 等比例缩放到pix
                    scale = pix / height
                    # 缩放后分辨率
                    height = pix
                    width = int(width * scale)
                else:
                    # 等比例缩放到pix
                    scale = pix / width
                    # 缩放后分辨率
                    width = pix
                    height = int(height * scale)
            img = cv2.resize(cropped, (width, height))
            # 保存截取的图片
            cv2.imwrite(save_path, img)
    else:
        cv2.imwrite(save_path, cropped)


def single_img_cut(json_path, label, img_save_path, cnt, pix):
    with open(json_path, "r", encoding='utf-8') as f:
        cnt1 = 1
        jf = json.load(f)
        shapes = jf["shapes"]
        for shape in shapes:
            if shape["label"] == label:
                time = datetime.datetime.now()
                t = str(time.month).zfill(2) + str(time.day).zfill(2) + str(time.hour).zfill(2) + str(
                    time.minute).zfill(2) + \
                    str(time.second).zfill(2)
                image_cut_save(json_path.replace('json', 'jpg'), int(shape["points"][0][0]), int(shape["points"][0][1]),
                               int(shape["points"][1][0]), int(shape["points"][1][1]),
                               img_save_path + '/' + 'kdan' + '_' + t + '_' + str(uuid.uuid1())[0:8] + '_' + str(
                                   cnt).zfill(5) + str(cnt1).zfill(3) + '.jpg', pix)
                cnt1 += 1
    f.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_dirs', type=str, nargs="+", default=['./'])
    parser.add_argument('--label', type=str, default='Table')
    parser.add_argument('--save_dir', type=str, default='./imgs')
    parser.add_argument('--size', type=int, default='0')
    args = parser.parse_args()
    cnt = 1
    input_list = args.input_dirs
    # 文件夹列表
    json_list = []
    # 遍历input_list
    for jp in input_list:
        json_img_list = os.listdir(jp)
        for path in json_img_list:
            if path.endswith('.json'):
                path = os.path.join(jp, path)
                json_list.append(path)
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    for path in tqdm(json_list):
        single_img_cut(path, args.label, args.save_dir, cnt, args.size)
        cnt += 1