maskrcnn_benchmark/engine/text_inference.py

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import datetime
import logging
import os
import pickle
import subprocess
import time

import cv2
import numpy as np
import torch
from maskrcnn_benchmark.utils.chars import char2num, get_tight_rect, getstr_grid
from PIL import Image, ImageDraw
from tqdm import tqdm

from ..utils.comm import is_main_process, scatter_gather, synchronize
import pdb

# TO DO: format output with dictionnary
def compute_on_dataset(model, data_loader, device, cfg):
    model.eval()
    results_dict = {}
    seg_results = []
    cpu_device = torch.device("cpu")
    total_time = 0
    for _, batch in tqdm(enumerate(data_loader)):
        images, targets, image_paths = batch
        images = images.to(device)
        with torch.no_grad():
            if cfg.MODEL.SEG_ON:
                predictions, proposals, seg_results_dict = model(
                    images
                )
                seg_results.append(
                    [image_paths, proposals, seg_results_dict['rotated_boxes'], seg_results_dict['polygons'], seg_results_dict['preds'], seg_results_dict['scores']]
                ) 
                # if cfg.MODEL.MASK_ON and predictions is not None:
                if predictions is not None:
                    if cfg.MODEL.CHAR_MASK_ON or cfg.SEQUENCE.SEQ_ON:
                        global_predictions = predictions[0]
                        char_predictions = predictions[1]
                        char_mask = char_predictions["char_mask"]
                        boxes = char_predictions["boxes"]
                        seq_words = char_predictions["seq_outputs"]
                        seq_scores = char_predictions["seq_scores"]
                        detailed_seq_scores = char_predictions["detailed_seq_scores"]
                        global_predictions = [o.to(cpu_device) for o in global_predictions]
                        results_dict.update(
                            {
                                image_paths[0]: [
                                    global_predictions[0],
                                    char_mask,
                                    boxes,
                                    seq_words,
                                    seq_scores,
                                    detailed_seq_scores,
                                ]
                            }
                        )
                    else:
                        global_predictions = [o.to(cpu_device) for o in predictions]
                        results_dict.update(
                            {
                                image_paths[0]: [
                                    global_predictions[0],
                                ]
                            }
                        )
            else:
                predictions = model(images)
                if predictions is not None:
                    if not (cfg.MODEL.CHAR_MASK_ON and cfg.SEQUENCE.SEQ_ON):
                        global_predictions = predictions
                        global_predictions = [o.to(cpu_device) for o in global_predictions]
                        results_dict.update(
                            {
                                image_paths[0]: [
                                    global_predictions[0],
                                ]
                            }
                        )
                    else:
                        global_predictions = predictions[0]
                        char_predictions = predictions[1]
                        if cfg.MODEL.CHAR_MASK_ON:
                            char_mask = char_predictions["char_mask"]
                        else:
                            char_mask = None
                        boxes = char_predictions["boxes"]
                        seq_words = char_predictions["seq_outputs"]
                        seq_scores = char_predictions["seq_scores"]
                        detailed_seq_scores = char_predictions["detailed_seq_scores"]
                        global_predictions = [o.to(cpu_device) for o in global_predictions]
                        results_dict.update(
                            {
                                image_paths[0]: [
                                    global_predictions[0],
                                    char_mask,
                                    boxes,
                                    seq_words,
                                    seq_scores,
                                    detailed_seq_scores,
                                ]
                            }
                        )
    return results_dict, seg_results


def polygon2rbox(polygon, image_height, image_width):
    poly = np.array(polygon).reshape((-1, 2))
    rect = cv2.minAreaRect(poly)
    corners = cv2.boxPoints(rect)
    corners = np.array(corners, dtype="int")
    pts = get_tight_rect(corners, 0, 0, image_height, image_width, 1)
    pts = list(map(int, pts))
    return pts


def mask2polygon(mask, box, im_size, threshold=0.5, output_folder=None):
    # mask 32*128
    image_width, image_height = im_size[0], im_size[1]
    box_h = box[3] - box[1]
    box_w = box[2] - box[0]
    cls_polys = (mask * 255).astype(np.uint8)
    poly_map = np.array(Image.fromarray(cls_polys).resize((box_w, box_h)))
    poly_map = poly_map.astype(np.float32) / 255
    poly_map = cv2.GaussianBlur(poly_map, (3, 3), sigmaX=3)
    ret, poly_map = cv2.threshold(poly_map, threshold, 1, cv2.THRESH_BINARY)
    if "total_text" in output_folder or "cute80" in output_folder:
        SE1 = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        poly_map = cv2.erode(poly_map, SE1)
        poly_map = cv2.dilate(poly_map, SE1)
        poly_map = cv2.morphologyEx(poly_map, cv2.MORPH_CLOSE, SE1)
        try:
            _, contours, _ = cv2.findContours(
                (poly_map * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE
            )
        except:
            contours, _ = cv2.findContours(
                (poly_map * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE
            )
        if len(contours) == 0:
            # print(contours)
            # print(len(contours))
            return None
        max_area = 0
        max_cnt = contours[0]
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > max_area:
                max_area = area
                max_cnt = cnt
        # perimeter = cv2.arcLength(max_cnt, True)
        epsilon = 0.01 * cv2.arcLength(max_cnt, True)
        approx = cv2.approxPolyDP(max_cnt, epsilon, True)
        pts = approx.reshape((-1, 2))
        pts[:, 0] = pts[:, 0] + box[0]
        pts[:, 1] = pts[:, 1] + box[1]
        polygon = list(pts.reshape((-1,)))
        polygon = list(map(int, polygon))
        if len(polygon) < 6:
            return None
    else:
        SE1 = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        poly_map = cv2.erode(poly_map, SE1)
        poly_map = cv2.dilate(poly_map, SE1)
        poly_map = cv2.morphologyEx(poly_map, cv2.MORPH_CLOSE, SE1)
        idy, idx = np.where(poly_map == 1)
        xy = np.vstack((idx, idy))
        xy = np.transpose(xy)
        hull = cv2.convexHull(xy, clockwise=True)
        # reverse order of points.
        if hull is None:
            return None
        hull = hull[::-1]
        # find minimum area bounding box.
        rect = cv2.minAreaRect(hull)
        corners = cv2.boxPoints(rect)
        corners = np.array(corners, dtype="int")
        pts = get_tight_rect(corners, box[0], box[1], image_height, image_width, 1)
        polygon = [x * 1.0 for x in pts]
        polygon = list(map(int, polygon))
    return polygon


def _accumulate_predictions_from_multiple_gpus(predictions_per_gpu):
    all_predictions = scatter_gather(predictions_per_gpu)
    if not is_main_process():
        return
    # merge the list of dicts
    predictions = {}
    for p in all_predictions:
        predictions.update(p)
    return predictions


def format_output(out_dir, boxes, img_name):
    with open(
        os.path.join(out_dir, "res_" + img_name.split(".")[0] + ".txt"), "wt"
    ) as res:
        ## char score save dir
        ssur_name = os.path.join(out_dir, "res_" + img_name.split(".")[0])
        for i, box in enumerate(boxes):
            save_name = ssur_name + "_" + str(i) + ".pkl"
            save_dict = {}
            if "total_text" in out_dir or "cute80" in out_dir:
                # np.save(save_name, box[-2])
                save_dict["seg_char_scores"] = box[-3]
                save_dict["seq_char_scores"] = box[-2]
                box = (
                    ",".join([str(x) for x in box[:4]])
                    + ";"
                    + ",".join([str(x) for x in box[4 : 4 + int(box[-1])]])
                    + ";"
                    + ",".join([str(x) for x in box[4 + int(box[-1]) : -3]])
                    + ","
                    + save_name
                )
            else:
                save_dict["seg_char_scores"] = box[-2]
                save_dict["seq_char_scores"] = box[-1]
                np.save(save_name, box[-1])
                box = ",".join([str(x) for x in box[:-2]]) + "," + save_name
            with open(save_name, "wb") as f:
                pickle.dump(save_dict, f, protocol=2)
            res.write(box + "\n")

def format_seg_output(results_dir, rotated_boxes_this_image, polygons_this_image, scores, img_name, ratio):
    height_ratio, width_ratio = ratio
    with open(
        os.path.join(results_dir, "res_" + img_name.split(".")[0] + ".txt"), "wt"
    ) as res:
        if "total_text" in results_dir or "cute80" in results_dir:
            for i, box in enumerate(polygons_this_image):
                box = box[0]
                box[0::2] = box[0::2] * width_ratio
                box[1::2] = box[1::2] * height_ratio
                save_dict = {}
                # result = ",".join([str(int(x[0])) + ',' +str(int(x[1])) for x in box])
                result = ",".join([str(int(x)) for x in box])
                score = scores[i].item()
                res.write(result + ',' + str(score) + "\n")
        else:
            for i, box in enumerate(rotated_boxes_this_image):
                box[0::2] = box[0::2] * width_ratio
                box[1::2] = box[1::2] * height_ratio
                save_dict = {}
                result = ",".join([str(int(x[0])) + ',' +str(int(x[1])) for x in box])
                score = scores[i].item()
                res.write(result + ',' + str(score) + "\n")



def process_char_mask(char_masks, boxes, threshold=192):
    texts, rec_scores, rec_char_scores, char_polygons = [], [], [], []
    for index in range(char_masks.shape[0]):
        box = list(boxes[index])
        box = list(map(int, box))
        text, rec_score, rec_char_score, char_polygon = getstr_grid(
            char_masks[index, :, :, :].copy(), box, threshold=threshold
        )
        texts.append(text)
        rec_scores.append(rec_score)
        rec_char_scores.append(rec_char_score)
        char_polygons.append(char_polygon)
        # segmss.append(segms)
    return texts, rec_scores, rec_char_scores, char_polygons


def creat_color_map(n_class, width):
    splits = int(np.ceil(np.power((n_class * 1.0), 1.0 / 3)))
    maps = []
    for i in range(splits):
        r = int(i * width * 1.0 / (splits - 1))
        for j in range(splits):
            g = int(j * width * 1.0 / (splits - 1))
            for k in range(splits - 1):
                b = int(k * width * 1.0 / (splits - 1))
                maps.append((r, g, b, 200))
    return maps


def visualization(image, polygons, resize_ratio, colors, char_polygons=None, words=None):
    draw = ImageDraw.Draw(image, "RGBA")
    for polygon in polygons:
        # draw.polygon(polygon, fill=None, outline=(0, 255, 0, 255))
        # print(polygon)
        polygon.append(polygon[0])
        polygon.append(polygon[1])
        # print(polygon)
        color = '#33FF33'
        draw.line(polygon, fill=color, width=5)
    # if char_polygons is not None:
    #     for i, char_polygon in enumerate(char_polygons):
    #         for j, polygon in enumerate(char_polygon):
    #             polygon = [int(x * resize_ratio) for x in polygon]
    #             char = words[i][j]
    #             color = colors[char2num(char)]
    #             draw.polygon(polygon, fill=color, outline=color)


def vis_seg_map(image_path, seg_map, rotated_boxes, polygons_this_image, proposals, vis_dir):
    img_name = image_path.split("/")[-1]
    image = cv2.imread(image_path)
    height, width, _ = image.shape
    seg_map = seg_map.data.cpu().numpy()
    img = Image.fromarray(image).convert("RGB")
    # height_ratio = height / seg_map.shape[1]
    # width_ratio = width / seg_map.shape[2]
    # print('seg_map.shape:', seg_map.shape)
    # print('image.shape:', image.shape)
    seg_image = (
        Image.fromarray((seg_map[0, :proposals.size[1], :proposals.size[0]] * 255).astype(np.uint8))
        .convert("RGB")
        .resize((width, height))
    )
    visu_image = Image.blend(seg_image, img, 0.5)
    img_draw = ImageDraw.Draw(visu_image)
    if "total_text" in vis_dir or "cute80" in vis_dir:
        for box in polygons_this_image:
            # box[:, 0] = box[:, 0]
            # box[:, 1] = box[:, 1]
            tuple_box = [tuple(x) for x in box[0].reshape(-1, 2).tolist()]
            tuple_box.append(tuple_box[0])
            img_draw.line(tuple_box, fill=(0, 255, 0), width=5)
    else:
        for box in rotated_boxes:
            # box[:, 0] = box[:, 0]
            # box[:, 1] = box[:, 1]
            tuple_box = [tuple(x) for x in box.tolist()]
            tuple_box.append(tuple_box[0])
            img_draw.line(tuple_box, fill=(0, 255, 0), width=5)
    visu_image.save(vis_dir + "/seg_" + img_name)


def prepare_results_for_evaluation(
    predictions, output_folder, model_name, seg_predictions=None, vis=False, cfg=None
):
    results_dir = os.path.join(output_folder, model_name + "_results")
    if not os.path.isdir(results_dir):
        os.mkdir(results_dir)
    seg_results_dir = os.path.join(output_folder, model_name + "_seg_results")
    if not os.path.isdir(seg_results_dir):
        os.mkdir(seg_results_dir)
    if vis:
        visu_dir = os.path.join(output_folder, model_name + "_visu")
        if not os.path.isdir(visu_dir):
            os.mkdir(visu_dir)
        seg_visu_dir = os.path.join(output_folder, model_name + "_seg_visu")
        if not os.path.isdir(seg_visu_dir):
            os.mkdir(seg_visu_dir)
    if len(seg_predictions) > 0:
        for seg_prediction in seg_predictions:
            image_paths, proposals, rotated_boxes, polygons, seg_maps, seg_scores = (
                seg_prediction[0],
                seg_prediction[1],
                seg_prediction[2],
                seg_prediction[3],
                seg_prediction[4],
                seg_prediction[5],
            )
            for batch_id in range(len(image_paths)):
                image_path = image_paths[batch_id]
                im_name = image_path.split("/")[-1]
                image = cv2.imread(image_path)
                height, width, _ = image.shape
                rotated_boxes_this_image = rotated_boxes[batch_id]
                polygons_this_image = polygons[batch_id]
                proposals_this_image = proposals[batch_id]
                seg_map = seg_maps[batch_id]
                seg_score = seg_scores[batch_id]
                height, width, _ = image.shape
                height_ratio = height / proposals_this_image.size[1]
                width_ratio = width / proposals_this_image.size[0]
                format_seg_output(seg_results_dir, rotated_boxes_this_image, polygons_this_image, seg_score, im_name, (height_ratio, width_ratio))
                if vis:
                    vis_seg_map(image_path, seg_map, rotated_boxes_this_image, polygons_this_image, proposals_this_image, seg_visu_dir)
    if (not cfg.MODEL.TRAIN_DETECTION_ONLY):
        for image_path, prediction in predictions.items():
            im_name = image_path.split("/")[-1]
            if cfg.MODEL.CHAR_MASK_ON or cfg.SEQUENCE.SEQ_ON:
                global_prediction, char_mask, boxes_char, seq_words, seq_scores, detailed_seq_scores = (
                    prediction[0],
                    prediction[1],
                    prediction[2],
                    prediction[3],
                    prediction[4],
                    prediction[5],
                )
                if char_mask is not None:
                    words, rec_scores, rec_char_scoress, char_polygons = process_char_mask(
                        char_mask, boxes_char
                    )
            else:
                global_prediction = prediction[0]
            test_image_width, test_image_height = global_prediction.size
            img = Image.open(image_path)
            width, height = img.size
            resize_ratio = float(height) / test_image_height
            global_prediction = global_prediction.resize((width, height))
            boxes = global_prediction.bbox.tolist()
            if cfg.MODEL.ROI_BOX_HEAD.INFERENCE_USE_BOX:
                scores = global_prediction.get_field("scores").tolist()
            if not cfg.MODEL.SEG.USE_SEG_POLY:
                masks = global_prediction.get_field("mask").cpu().numpy()
            else:
                masks = global_prediction.get_field("masks").get_polygons()
            result_logs = []
            polygons = []
            for k, box in enumerate(boxes):
                if box[2] - box[0] < 1 or box[3] - box[1] < 1:
                    continue
                box = list(map(int, box))
                if not cfg.MODEL.SEG.USE_SEG_POLY:
                    mask = masks[k, 0, :, :]
                    polygon = mask2polygon(
                        mask, box, img.size, threshold=0.5, output_folder=output_folder
                    )
                else:
                    polygon = list(masks[k].get_polygons()[0].cpu().numpy())
                    if not ("total_text" in output_folder or "cute80" in output_folder):
                        polygon = polygon2rbox(polygon, height, width)
                if polygon is None:
                    polygon = [
                        box[0],
                        box[1],
                        box[2],
                        box[1],
                        box[2],
                        box[3],
                        box[0],
                        box[3],
                    ]
                    continue
                polygons.append(polygon)
                if cfg.MODEL.ROI_BOX_HEAD.INFERENCE_USE_BOX:
                    score = scores[k]
                else:
                    score = 1.0
                if cfg.MODEL.CHAR_MASK_ON or cfg.SEQUENCE.SEQ_ON:
                    if char_mask is None:
                        word = 'aaa'
                        rec_score = 1.0
                        char_score = None
                    else:
                        word = words[k]
                        rec_score = rec_scores[k]
                        char_score = rec_char_scoress[k]
                    seq_word = seq_words[k]
                    seq_char_scores = seq_scores[k]
                    seq_score = sum(seq_char_scores) / float(len(seq_char_scores))
                    detailed_seq_score = detailed_seq_scores[k]
                    detailed_seq_score = np.squeeze(np.array(detailed_seq_score), axis=1)
                else:
                    word = 'aaa'
                    rec_score = 1.0
                    char_score = [1.0, 1.0, 1.0]
                    seq_word = 'aaa'
                    seq_char_scores = [1.0, 1.0, 1.0]
                    seq_score = 1.0
                    detailed_seq_score = None
                if "total_text" in output_folder or "cute80" in output_folder:
                    result_log = (
                        [int(x * 1.0) for x in box[:4]]
                        + polygon
                        + [word]
                        + [seq_word]
                        + [score]
                        + [rec_score]
                        + [seq_score]
                        + [char_score]
                        + [detailed_seq_score]
                        + [len(polygon)]
                    )
                else:
                    result_log = (
                        [int(x * 1.0) for x in box[:4]]
                        + polygon
                        + [word]
                        + [seq_word]
                        + [score]
                        + [rec_score]
                        + [seq_score]
                        + [char_score]
                        + [detailed_seq_score]
                    )
                result_logs.append(result_log)
            if vis:
                colors = creat_color_map(37, 255)
                if cfg.MODEL.CHAR_MASK_ON:
                    visualization(img, polygons, resize_ratio, colors, char_polygons, words)
                else:
                    visualization(img, polygons, resize_ratio, colors)
                img.save(os.path.join(visu_dir, im_name))
            format_output(results_dir, result_logs, im_name)


def inference(
    model,
    data_loader,
    iou_types=("bbox",),
    box_only=False,
    device="cuda",
    expected_results=(),
    expected_results_sigma_tol=4,
    output_folder=None,
    model_name=None,
    cfg=None,
):

    # convert to a torch.device for efficiency
    model_name = model_name.split(".")[0] + "_" + str(cfg.INPUT.MIN_SIZE_TEST)
    predictions_path = os.path.join(output_folder, model_name + "_predictions.pth")
    seg_predictions_path = os.path.join(
        output_folder, model_name + "_seg_predictions.pth"
    )
    # if os.path.isfile(predictions_path) and os.path.isfile(seg_predictions_path):
    if False:
        predictions = torch.load(predictions_path)
        seg_predictions = torch.load(seg_predictions_path)
    else:
        device = torch.device(device)
        num_devices = (
            torch.distributed.get_world_size()
            if torch.distributed.is_initialized()
            else 1
        )
        logger = logging.getLogger("maskrcnn_benchmark.inference")
        dataset = data_loader.dataset
        logger.info("Start evaluation on {} images".format(len(dataset)))
        start_time = time.time()
        predictions, seg_predictions = compute_on_dataset(
            model, data_loader, device, cfg
        )
        # wait for all processes to complete before measuring the time
        synchronize()
        total_time = time.time() - start_time
        total_time_str = str(datetime.timedelta(seconds=total_time))
        logger.info(
            "Total inference time: {} ({} s / img per device, on {} devices)".format(
                total_time_str, total_time * num_devices / len(dataset), num_devices
            )
        )

        # predictions = _accumulate_predictions_from_multiple_gpus(predictions)
        # if not is_main_process():
        # 	return

        if output_folder:
            torch.save(predictions, predictions_path)
            torch.save(seg_predictions, seg_predictions_path)

    prepare_results_for_evaluation(
        predictions,
        output_folder,
        model_name,
        seg_predictions=seg_predictions,
        vis=cfg.TEST.VIS,
        cfg=cfg
    )