FORMAT_CONVERSION_MAXCOMPUTE/idc/pre_process.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jun 21 10:53:51 2022
pre_process.py
@author: fangjiasheng
"""
import copy
import json
import base64
import math
import os
import random
import sys
import traceback
from glob import glob
import numpy as np
import six
import cv2
from PIL import Image
import fitz
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
from idc.utils import pil_resize, pil_rotate

Image.MAX_IMAGE_PIXELS = 2300000000


def get_img_label(img_np, size, cls_num=4):
    height, width = size
    h, w = get_best_predict_size2(img_np, threshold=1080)
    img_np = pil_resize(img_np, h, w)

    # cv2.namedWindow("origin", 0)
    # cv2.resizeWindow("origin", 1000, 800)
    # cv2.imshow("origin", img_np)

    # 获取合适的文字区域
    result_list, img_np = get_text_region(img_np, size, 1)
    # print(len(result_list), img_np.shape)
    if not result_list:
        return []

    if img_np.shape[0] != height or img_np.shape[1] != width:
        img_np = pil_resize(img_np, height, width)

    # 生成旋转后的图片及其角度
    img_label_list = [[img_np, 0]]
    # 图片旋转
    angle_first = int(360/cls_num)
    i = 1
    for angle in range(angle_first, 360, angle_first):
        img_rotate = pil_rotate(img_np, angle)
        img_label_list.append([img_rotate, i])
        i += 1

    # for _img, _label in img_label_list:
    #     cv2.imshow("img", _img)
    #     cv2.waitKey(0)
    return img_label_list


def get_text_region(img_np, size, top_n=1):
    img_np = remove_black_border(img_np)
    origin_h, origin_w = img_np.shape[:2]
    gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)

    h, w = get_best_predict_size2(img_np, threshold=640)
    img_np = pil_resize(img_np, h, w)

    # 1.  转化成灰度图
    img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)

    # 2. 二值化
    ret, binary = cv2.threshold(img_np, 125, 255, cv2.THRESH_BINARY_INV)

    # 3. 膨胀和腐蚀操作的核函数
    kernal = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))

    # 4. 膨胀一次，让轮廓突出
    dilation = cv2.dilate(binary, kernal, iterations=2)
    dilation1 = copy.deepcopy(dilation)
    dilation1 = np.add(np.int0(np.full(dilation1.shape, 255)), -1 * np.int0(dilation))
    dilation1 = np.uint8(dilation1)
    # cv2.imshow("dilation1", dilation1)

    # 5. 腐蚀一次，去掉细节，如表格线等。注意这里去掉的是竖直的线
    erosion = cv2.erode(dilation, kernal, iterations=1)

    dilation = cv2.dilate(erosion, kernal, iterations=3)

    # 颜色反转
    dilation = np.add(np.int0(np.full(dilation.shape, 255)), -1 * np.int0(dilation))
    dilation = np.uint8(dilation)

    # 1. 查找轮廓
    contours, hierarchy = cv2.findContours(dilation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    region_list = []
    for cont in contours:
        rect = cv2.minAreaRect(cont)

        # box是四个点的坐标
        box = cv2.boxPoints(rect)
        box = np.int0(box)
        box = box.tolist()
        box.sort(key=lambda x: (x[0], x[1]))
        if box[0][0] > box[3][0]:
            temp = box[0][0]
            box[0][0] = box[3][0]
            box[3][0] = temp
        if box[0][1] > box[3][1]:
            temp = box[0][1]
            box[0][1] = box[3][1]
            box[3][1] = temp
        region_list.append(box)

    # dilation = np.expand_dims(dilation, axis=-1)
    # dilation = np.concatenate([dilation, dilation, dilation], axis=-1)
    # for box in region_list:
    # #     cv2.drawContours(dilation, [box], 0, (0, 255, 0), 2)
    #     cv2.rectangle(dilation, (box[0][0], box[0][1]), (box[3][0], box[3][1]), (0, 255, 0), 2)

    region_list.sort(key=lambda x: abs((x[3][0] - x[0][0])*(x[3][1] - x[0][1])), reverse=True)

    # 筛选文字区域
    result_list = []
    h_scale = origin_h / h
    w_scale = origin_w / w
    i = 0
    for box in region_list:
        if i >= top_n:
            break

        p1 = box[0]
        p2 = box[3]
        # print(p1, p2, abs((p2[1]-p1[1])*(p2[1]-p1[1])), h*w)

        # 旋转的box忽略
        if p1[0] >= p2[0] or p1[1] >= p2[1]:
            # print(p1[0], ">=", p2[0], p1[1], ">=", p2[1])
            continue

        # 太大的box忽略
        if abs(p2[0] - p1[0]) >= 0.7 * w and abs(p2[1] - p1[1]) >= 0.7 * h:
            # print("too large", abs(p2[0] - p1[0]), abs(p2[1] - p1[1]), 0.7 * w, 0.7 * h)
            continue

        # 黑色点不够的忽略
        if not count_black(dilation1[p1[1]:p2[1], p1[0]:p2[0]]):
            # print("black not enough")
            continue

        p1[1] = int(p1[1] * h_scale)
        p1[0] = int(p1[0] * w_scale)
        p2[1] = int(p2[1] * h_scale)
        p2[0] = int(p2[0] * w_scale)
        result_list.append([p1, p2])
        i += 1

    # cv2.imshow("result", dilation)
    # cv2.waitKey(0)

    if not result_list:
        return [], None

    # 裁剪
    if top_n > 1:
        result = random.sample(result_list, 1)[0]
        height_start = result[0][1]
        width_start = result[0][0]
    else:
        height_start = result_list[0][0][1]
        width_start = result_list[0][0][0]
    height, width = size
    gray = gray[height_start:height_start+height, width_start:width_start+width]

    # cv2.imshow("gray", gray)
    # cv2.waitKey(0)
    return result_list, gray


def gen(paths, batch_size=2, shape=(640, 640), cls_num=4, is_test=False):
    def choose(_paths, _i):
        while True:
            if _i >= len(_paths):
                _i = 0
                np.random.shuffle(paths)
            _p = _paths[_i]
            # read error
            try:
                _img_np = cv2.imread(_p)

                # limit h, w > 150
                if _img_np.shape[0] <= 150 or _img_np.shape[1] <= 150:
                    _i += 1
                    continue
                # limit pixels 89478485
                if _img_np.shape[0] * _img_np.shape[1] * _img_np.shape[2] >= 89478485:
                    _i += 1
                    continue

                _img_label_list = get_img_label(_img_np, size=(height, width), cls_num=cls_num)
                if not _img_label_list:
                    _i += 1
                    continue
            except:
                _i += 1
                continue
            _i += 1
            return _img_label_list, _i

    num = len(paths)

    i = 0
    all_cnt = 0
    while True:
        height, width = shape
        if is_test:
            X = np.zeros((batch_size, height, width, 1))
            Y = np.zeros((batch_size, cls_num))
        else:
            X = np.zeros((batch_size, height, width, 1))
            Y = np.zeros((batch_size, cls_num))
        img_np_list = []
        batch_list = []

        if batch_size % cls_num != 0:
            print("batch_size % cls_num != 0")
            raise

        for j in range(batch_size//cls_num):
            # 生成标注数据
            img_label_list, i = choose(paths, i)
            random.shuffle(img_label_list)

            if is_test:
                img_label_list = random.sample(img_label_list, 1)

            for c in range(cls_num):
                if c >= len(img_label_list):
                    break

                img = img_label_list[c][0]
                img_np_list.append(img)

                # 模糊
                if_blur = random.choice([0, 0])
                if if_blur:
                    # 高斯模糊
                    sigmaX = random.randint(1, 2)
                    sigmaY = random.randint(1, 2)
                    img = cv2.GaussianBlur(img, (5, 5), sigmaX, sigmaY)

                # label
                label_list = [0]*cls_num
                label_list[img_label_list[c][1]] = 1
                label = np.array(label_list)

                if len(img.shape) < 3:
                    img = np.expand_dims(img, axis=-1)
                img = np.expand_dims(img[:, :, 0], axis=-1)

                batch_list.append([img, label])

                # X[j+c] = img
                # Y[j+c] = label

        random.shuffle(batch_list)
        for j in range(0, batch_size, cls_num):
            for c in range(cls_num):
                img, label = batch_list[j+c]
                # print("label", label)
                # cv2.imshow("gen", img)
                # cv2.waitKey(0)
                # cv2.imwrite("data/3/"+str(all_cnt)+"_"+str(label)+".jpg", img)
                all_cnt += 1
                X[j+c] = img
                Y[j+c] = label

        if is_test:
            yield X, Y, img_np_list
        else:
            yield X, Y


def get_image_from_pdf():
    paths = glob("C:/Users/Administrator/Desktop/test_pdf/*")
    save_dir = "D:/Project/image_direction_classification/data/1/"

    i = 0
    for path in paths:
        try:
            doc = fitz.open(path)
            output_image_dict = {}
            page_count = doc.page_count
            for page_no in range(page_count):
                try:
                    page = doc.loadPage(page_no)
                    output = save_dir + "pdf_" + str(i) + ".png"
                    i += 1
                    rotate = int(0)
                    # 每个尺寸的缩放系数为1.3，这将为我们生成分辨率提高2.6的图像。
                    # 此处若是不做设置，默认图片大小为：792X612, dpi=96
                    # (1.33333333 --> 1056x816)   (2 --> 1584x1224)
                    # (1.183, 2.28 --> 1920x1080)
                    zoom_x = 1.3
                    zoom_y = 1.3
                    mat = fitz.Matrix(zoom_x, zoom_y).preRotate(rotate)
                    pix = page.getPixmap(matrix=mat, alpha=False)
                    pix.writePNG(output)
                except:
                    continue
        except Exception as e:
            print("pdf2Image", traceback.print_exc())
            continue


def get_best_predict_size2(image_np, threshold=640):
    h, w = image_np.shape[:2]
    scale = threshold / max(h, w)
    h = int(math.ceil(h * scale))
    w = int(math.ceil(w * scale))
    return h, w


def count_black(image_np):
    try:
        if len(image_np.shape) == 3:
            lower_black = np.array([0, 0, 0])
            upper_black = np.array([10, 10, 10])
        else:
            lower_black = np.array([0])
            upper_black = np.array([10])

        mask = cv2.inRange(image_np, lower_black, upper_black)
        cnt = np.sum(mask != 0)
        # print("black count", cnt, image_np.shape[0]*image_np.shape[1])
        if cnt >= image_np.shape[0]*image_np.shape[1]*0.3:
            return True
        else:
            return False
    except:
        return False


def remove_black_border(img_np):
    try:
        # 阈值
        threshold = 100

        # 转换为灰度图像
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)

        # 获取图片尺寸
        h, w = gray.shape[:2]

        # 无法区分黑色区域超过一半的情况
        rowc = gray[:, int(1/2*w)]
        colc = gray[int(1/2*h), :]

        rowflag = np.argwhere(rowc > threshold)
        colflag = np.argwhere(colc > threshold)

        left, bottom, right, top = rowflag[0, 0], colflag[-1, 0], rowflag[-1, 0], colflag[0, 0]

        # cv2.imshow('remove_black_border', img_np[left:right, top:bottom, :])
        # cv2.waitKey()
        return img_np[left:right, top:bottom, :]
    except:
        return img_np


if __name__ == '__main__':
    # get_img_label("data/0/7248_fe52d616989e19e6967e0461ef19b149.jpg", (640, 640))
    # get_image_from_pdf()

    # paths = glob("C:\\Users\\Administrator\\Desktop\\test_image\\*")
    # for path in paths:
    #     get_text_region(cv2.imread(path))

    # path = "C:\\Users\\Administrator\\Desktop\\test_image\\error17.jpg"
    # get_img_label(cv2.imread(path), (192, 192))

    table_image_path = "data/0/*"
    pdf_image_path = "data/1/*"
    no_table_image_path = 'data/2/*'
    paths = glob(table_image_path) + glob(pdf_image_path) + glob(no_table_image_path)
    _gen = gen(paths, batch_size=32, shape=(192, 192))

    for g in _gen:
        print("ok")
        break