answer-sheet-scan/sheet.py at master · inuyasha2012/answer-sheet-scan · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
# coding=utf-8
import cPickle
import cv2
from imutils import contours
from e import ContourCountError, ContourPerimeterSizeError, PolyNodeCountError
from settings import ANS_IMG_THRESHOLD, CNT_PERIMETER_THRESHOLD, CHOICE_IMG_THRESHOLD, ANS_IMG_DILATE_ITERATIONS, \
    ANS_IMG_ERODE_ITERATIONS, CHOICE_IMG_DILATE_ITERATIONS, CHOICE_IMG_ERODE_ITERATIONS, CHOICE_MAX_AREA, \
    CHOICE_CNT_COUNT, ANS_IMG_KERNEL, CHOICE_IMG_KERNEL, CHOICE_MIN_AREA
from utils import detect_cnt_again, get_init_process_img, get_max_area_cnt, get_ans, \
    get_left_right, get_top_bottom, sort_by_row, sort_by_col, insert_null_2_rows, get_vertical_projective, \
    get_h_projective


def get_answer_from_sheet(base_img):

    # cv2.imshow('temp', base_img)
    # cv2.waitKey(0)

    # 灰度化然后进行边缘检测、二值化等等一系列处理
    img = cv2.cvtColor(base_img, cv2.COLOR_BGR2GRAY)
    img = get_init_process_img(img)

    # cv2.imshow('temp', img)
    # cv2.waitKey(0)

    # 获取最大面积轮廓并和图片大小作比较，看轮廓周长大小判断是否是答题卡的轮廓
    cnt = get_max_area_cnt(img)
    cnt_perimeter = cv2.arcLength(cnt, True)
    base_img_perimeter = (base_img.shape[0] + base_img.shape[1]) * 2
    if not cnt_perimeter > CNT_PERIMETER_THRESHOLD * base_img_perimeter:
        raise ContourPerimeterSizeError

    # cv2.drawContours(base_img, [cnt], 0, (0, 255, 0), 1)
    # cv2.imshow('temp', base_img)
    # cv2.waitKey(0)

    # 计算多边形的顶点，并看是否是四个顶点
    poly_node_list = cv2.approxPolyDP(cnt, cv2.arcLength(cnt, True) * 0.1, True)
    if not poly_node_list.shape[0] == 4:
        raise PolyNodeCountError

    # 根据计算的多边形顶点继续处理图片，主要是是纠偏
    processed_img = detect_cnt_again(poly_node_list, base_img)

    # cv2.imshow('temp', processed_img)
    # cv2.waitKey(0)

    # 调整图片的亮度
    processed_img = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY)
    processed_img = cv2.GaussianBlur( processed_img, (9, 9), 0)
    # cv2.imshow('temp', processed_img)
    # cv2.waitKey(0)

    # 通过二值化和膨胀腐蚀获得填涂区域
    # ret, ans_img = cv2.threshold(processed_img, ANS_IMG_THRESHOLD[0], ANS_IMG_THRESHOLD[1], cv2.THRESH_BINARY_INV)
    ans_img = cv2.adaptiveThreshold(processed_img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 41, 35)
    ans_img = cv2.dilate(ans_img, ANS_IMG_KERNEL, iterations=ANS_IMG_DILATE_ITERATIONS)
    ans_img = cv2.erode(ans_img, ANS_IMG_KERNEL, iterations=ANS_IMG_ERODE_ITERATIONS)
    ret, ans_img = cv2.threshold(ans_img, ANS_IMG_THRESHOLD[0], ANS_IMG_THRESHOLD[1], cv2.THRESH_BINARY_INV)

    # cv2.imshow('temp', ans_img)
    # cv2.waitKey(0)

    # 通过二值化和膨胀腐蚀获得选项框区域
    choice_img = cv2.adaptiveThreshold(processed_img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 51, 10)
    # ret, choice_img = cv2.threshold(processed_img, CHOICE_IMG_THRESHOLD[0], CHOICE_IMG_THRESHOLD[1],
    #                                 cv2.THRESH_BINARY_INV)
    # cv2.imshow('temp', choice_img)
    # cv2.waitKey(0)
    # get_vertical_projective(choice_img)
    # get_h_projective(choice_img)

    for i in range(1):
        choice_img = cv2.dilate(choice_img, CHOICE_IMG_KERNEL, iterations=CHOICE_IMG_DILATE_ITERATIONS)
        choice_img = cv2.erode(choice_img, CHOICE_IMG_KERNEL, iterations=CHOICE_IMG_ERODE_ITERATIONS)

    # get_vertical_projective(choice_img)
    # get_h_projective(choice_img)
    # cv2.imshow('temp', choice_img)
    # cv2.waitKey(0)

    # 查找选项框以及前面题号的轮廓
    cnts, h = cv2.findContours(choice_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    question_cnts = []

    temp1_ans_img = ans_img.copy()
    for i, c in enumerate(cnts):
        # 如果面积小于某值，则认为这个轮廓是选项框或题号
        if CHOICE_MIN_AREA < cv2.contourArea(c) < CHOICE_MAX_AREA:
            cv2.drawContours(temp1_ans_img, cnts, i, (0, 0, 0), 1)
            question_cnts.append(c)

    # cv2.imshow('temp', temp1_ans_img)
    # cv2.waitKey(0)

    question_cnts = get_left_right(question_cnts)
    question_cnts = get_top_bottom(question_cnts)

    temp2_ans_img = ans_img.copy()
    for i in range(len(question_cnts)):
        cv2.drawContours(temp2_ans_img, question_cnts, i, (0, 0, 0), 1)

    # cv2.imshow('temp', temp2_ans_img)
    # cv2.waitKey(0)

    # 如果轮廓小于特定值，重新扫描
    # TODO 运用统计分析排除垃圾轮廓
    # if len(question_cnts) != CHOICE_CNT_COUNT:
    #     raise ContourCountError

    # cv2.imshow('temp', temp_ans_img)
    # cv2.waitKey(0)

    # 对轮廓之上而下的排序
    question_cnts, cnts_pos = contours.sort_contours(question_cnts, method="top-to-bottom")
    rows = sort_by_row(list(cnts_pos))
    cols = sort_by_col(list(cnts_pos))

    # cv2.imshow('temp', temp2_ans_img)
    # cv2.waitKey(0)

    insert_null_2_rows(cols, rows)
    # 获得答案
    rows, res = get_ans(ans_img, rows)
    if not res[0]:
        print res[1]
        cv2.imshow('temp1', temp2_ans_img)
        cv2.waitKey(0)
        print 'end'
    else:
        print res