tools/demo.py

import os
import cv2
import torch
from torchvision import transforms as T

from maskrcnn_benchmark.modeling.detector import build_detection_model
from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer
from maskrcnn_benchmark.structures.image_list import to_image_list
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.utils.chars import getstr_grid, get_tight_rect

from PIL import Image
import numpy as np
import argparse

class TextDemo(object):
    def __init__(
        self,
        cfg,
        confidence_threshold=0.7,
        min_image_size=224,
        output_polygon=True
    ):
        self.cfg = cfg.clone()
        self.model = build_detection_model(cfg)
        self.model.eval()
        self.device = torch.device(cfg.MODEL.DEVICE)
        self.model.to(self.device)
        self.min_image_size = min_image_size

        checkpointer = DetectronCheckpointer(cfg, self.model)
        _ = checkpointer.load(cfg.MODEL.WEIGHT)

        self.transforms = self.build_transform()
        self.cpu_device = torch.device("cpu")
        self.confidence_threshold = confidence_threshold
        self.output_polygon = output_polygon

    def build_transform(self):
        """
        Creates a basic transformation that was used to train the models
        """
        cfg = self.cfg
        # we are loading images with OpenCV, so we don't need to convert them
        # to BGR, they are already! So all we need to do is to normalize
        # by 255 if we want to convert to BGR255 format, or flip the channels
        # if we want it to be in RGB in [0-1] range.
        if cfg.INPUT.TO_BGR255:
            to_bgr_transform = T.Lambda(lambda x: x * 255)
        else:
            to_bgr_transform = T.Lambda(lambda x: x[[2, 1, 0]])

        normalize_transform = T.Normalize(
            mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD
        )

        transform = T.Compose(
            [
                T.ToPILImage(),
                T.Resize(self.min_image_size),
                T.ToTensor(),
                to_bgr_transform,
                normalize_transform,
            ]
        )
        return transform

    def run_on_opencv_image(self, image):
        """
        Arguments:
            image (np.ndarray): an image as returned by OpenCV
        Returns:
            result_polygons (list): detection results
            result_words (list): recognition results
        """
        result_polygons, result_words = self.compute_prediction(image)
        return result_polygons, result_words

    def compute_prediction(self, original_image):
        # apply pre-processing to image
        image = self.transforms(original_image)
        # convert to an ImageList, padded so that it is divisible by
        # cfg.DATALOADER.SIZE_DIVISIBILITY
        image_list = to_image_list(image, self.cfg.DATALOADER.SIZE_DIVISIBILITY)
        image_list = image_list.to(self.device)
        # compute predictions
        with torch.no_grad():
            predictions, _, _ = self.model(image_list)
        global_predictions = predictions[0]
        char_predictions = predictions[1]
        char_mask = char_predictions['char_mask']
        char_boxes = char_predictions['boxes']
        words, rec_scores = self.process_char_mask(char_mask, char_boxes)
        seq_words = char_predictions['seq_outputs']
        seq_scores = char_predictions['seq_scores']

        global_predictions = [o.to(self.cpu_device) for o in global_predictions]

        # always single image is passed at a time
        global_prediction = global_predictions[0]

        # reshape prediction (a BoxList) into the original image size
        height, width = original_image.shape[:-1]
        global_prediction = global_prediction.resize((width, height))
        boxes = global_prediction.bbox.tolist()
        scores = global_prediction.get_field("scores").tolist()
        masks = global_prediction.get_field("mask").cpu().numpy()

        result_polygons = []
        result_words = []
        for k, box in enumerate(boxes):
            score = scores[k]
            if score < self.confidence_threshold:
                continue
            box = list(map(int, box))
            mask = masks[k,0,:,:]
            polygon = self.mask2polygon(mask, box, original_image.shape, threshold=0.5, output_polygon=self.output_polygon)
            if polygon is None:
                polygon = [box[0], box[1], box[2], box[1], box[2], box[3], box[0], box[3]]
            result_polygons.append(polygon)
            word = words[k]
            rec_score = rec_scores[k]
            seq_word = seq_words[k]
            seq_char_scores = seq_scores[k]
            seq_score = sum(seq_char_scores) / float(len(seq_char_scores))
            if seq_score > rec_score:
                result_words.append(seq_word)
            else:
                result_words.append(word)
        return result_polygons, result_words

    def process_char_mask(self, char_masks, boxes, threshold=192):
        texts, rec_scores = [], []
        for index in range(char_masks.shape[0]):
            box = list(boxes[index])
            box = list(map(int, box))
            text, rec_score, _, _ = getstr_grid(char_masks[index,:,:,:].copy(), box, threshold=threshold)
            texts.append(text)
            rec_scores.append(rec_score)
        return texts, rec_scores

    def mask2polygon(self, mask, box, im_size, threshold=0.5, output_polygon=True):
        # mask 32*128
        image_width, image_height = im_size[1], im_size[0]
        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,0.5,1,cv2.THRESH_BINARY)
        if output_polygon:
            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 visualization(self, image, polygons, words):
        for polygon, word in zip(polygons, words):
            pts = np.array(polygon, np.int32)
            pts = pts.reshape((-1,1,2))
            xmin = min(pts[:,0,0])
            ymin = min(pts[:,0,1])
            cv2.polylines(image,[pts],True,(0,0,255))
            cv2.putText(image, word, (xmin, ymin), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255), 2)


def main(args):
    # update the config options with the config file
    cfg.merge_from_file(args.config_file)
    # manual override some options
    # cfg.merge_from_list(["MODEL.DEVICE", "cpu"])

    text_demo = TextDemo(
        cfg,
        min_image_size=800,
        confidence_threshold=0.7,
        output_polygon=True
    )
    # load image and then run prediction
    
    image = cv2.imread(args.image_path)
    result_polygons, result_words = text_demo.run_on_opencv_image(image)
    text_demo.visualization(image, result_polygons, result_words)
    cv2.imwrite(args.visu_path, image)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='parameters for demo')
    parser.add_argument("--config-file", type=str, default='configs/mixtrain/seg_rec_poly_fuse_feature.yaml')
    parser.add_argument("--image_path", type=str, default='./demo_images/demo.jpg')
    parser.add_argument("--visu_path", type=str, default='./demo_images/demo_results.jpg')
    args = parser.parse_args()
    main(args)