exam-segment-django/segment/sheet_resolve/analysis/choice/analysis_choice.py

# @Author  : lightXu
# @File    : analysis_choice.py
import time

import numpy as np

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


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

    size = im.shape

    # Detect all object classes and regress object bounds
    timer = Timer()
    timer.tic()
    im, radio = utils.img_resize(analysis_type, im)
    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 = []
    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]) * radio[0]) + xmin_bias
                ymin = int(int(bbox[1]) * radio[1]) + ymin_bias
                xmax = int(int(bbox[2]) * radio[0]) + xmax_bias
                ymax = int(int(bbox[3]) * radio[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)
                xavg = int(xmin + (xmax - xmin) / 2)
                yavg = int(ymin + (ymax - ymin) / 2)

                bbox_dict = {"xmin": xmin, "ymin": ymin, "xmax": xmax, "ymax": ymax, 'x_center': xavg, 'y_center': yavg}
                class_dict = {"class_name": cls, "bounding_box": bbox_dict, "score": score, 'x_center': xavg,
                              'y_center': yavg}
                content_list.append(class_dict)

                analysis_cls_list.append(cls)

    return content_list, analysis_cls_list


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

    content, cls = analysis_single_image_with_regions(analysis_type, classes,
                                                      sess, net,
                                                      im, conf_thresh, mns_thresh,
                                                      coordinate_bias_dict)

    img_dict = {"img_name": 'choice_m',
                'analysis_type': analysis_type,
                "regions": content,
                }

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

    return img_dict