exam-segment-django/segment/sheet_resolve/analysis/sheet/analysis_sheet.py

# @Author  : lightXu
# @File    : analysis_sheet.py
import time
import os
import traceback

import numpy as np
import cv2

from segment.sheet_resolve.lib.model.test import im_detect
from segment.sheet_resolve.lib.model.nms_wrapper import nms
from segment.sheet_resolve.lib.utils.timer import Timer
from segment.sheet_resolve.tools import utils
from segment.sheet_resolve.analysis.solve.optional_solve import find_contours, resolve_optional_choice


def analysis_single_image_with_regions(analysis_type, classes,
                                       sess, net,
                                       im_raw, conf_thresh, mns_thresh,
                                       coordinate_bias_dict):
    """Detect object classes in an image using pre-computed object proposals."""

    size = im_raw.shape

    # Detect all object classes and regress object bounds
    timer = Timer()
    timer.tic()
    if analysis_type in ['unknown_subject', 'math', 'math_zxhx', 'english', 'chinese',
                         'physics', 'chemistry', 'biology', 'politics', 'history',
                         'geography', 'science_comprehensive', 'arts_comprehensive'
                         ]:
        analysis_type = 'sheet'
    im, ratio = utils.img_resize(analysis_type, im_raw)
    scores, boxes = im_detect(analysis_type, sess, net, im)
    timer.toc()
    print('Detection took {:.3f}s for {:d} object proposals'.format(timer.total_time, boxes.shape[0]))

    content_list = []
    analysis_cls_list = []
    qr_code_info = 'Nan'

    for cls_ind, cls in enumerate(classes[1:]):  # classes
        cls_ind += 1  # because we skipped background
        cls_boxes = boxes[:, 4 * cls_ind:4 * (cls_ind + 1)]
        cls_scores = scores[:, cls_ind]
        dets = np.hstack((cls_boxes,
                          cls_scores[:, np.newaxis])).astype(np.float32)
        keep = nms(dets, mns_thresh)
        dets = dets[keep, :]
        # vis_detections(im, cls, dets, ax, thresh=conf_thresh)
        inds = np.where(dets[:, -1] >= conf_thresh)[0]
        if len(inds) > 0:
            if cls in list(coordinate_bias_dict.keys()):
                xmin_bias = coordinate_bias_dict[cls]['xmin_bias']
                ymin_bias = coordinate_bias_dict[cls]['ymin_bias']
                xmax_bias = coordinate_bias_dict[cls]['xmax_bias']
                ymax_bias = coordinate_bias_dict[cls]['ymax_bias']
            else:
                xmin_bias = 0
                ymin_bias = 0
                xmax_bias = 0
                ymax_bias = 0
            for i in inds:
                bbox = dets[i, :4]
                score = '{:.4f}'.format(dets[i, -1])

                xmin = int(int(bbox[0]) * ratio[0]) + xmin_bias
                ymin = int(int(bbox[1]) * ratio[1]) + ymin_bias
                xmax = int(int(bbox[2]) * ratio[0]) + xmax_bias
                ymax = int(int(bbox[3]) * ratio[1]) + ymax_bias

                xmin = (xmin if (xmin > 0) else 1)
                ymin = (ymin if (ymin > 0) else 1)
                xmax = (xmax if (xmax < size[1]) else size[1] - 1)
                ymax = (ymax if (ymax < size[0]) else size[0] - 1)

                if cls in ['solve0', ]:
                    cls = 'solve'

                bbox_dict = {"xmin": xmin, "ymin": ymin, "xmax": xmax, "ymax": ymax}
                # class_dict = {"class_name": cls, "bounding_box": bbox_dict, "score": score}
                class_dict = {"class_name": cls, "bounding_box": bbox_dict}

                # if cls == 'qr_code':
                #     qr_img = utils.crop_region(im_raw, bbox_dict)
                #     qr_path = r'./qr_code.jpg'
                #     cv2.imwrite(qr_path, qr_img)
                #     qr_code_info = utils.check_qr_code_with_region_img(qr_path)
                #     os.remove(qr_path)

                content_list.append(class_dict)

    return content_list, analysis_cls_list, qr_code_info


def get_single_image_sheet_regions(analysis_type, img_path, img, classes,
                                   sess, net, conf_thresh, mns_thresh,
                                   coordinate_bias_dict):
    start_time = time.time()

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('analysis for JPG {}'.format(img_path))

    content, cls, qr_code_info = \
        analysis_single_image_with_regions(analysis_type, classes, sess, net,
                                           img, conf_thresh, mns_thresh,
                                           coordinate_bias_dict)

    img_dict = {"img_name": img_path,
                # 'qr_code': qr_code_info,
                'subject': analysis_type,
                "regions": content,
                }

    end_time = time.time()
    print(end_time - start_time)

    return img_dict


def question_number_format(init_number, crt_numbers, sheet_dict):
    for region in sheet_dict['regions']:
        numbers = region.get("number")
        if numbers and isinstance(numbers, int):
            if numbers <= 0 or numbers in crt_numbers:
                numbers = init_number
                crt_numbers.append(numbers)
                init_number += 1
            region.update({"number": numbers})
            crt_numbers.append(numbers)
        if numbers and isinstance(numbers, list):
            for i, num in enumerate(numbers):
                if num <= 0 or num in crt_numbers:
                    numbers[i] = init_number
                    crt_numbers.append(init_number)
                    init_number += 1

            region.update({"number": numbers})
            crt_numbers.extend(numbers)

    return sheet_dict, init_number, crt_numbers


def box_region_format(sheet_dict, image, subject, shrink=True):
    include_class = ['anchor_point',
                     'bar_code',
                     'choice_m',
                     'cloze',
                     'cloze_s',
                     'exam_number_col_row',
                     'optional_choice',
                     'optional_solve',
                     # 'qr_code',
                     'solve',
                     'optional_solve',
                     'composition',
                     # 'correction'
                     ]

    sheet_regions = sheet_dict['regions']
    optional_solve_tmp = []
    default_points_dict = {'choice_m': 5, "cloze": 5, 'solve': 12, 'cloze_s': 5, "composition": 60}
    if subject == "english":
        default_points_dict = {'choice_m': 2, "cloze": 2, 'solve': 2, 'cloze_s': 2, "composition": 25}

    for i in range(len(sheet_regions) - 1, -1, -1):
        if subject == "math":
            if sheet_regions[i]['class_name'] == 'cloze':
                sheet_regions[i]['class_name'] = 'cloze_big'  # math exclude cloze big
            if sheet_regions[i]['class_name'] == 'cloze_s':
                sheet_regions[i]['class_name'] = 'cloze'  # math exclude cloze big
        if subject == "english":
            if sheet_regions[i]['class_name'] == 'solve':
                sheet_regions[i]['class_name'] = 'cloze'
            if sheet_regions[i]['class_name'] == 'correction':
                sheet_regions[i]['class_name'] = 'solve'

    for i in range(len(sheet_regions) - 1, -1, -1):
        if sheet_regions[i]['class_name'] in ['solve0']:
            sheet_regions[i]['class_name'] = 'solve'
        if sheet_regions[i]['class_name'] in ['composition0']:
            sheet_regions[i]['class_name'] = 'composition'

        if sheet_regions[i]['class_name'] == 'select_s':
            sheet_regions[i]['class_name'] = 'optional_choice'
            optional_solve_tmp.append(sheet_regions[i])
            sheet_regions.pop(i)

        if shrink:
            if sheet_regions[i]['class_name'] not in include_class:
                sheet_regions.pop(i)

    for ele in sheet_regions:
        if ele['class_name'] == 'solve':
            solve_box = (ele['bounding_box']['xmin'], ele['bounding_box']['ymin'],
                         ele['bounding_box']['xmax'], ele['bounding_box']['ymax'])
            for optional_solve in optional_solve_tmp:
                optional_solve_box = (optional_solve['bounding_box']['xmin'], optional_solve['bounding_box']['ymin'],
                                      optional_solve['bounding_box']['xmax'], optional_solve['bounding_box']['ymax'])
                if utils.decide_coordinate_contains(optional_solve_box, solve_box):
                    ele['class_name'] = 'optional_solve'
                    break
                else:
                    continue

        if ele['class_name'] == "composition":
            if isinstance(ele['default_points'], list):
                for i, dp in enumerate(ele['default_points']):
                    if dp != default_points_dict[ele['class_name']]:
                        ele['default_points'][i] = default_points_dict[ele['class_name']]

            if isinstance(ele['default_points'], int):
                if ele['default_points'] != default_points_dict[ele['class_name']]:
                    ele['default_points'] = default_points_dict[ele['class_name']]

        if ele['class_name'] in ["choice_m", "cloze", "cloze_s", "solve"]:
            if isinstance(ele['default_points'], list):
                for i, dp in enumerate(ele['default_points']):
                    if dp == -1:
                        ele['default_points'][i] = default_points_dict[ele['class_name']]

            if isinstance(ele['default_points'], int):
                if ele['default_points'] == -1:
                    ele['default_points'] = default_points_dict[ele['class_name']]

    for ele in optional_solve_tmp:  # 选做题
        bbox = ele['bounding_box']
        box_region = utils.crop_region(image, bbox)
        left = bbox['xmin']
        top = bbox['ymin']
        right = bbox['xmax']
        bottom = bbox['ymax']

        if (right - left) >= (bottom-top):
            direction = 180
        else:
            direction = 90

        # res = find_contours(left, top, box_region)
        try:
            res = resolve_optional_choice(left, top, direction, box_region)
        except Exception as e:
            res = {'rows': 1, 'cols': 2,
                   'option': 'A, B',
                   'single_width': (right - left) // 3,
                   'single_height': bottom - top,
                   'bounding_box': {'xmin': left,
                                    'ymin': top,
                                    'xmax': right,
                                    'ymax': bottom}}
        res['class_name'] = 'optional_choice'

        sheet_regions.append(res)

    # iou
    sheet_tmp = sheet_regions.copy()
    remove_index = []
    for i, region in enumerate(sheet_tmp):
        if i not in remove_index:
            box = region['bounding_box']
            for j, region_in in enumerate(sheet_tmp):
                box_in = region_in['bounding_box']
                iou = utils.cal_iou(box, box_in)
                if iou[0] > 0.75 and i != j:
                    sheet_regions.remove(region)
                    remove_index.append(j)
                    break

    sheet_dict.update({'regions': sheet_regions})
    return sheet_dict