xzbautomarkapi/admin/monitor/recpaper.py

# -*-coding:utf-8 -*-

import os

import cv2
import numpy as np
from imutils.perspective import four_point_transform


class RecPaper:
    """识别出试卷上的类答题卡区域
    """

    def __init__(self, image_path="test.jpg", size="A4", resize=False):
        self.image_path = image_path
        self.image = cv2.imread(self.image_path)

        # 图片的宽高度
        self.width = self.image.shape[1]
        self.height = self.image.shape[0]

        # 图片转成灰度图
        self.gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)

        # 边缘检测后的图
        self.edged = cv2.Canny(self.gray, 75, 200)

        # 对灰度图进行二值化处理
        ret, thresh = cv2.threshold(self.gray, 0, 250, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
        self.thresh = thresh

        if resize:
            if size == "A4":
                self.thresh = cv2.resize(thresh, (2480, 3480), interpolation=cv2.INTER_LANCZOS4)
            else:
                self.thresh = cv2.resize(thresh, (4870, 3480), interpolation=cv2.INTER_LANCZOS4)

        # 输出二值化操作后的图像
        spt_lst = os.path.splitext(self.image_path)
        close_path = spt_lst[0] + '_thresh' + spt_lst[1]
        cv2.imwrite(close_path, self.thresh)

    def transform(self, point4):
        """透视变换
        """
        four = [[point4[0]["x"], point4[0]["y"]], [point4[1]["x"], point4[1]["y"]],
                [point4[3]["x"], point4[3]["y"]], [point4[2]["x"], point4[2]["y"]]]
        docCnt = np.array(four)
        warped = four_point_transform(self.image, docCnt.reshape(4, 2))

        warped = cv2.resize(warped, (2480, 3480), interpolation=cv2.INTER_LANCZOS4)

        spt_lst = os.path.splitext(self.image_path)
        trpath = spt_lst[0] + '_transform' + spt_lst[1]
        cv2.imwrite(trpath, warped)
        return trpath

    def _erode(self, w=None, h=None):
        '''
        图像腐蚀操作
        '''
        # 腐蚀核大小
        edsize = cv2.getStructuringElement(cv2.MORPH_RECT, (w, h))
        self.erode_image = cv2.erode(self.thresh, edsize)
        # 输出腐蚀操作后的图像
        spt_lst = os.path.splitext(self.image_path)
        close_path = spt_lst[0] + '_erode' + spt_lst[1]
        cv2.imwrite(close_path, self.erode_image)
        return self.erode_image

    def _erode_dilate(self, w=None, h=None):
        '''
        图像腐蚀操作
        '''
        # 腐蚀核大小
        edsize = cv2.getStructuringElement(cv2.MORPH_RECT, (w, h))
        self.erode_dilate_image = cv2.dilate(self.erode_image, edsize)
        # 输出腐蚀操作后的图像
        spt_lst = os.path.splitext(self.image_path)
        close_path = spt_lst[0] + '_erode_dilate' + spt_lst[1]
        cv2.imwrite(close_path, self.erode_dilate_image)
        return self.erode_dilate_image

    def _dilate(self, w=None, h=None):
        '''
        图像膨胀操作
        '''
        dilsize = cv2.getStructuringElement(cv2.MORPH_RECT, (w, h))
        self.dilate_image = cv2.dilate(self.thresh, dilsize)
        # 输出膨胀操作后的图像
        spt_lst = os.path.splitext(self.image_path)
        close_path = spt_lst[0] + '_dilate' + spt_lst[1]
        cv2.imwrite(close_path, self.dilate_image)
        return self.dilate_image

    def _dilate_erode(self, w=None, h=None):
        '''
        图像膨胀操作
        '''
        dilsize = cv2.getStructuringElement(cv2.MORPH_RECT, (w, h))
        self.dilate_erode_image = cv2.erode(self.dilate_image, dilsize)
        # 输出膨胀操作后的图像
        spt_lst = os.path.splitext(self.image_path)
        close_path = spt_lst[0] + '_dilate_erode' + spt_lst[1]
        cv2.imwrite(close_path, self.dilate_erode_image)
        return self.dilate_erode_image

    def _open(self, w=None, h=None):
        """开操作先腐蚀后膨胀
        """
        self._erode(w, h)
        # 自定义开操作
        self.open_image = self._erode_dilate(w, h)
        # 输出闭操作后的图像
        spt_lst = os.path.splitext(self.image_path)
        close_path = spt_lst[0] + '_open' + spt_lst[1]
        cv2.imwrite(close_path, self.open_image)
        return self.open_image

    def _close(self, w=20, h=40):
        """闭操作先膨胀后腐蚀
        """
        self._dilate()
        #自定义闭操作
        self.close_image = self._dilate_erode(w, h)
        # 输出闭操作后的图像
        spt_lst = os.path.splitext(self.image_path)
        close_path = spt_lst[0] + '_close' + spt_lst[1]
        cv2.imwrite(close_path, self.close_image)
        return self.close_image

    def get_cnts(self, std_w=None, std_h=None):
        """识别出答案轮廓
        """
        self._open(std_w, std_h)
        cnts = cv2.findContours(self.open_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[0]
        all_c = []
        if len(cnts) > 0:
            cnts = sorted(cnts, key=cv2.contourArea, reverse=True)

            all_c = []
            for c in cnts:
                peri = std_h / 2
                approx = cv2.approxPolyDP(c, peri, True)
                if len(approx) == 4 or True:
                    (x, y, w, h) = cv2.boundingRect(c)
                    # 标识出识别出的机型轮廓并输出
                    cv2.drawContours(self.image, [c], -1, (255, 0, 255), 3)
                    spt_lst = os.path.splitext(self.image_path)
                    draw_path = spt_lst[0] + '_draw' + spt_lst[1]
                    cv2.imwrite(draw_path, self.image)
                    all_c.append({'x': x, 'y': y, 'w': w, 'h': h})
        return all_c

    def get_std_points(self, ow=10, oh=10, std_w=(10, 10), std_h=(10, 10)):
        """
        """
        std_points = []
        cnts = self.get_cnts(ow, oh)
        for c in cnts:
            if c["w"] > std_w[0] and c["w"] < std_w[1]:
                std_points.append(c)
        std_points = sorted(std_points, key=lambda x: x["x"] + x["y"])
        return std_points

    def get_ans_points(self, ow=20, oh=10):
        """识别考号和客观题答案
        """
        std_w = (ow,ow*2)
        std_h = (oh,oh*2)
        std_points = []
        cnts = self.get_cnts(ow, oh)
        for c in cnts:
            #if c["w"] > std_w[0] and c["w"] < std_w[1]:
            if c["w"] > std_w[0]:
                std_points.append(c)
        #std_points = sorted(std_points, key=lambda x: x["x"])

        ans_x_list = sorted(std_points, key=lambda x: x["x"])
        ans_y_list = sorted(std_points, key=lambda x: x["y"])
        return ans_x_list, ans_y_list

    def rec_kh(self,kh_y_list,std_x):
        """识别和计算考号
        """
        kh_list = []
        #计算考号区域的总长度
        kh_w_list = [x["w"] for x in kh_y_list]
        kh_w_list.sort()
        if kh_w_list[1:-2]:
            std_w = sum(kh_w_list[1:-2])/len(kh_w_list[1:-2])
            kh_length = std_w * 1.5 * 10
            for kh in kh_y_list:
                kh_x = kh["x"] + kh_length
                dis = kh_x-std_x
                kh = round(dis/(std_w*1.5))
                kh_list.append(str(kh))
        return "".join(kh_list)

    def get_kh_ans(self,ow,oh,std_x_list,std_y_list,qno_list):
        """
        """
        ans_x_list,ans_y_list = self.get_ans_points(ow,oh)
        std_x = std_x_list[0]["x"]
        std_w = std_x_list[0]["w"]

        #识别考号
        kh_x_list = filter(lambda x:x["x"]<std_x-std_w/2,ans_x_list)
        kh_y_list = sorted(kh_x_list,key=lambda x: x["y"])
        kh = self.rec_kh(kh_y_list,std_x-std_w)

        #识别客观题答案
        ans_dct = {}
        #遍历题号
        for i,qno in enumerate(qno_list):
            std_x = std_x_list[i]["x"]
            std_w = std_x_list[i]["w"]
            #遍历识别出的轮廓
            ans_list = []
            for ans in ans_x_list:
                #判断当前题目是否有填涂
                if ans["x"] > std_x - std_w and ans["x"] < std_x + std_w:
                    #计算填图的答案是ABCD
                    for j,stdy in enumerate(std_y_list):
                        std_y = stdy["y"]
                        std_h = stdy["h"]
                        if ans["y"] > std_y - std_h and ans["y"] < std_y + std_h:
                            std_choices = ["A","B","C","D","E","F","G","H"]
                            ans_list.append(std_choices[j])
            ans_dct[qno] = ",".join(ans_list)
        return kh,ans_dct



def main():
    stdans_config = {"std_x_list": [{"h": 19, "w": 38, "x": 968, "y": 318}, {"h": 20, "w": 38, "x": 1027, "y": 318}, {"h": 20, "w": 39, "x": 1086, "y": 318}, {"h": 20, "w": 38, "x": 1145, "y": 319}, {"h": 20, "w": 38, "x": 1204, "y": 319}, {"h": 20, "w": 38, "x": 1263, "y": 319}, {"h": 20, "w": 39, "x": 1322, "y": 319}, {"h": 20, "w": 37, "x": 1382, "y": 319}, {"h": 21, "w": 38, "x": 1440, "y": 319}, {"h": 20, "w": 38, "x": 1499, "y": 319}, {"h": 21, "w": 38, "x": 1558, "y": 319}, {"h": 21, "w": 38, "x": 1617, "y": 319}, {"h": 21, "w": 38, "x": 1676, "y": 319}, {"h": 20, "w": 38, "x": 1735, "y": 320}, {"h": 20, "w": 38, "x": 1794, "y": 320}], "std_y_list": [{"h": 20, "w": 37, "x": 1869, "y": 128}, {"h": 21, "w": 38, "x": 1868, "y": 177}, {"h": 21, "w": 38, "x": 1868, "y": 227}, {"h": 21, "w": 37, "x": 1868, "y": 277}]}
    qno_list = [1,2,3,5,6,9,10,13,14,15,16,17,18,19,20]

    std_x_list = stdans_config.get("std_x_list")
    std_y_list = stdans_config.get("std_y_list")

    img_dect = RecPaper("/tmp/1_0011_crop.jpg")
    img_dect.get_kh_ans(20,10,std_x_list,std_y_list,qno_list) 


if __name__ == "__main__":
    main()