DocumentAIKit/correct_imgs_rotation/img_correct.py

import argparse
import os

import cv2
import numpy as np
from PIL import Image
import glob


def save_img(img_Path, save_Path, cnt):
    img_path = img_Path
    img = cv2.imread(img_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    ret, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
    coords = np.column_stack(np.where(thresh > 0))
    angle = cv2.minAreaRect(coords)[-1]
    print('{}:{}'.format(img_path, angle))
    # #调整角度
    if angle < -45:
        angle = -(90 + angle)
    else:
        angle = -angle
    im = Image.open(img_path)
    im_rotate = im.rotate(angle, expand=0, fillcolor='#FFFFFF')
    if cnt == 0:
        save_path = save_Path + '/' + img_path.split('\\')[-1]
    else:
        save_path = save_Path + '/' + img_path.split('/')[-1]
    im_rotate.save(save_path)

def correct(img_fold, img_Path, save_fold):
    if not os.path.exists(save_fold):
        os.makedirs(save_fold)
    if img_Path == '':
        path_list = glob.glob(img_fold + '/*.png')
        size = len(path_list)
        print('total: {} images'.format(size))
        for path in path_list:
            save_img(path, save_fold, 0)
    else:
        save_img(img_Path, save_fold, 1)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_dir', type=str, default='')
    parser.add_argument('--img_path', type=str, default='')
    parser.add_argument('--save_dir', type=str, default='')
    args = parser.parse_args()
    correct(args.img_dir, args.img_path, args.save_dir)