tools/convert_dataset.py

import os
import numpy as np
import cv2
from shapely.geometry import box, Polygon 
from shapely import affinity
import math


def _rect2quad(boxes):
    x_min, y_min, x_max, y_max = boxes[:, 0].reshape((-1, 1)), boxes[:, 1].reshape((-1, 1)), boxes[:, 2].reshape((-1, 1)), boxes[:, 3].reshape((-1, 1))
    return np.hstack((x_min, y_min, x_max, y_min, x_max, y_max, x_min, y_max))

def _quad2rect(boxes):
    ## only support rectangle
    return np.hstack((boxes[:, 0].reshape((-1, 1)), boxes[:, 1].reshape((-1, 1)), boxes[:, 4].reshape((-1, 1)), boxes[:, 5].reshape((-1, 1))))

def _quad2minrect(boxes):
    ## trans a quad(N*4) to a rectangle(N*4) which has miniual area to cover it
    return np.hstack((boxes[:, ::2].min(axis=1).reshape((-1, 1)), boxes[:, 1::2].min(axis=1).reshape((-1, 1)), boxes[:, ::2].max(axis=1).reshape((-1, 1)), boxes[:, 1::2].max(axis=1).reshape((-1, 1))))


def _quad2boxlist(boxes):
    res = []
    for i in range(boxes.shape[0]):
        res.append([[boxes[i][0], boxes[i][1]], [boxes[i][2], boxes[i][3]], [boxes[i][4], boxes[i][5]], [boxes[i][6], boxes[i][7]]])
    return res

def _boxlist2quads(boxlist):
    res = np.zeros((len(boxlist), 8))
    for i, box in enumerate(boxlist):
        # print(box)
        res[i] = np.array([box[0][0], box[0][1], box[1][0], box[1][1], box[2][0], box[2][1], box[3][0], box[3][1]])
    return res

def _rotate_image(im, polygons, angle):
    new_polygons = polygons
    ## rotate image first
    height, width, _ = im.shape
    ## get the minimal rect to cover the rotated image
    img_box = np.array([[0, 0, width, 0, width, height, 0, height]])
    rotated_img_box = _quad2minrect(_rotate_polygons(img_box, -1*angle, (width/2, height/2)))
    r_height = int(max(rotated_img_box[0][3], rotated_img_box[0][1]) - min(rotated_img_box[0][3], rotated_img_box[0][1]))
    r_width = int(max(rotated_img_box[0][2], rotated_img_box[0][0]) - min(rotated_img_box[0][2], rotated_img_box[0][0]))
    r_height_padding = max(r_height, height)
    r_width_padding = max(r_width, width)
    ## padding im
    im_padding = np.zeros((r_height_padding, r_width_padding, 3))
    start_h, start_w = int((r_height_padding - height)/2.0), int((r_width_padding - width)/2.0)
    # start_h = max(start_h, 0)
    # start_w = max(start_w, 0)
    end_h, end_w = start_h + height, start_w + width
    # print(start_h, end_h, start_w, end_w, im.shape)
    im_padding[start_h:end_h, start_w:end_w, :] = im

    M = cv2.getRotationMatrix2D((r_width/2, r_height/2), angle, 1)
    im = cv2.warpAffine(im_padding, M, (r_width, r_height))
    
    ## polygons
    new_polygons = _rotate_segms(polygons, -1*angle, (r_width/2, r_height/2), start_h, start_w)

    return im, new_polygons

def _rotate_polygons(polygons, angle, r_c):
    ## polygons: N*8
    ## r_x: rotate center x
    ## r_y: rotate center y
    ## angle: -15~15

    poly_list = _quad2boxlist(polygons)
    rotate_boxes_list = []
    for poly in poly_list:
        box = Polygon(poly)
        rbox = affinity.rotate(box, angle, r_c)
        if len(list(rbox.exterior.coords))<5:
            print(poly)
            print(rbox)
        # assert(len(list(rbox.exterior.coords))>=5)
        rotate_boxes_list.append(rbox.boundary.coords[:-1])
    res = _boxlist2quads(rotate_boxes_list)
    return res

def _rotate_segms(polygons, angle, r_c, start_h, start_w):
    ## polygons: N*8
    ## r_x: rotate center x
    ## r_y: rotate center y
    ## angle: -15~15
    poly_list=[]
    for polygon in polygons:
        tmp=[]
        for i in range(int(len(polygon) / 2)):
            tmp.append([polygon[2*i] + start_w, polygon[2*i+1] + start_h])
        poly_list.append(tmp)

    rotate_boxes_list = []
    for poly in poly_list:
        box = Polygon(poly)
        rbox = affinity.rotate(box, angle, r_c)
        if len(list(rbox.exterior.coords))<5:
            print(poly)
            print(rbox)
        rotate_boxes_list.append(rbox.boundary.coords[:-1])
    res = []
    for i, box in enumerate(rotate_boxes_list):
        tmp = []
        for point in box:
            tmp.append(point[0])
            tmp.append(point[1])
        res.append([tmp])

    return res

def _read_gt(gt_path):
    polygons = []
    words = []
    with open(gt_path, 'r') as fid:
        lines = fid.readlines()
        for line in lines:
            line = line.strip()
            polygon = line.split(',')[:8]
            word = line.split(',')[8]
            polygon = [float(x) for x in polygon]
            polygons.append(polygon)
            words.append(word)
    return polygons, words

def format_new_gt(polygons, words, new_gt_path):
    with open(new_gt_path, 'wt') as fid:
        for polygon, word in zip(polygons, words):
            # print(polygon)
            polygon = [str(int(x)) for x in polygon[0]]
            # polygon = [str(int(x)) for x in polygon]
            line = ','.join(polygon) + ',' + word
            # print(line)
            fid.write(line+'\n')

def visu_gt(img, polygons, visu_path):
    for polygon in polygons:
        pts = np.array(polygon, np.int32)
        pts = pts.reshape((-1,1,2))
        cv2.polylines(img,[pts],True,(0,255,255))
    cv2.imwrite(visu_path, img)


img_dir = '../datasets/icdar2013/test_images'
gt_dir = '../datasets/icdar2013/test_gts'
angle = 45
new_img_dir = '../datasets/icdar2013/rotated_test_images'+'_'+str(angle)
new_gt_dir = '../datasets/icdar2013/rotated_test_gts'+'_'+str(angle)
if not os.path.isdir(new_img_dir):
    os.mkdir(new_img_dir)
if not os.path.isdir(new_gt_dir):
    os.mkdir(new_gt_dir)

visu_dir = '../output/visu/'

for i in range(233):
    img_name = 'img_' + str(i+1) + '.jpg'
    img_path = os.path.join(img_dir, img_name)
    img = cv2.imread(img_path)
    gt_path = os.path.join(gt_dir, img_name + '.txt')
    new_img_path = os.path.join(new_img_dir, img_name)
    visu_path = os.path.join(visu_dir, img_name)
    new_gt_path = os.path.join(new_gt_dir, 'gt_' + img_name.split('.')[0] + '.txt')
    polygons, words = _read_gt(gt_path)
    # print(img_name)
    if angle == 90:
        (h, w) = img.shape[:2]
        img = cv2.transpose(img)
        img = cv2.flip(img,flipCode=0)
        # M = cv2.getRotationMatrix2D(center, 90, 1)
        # img = cv2.warpAffine(img, M, (h, w))
        new_polygons = [[polygon[1], w-polygon[0], polygon[3], w-polygon[2], polygon[5], w-polygon[4], polygon[7], w-polygon[6]] for polygon in polygons]
    else:
        img, new_polygons = _rotate_image(img, polygons, angle)
    format_new_gt(new_polygons, words, new_gt_path)
    # visu_gt(img, new_polygons, visu_path)
    cv2.imwrite(new_img_path, img)