Geo_dataset.py

import multiprocessing
import os.path as osp
import math
import json
import numpy as np
from numpy import random
import cv2
import torch
from PIL import Image, ImageOps
from tqdm import tqdm
import albumentations as A
from torch.utils.data import Dataset
from shapely.geometry import Polygon

from transform import ComposedTransformation

def cal_distance(x1, y1, x2, y2):
    '''calculate the Euclidean distance'''
    return math.sqrt((x1 - x2)**2 + (y1 - y2)**2)


def move_points(vertices, index1, index2, r, coef):
    '''move the two points to shrink edge
    Input:
        vertices: vertices of text region <numpy.ndarray, (8,)>
        index1  : offset of point1
        index2  : offset of point2
        r       : [r1, r2, r3, r4] in paper
        coef    : shrink ratio in paper
    Output:
        vertices: vertices where one edge has been shinked
    '''
    index1 = index1 % 4
    index2 = index2 % 4
    x1_index = index1 * 2 + 0
    y1_index = index1 * 2 + 1
    x2_index = index2 * 2 + 0
    y2_index = index2 * 2 + 1

    r1 = r[index1]
    r2 = r[index2]
    length_x = vertices[x1_index] - vertices[x2_index]
    length_y = vertices[y1_index] - vertices[y2_index]
    length = cal_distance(vertices[x1_index], vertices[y1_index], vertices[x2_index], vertices[y2_index])
    if length > 1:
        ratio = (r1 * coef) / length
        vertices[x1_index] += ratio * (-length_x)
        vertices[y1_index] += ratio * (-length_y)
        ratio = (r2 * coef) / length
        vertices[x2_index] += ratio * length_x
        vertices[y2_index] += ratio * length_y
    return vertices


def shrink_poly(vertices, coef=0.3):
    '''shrink the text region
    Input:
        vertices: vertices of text region <numpy.ndarray, (8,)>
        coef    : shrink ratio in paper
    Output:
        v       : vertices of shrinked text region <numpy.ndarray, (8,)>
    '''
    x1, y1, x2, y2, x3, y3, x4, y4 = vertices
    r1 = min(cal_distance(x1,y1,x2,y2), cal_distance(x1,y1,x4,y4))
    r2 = min(cal_distance(x2,y2,x1,y1), cal_distance(x2,y2,x3,y3))
    r3 = min(cal_distance(x3,y3,x2,y2), cal_distance(x3,y3,x4,y4))
    r4 = min(cal_distance(x4,y4,x1,y1), cal_distance(x4,y4,x3,y3))
    r = [r1, r2, r3, r4]

    # obtain offset to perform move_points() automatically
    if cal_distance(x1,y1,x2,y2) + cal_distance(x3,y3,x4,y4) > \
       cal_distance(x2,y2,x3,y3) + cal_distance(x1,y1,x4,y4):
        offset = 0 # two longer edges are (x1y1-x2y2) & (x3y3-x4y4)
    else:
        offset = 1 # two longer edges are (x2y2-x3y3) & (x4y4-x1y1)

    v = vertices.copy()
    v = move_points(v, 0 + offset, 1 + offset, r, coef)
    v = move_points(v, 2 + offset, 3 + offset, r, coef)
    v = move_points(v, 1 + offset, 2 + offset, r, coef)
    v = move_points(v, 3 + offset, 4 + offset, r, coef)
    return v


def get_rotate_mat(theta):
    '''positive theta value means rotate clockwise'''
    return np.array([[math.cos(theta), -math.sin(theta)], [math.sin(theta), math.cos(theta)]])


def rotate_vertices(vertices, theta, anchor=None):
    '''rotate vertices around anchor
    Input:
        vertices: vertices of text region <numpy.ndarray, (8,)>
        theta   : angle in radian measure
        anchor  : fixed position during rotation
    Output:
        rotated vertices <numpy.ndarray, (8,)>
    '''
    v = vertices.reshape((4,2)).T
    if anchor is None:
        anchor = v[:,:1]
    rotate_mat = get_rotate_mat(theta)
    res = np.dot(rotate_mat, v - anchor)
    return (res + anchor).T.reshape(-1)


def get_boundary(vertices):
    '''get the tight boundary around given vertices
    Input:
        vertices: vertices of text region <numpy.ndarray, (8,)>
    Output:
        the boundary
    '''
    x1, y1, x2, y2, x3, y3, x4, y4 = vertices
    x_min = min(x1, x2, x3, x4)
    x_max = max(x1, x2, x3, x4)
    y_min = min(y1, y2, y3, y4)
    y_max = max(y1, y2, y3, y4)
    return x_min, x_max, y_min, y_max


def cal_error(vertices):
    '''default orientation is x1y1 : left-top, x2y2 : right-top, x3y3 : right-bot, x4y4 : left-bot
    calculate the difference between the vertices orientation and default orientation
    Input:
        vertices: vertices of text region <numpy.ndarray, (8,)>
    Output:
        err     : difference measure
    '''
    x_min, x_max, y_min, y_max = get_boundary(vertices)
    x1, y1, x2, y2, x3, y3, x4, y4 = vertices
    err = cal_distance(x1, y1, x_min, y_min) + cal_distance(x2, y2, x_max, y_min) + \
          cal_distance(x3, y3, x_max, y_max) + cal_distance(x4, y4, x_min, y_max)
    return err


def find_min_rect_angle(vertices):
    '''find the best angle to rotate poly and obtain min rectangle
    Input:
        vertices: vertices of text region <numpy.ndarray, (8,)>
    Output:
        the best angle <radian measure>
    '''
    angle_interval = 1
    angle_list = list(range(-90, 90, angle_interval))
    area_list = []
    for theta in angle_list:
        rotated = rotate_vertices(vertices, theta / 180 * math.pi)
        x1, y1, x2, y2, x3, y3, x4, y4 = rotated
        temp_area = (max(x1, x2, x3, x4) - min(x1, x2, x3, x4)) * \
                    (max(y1, y2, y3, y4) - min(y1, y2, y3, y4))
        area_list.append(temp_area)

    sorted_area_index = sorted(list(range(len(area_list))), key=lambda k: area_list[k])
    min_error = float('inf')
    best_index = -1
    rank_num = 10
    # find the best angle with correct orientation
    for index in sorted_area_index[:rank_num]:
        rotated = rotate_vertices(vertices, angle_list[index] / 180 * math.pi)
        temp_error = cal_error(rotated)
        if temp_error < min_error:
            min_error = temp_error
            best_index = index
    return angle_list[best_index] / 180 * math.pi


def is_cross_text(start_loc, length, vertices):
    '''check if the crop image crosses text regions
    Input:
        start_loc: left-top position
        length   : length of crop image
        vertices : vertices of text regions <numpy.ndarray, (n,8)>
    Output:
        True if crop image crosses text region
    '''
    if vertices.size == 0:
        return False
    start_w, start_h = start_loc
    a = np.array([start_w, start_h, start_w + length, start_h, start_w + length, start_h + length,
                  start_w, start_h + length]).reshape((4, 2))
    p1 = Polygon(a).convex_hull
    for vertice in vertices:
        p2 = Polygon(vertice.reshape((4, 2))).convex_hull
        inter = p1.intersection(p2).area
        if 0.01 <= inter / p2.area <= 0.99:
            return True
    return False


def crop_img(img, vertices, labels, length):
    h, w = img.shape[:2]
    if h >= w and w < length:
        img = cv2.resize(img,(length, int(h * length / w)),interpolation=cv2.INTER_AREA)
    elif h < w and h < length:
        img = cv2.resize(img,(int(w * length / h), length),interpolation=cv2.INTER_AREA)
    ratio_w = img.shape[1] / w
    ratio_h = img.shape[0] / h
    assert(ratio_w >= 1 and ratio_h >= 1)

    new_vertices = np.zeros(vertices.shape)
    if vertices.size > 0:
        new_vertices[:,[0,2,4,6]] = vertices[:,[0,2,4,6]] * ratio_w
        new_vertices[:,[1,3,5,7]] = vertices[:,[1,3,5,7]] * ratio_h

    # find random position
    remain_h = img.shape[0] - length
    remain_w = img.shape[1] - length
    flag = True
    cnt = 0
    while flag and cnt < 1000:
        cnt += 1
        start_w = int(np.random.rand() * remain_w)
        start_h = int(np.random.rand() * remain_h)
        flag = is_cross_text([start_w, start_h], length, new_vertices[labels==1,:])
    box = (start_w, start_h, start_w + length, start_h + length)
    region = img[box]
    if new_vertices.size == 0:
        return region, new_vertices

    new_vertices[:,[0,2,4,6]] -= start_w
    new_vertices[:,[1,3,5,7]] -= start_h
    return region, new_vertices


def rotate_all_pixels(rotate_mat, anchor_x, anchor_y, length):
    '''get rotated locations of all pixels for next stages
    Input:
        rotate_mat: rotatation matrix
        anchor_x  : fixed x position
        anchor_y  : fixed y position
        length    : length of image
    Output:
        rotated_x : rotated x positions <numpy.ndarray, (length,length)>
        rotated_y : rotated y positions <numpy.ndarray, (length,length)>
    '''
    x = np.arange(length)
    y = np.arange(length)
    x, y = np.meshgrid(x, y)
    x_lin = x.reshape((1, x.size))
    y_lin = y.reshape((1, x.size))
    coord_mat = np.concatenate((x_lin, y_lin), 0)
    rotated_coord = np.dot(rotate_mat, coord_mat - np.array([[anchor_x], [anchor_y]])) + \
                                                   np.array([[anchor_x], [anchor_y]])
    rotated_x = rotated_coord[0, :].reshape(x.shape)
    rotated_y = rotated_coord[1, :].reshape(y.shape)
    return rotated_x, rotated_y


def resize_img(img, vertices, size):
    h, w = img.shape[:2]
    ratio = size / max(h, w)
    if w > h:
        # img = img.resize((size, int(h * ratio)), Image.BILINEAR)
        img = cv2.resize(img,(size, int(h * ratio)),interpolation=cv2.INTER_AREA)
    else:
        # img = img.resize((int(w * ratio), size), Image.BILINEAR)
        img = cv2.resize(img,(int(w * ratio), size),interpolation=cv2.INTER_AREA)
    new_vertices = vertices * ratio
    return img, new_vertices


def adjust_height(img, vertices, ratio=0.2):
    '''adjust height of image to aug data
    Input:
        img         : PIL Image
        vertices    : vertices of text regions <numpy.ndarray, (n,8)>
        ratio       : height changes in [0.8, 1.2]
    Output:
        img         : adjusted PIL Image
        new_vertices: adjusted vertices
    '''
    ratio_h = 1 + ratio * (np.random.rand() * 2 - 1)
    old_h, old_w = img.shape[:2]
    new_h = int(np.around(old_h * ratio_h))
    img = cv2.resize(img,(old_w, new_h),interpolation=cv2.INTER_AREA)

    new_vertices = vertices.copy()
    if vertices.size > 0:
        new_vertices[:,[1,3,5,7]] = vertices[:,[1,3,5,7]] * (new_h / old_h)
    return img, new_vertices


def rotate_img(img, vertices, angle_range=10):
    '''rotate image [-10, 10] degree to aug data
    Input:
        img         : PIL Image
        vertices    : vertices of text regions <numpy.ndarray, (n,8)>
        angle_range : rotate range
    Output:
        img         : rotated PIL Image
        new_vertices: rotated vertices
    '''
    h,w = img.shape[:2]
    center_x = (w - 1) / 2
    center_y = (h - 1) / 2
    angle = angle_range * (np.random.rand() * 2 - 1)
    met = cv2.getRotationMatrix2D((center_x,center_y),angle,1)
    img = cv2.warpAffine(img,met,(0,0))
    new_vertices = np.zeros(vertices.shape)
    for i, vertice in enumerate(vertices):
        new_vertices[i,:] = rotate_vertices(vertice, -angle / 180 * math.pi, np.array([[center_x],[center_y]]))
    return img, new_vertices


def generate_roi_mask(image, vertices, labels):
    mask = np.ones(image.shape[:2], dtype=np.float32)
    ignored_polys = []
    for vertice, label in zip(vertices, labels):
        if label == 0:
            ignored_polys.append(np.around(vertice.reshape((4, 2))).astype(np.int32))
    cv2.fillPoly(mask, ignored_polys, 0)
    return mask


def filter_vertices(vertices, labels, ignore_under=0, drop_under=0):
    if drop_under == 0 and ignore_under == 0:
        return vertices, labels

    new_vertices, new_labels = vertices.copy(), labels.copy()

    areas = np.array([Polygon(v.reshape((4, 2))).convex_hull.area for v in vertices])
    labels[areas < ignore_under] = 0

    if drop_under > 0:
        passed = areas >= drop_under
        new_vertices, new_labels = new_vertices[passed], new_labels[passed]

    return new_vertices, new_labels

class SceneTextDataset(Dataset):
    def __init__(self, root_dir, split='train', image_size=1024, crop_size=512, color_jitter=True,
                    normalize=True, augmentation=True):
        with open(osp.join(root_dir, 'ufo/{}.json'.format(split)), 'r') as f:
            anno = json.load(f)

        self.anno = anno
        self.image_fnames = sorted(anno['images'].keys())
        self.image_dir = osp.join(root_dir, 'images')

        self.augmentation = augmentation
        self.image_size, self.crop_size = image_size, crop_size
        self.color_jitter, self.normalize = color_jitter, normalize

        self.images = []
        self.vertices = []
        self.labels = []

                
        self.transforms = []
        for crop_size in [1.0,0.9,0.8,0.7,0.6,0.5,0.4,0.3,0.2]:
            self.transforms.append(ComposedTransformation(
                rotate_range=np.random.randint(30),
                crop_aspect_ratio=1.0, crop_size=(crop_size, crop_size),
                hflip=False, vflip=False, random_translate=True,
                resize_to=512,
                min_image_overlap=0.9, min_bbox_overlap=1.0, min_bbox_count=1, allow_partial_occurrence=True,
                max_random_trials=100,
                brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5,
                normalize=True, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), to_tensor=False
            ))
        self.load_image()
        
    def __load_item__(self, image_fname):
        image_fpath = osp.join(self.image_dir, image_fname)
        image = cv2.imread(image_fpath,cv2.IMREAD_COLOR)
        try:
            image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
        except:
            print('FAIL to load :',image_fpath)
            return None
        vertices, labels = [], []
        for word_info in self.anno['images'][image_fname]['words'].values():
            vertices.append(np.array(word_info['points']).flatten())
            labels.append(int(not word_info['illegibility']))
        vertices, labels = np.array(vertices, dtype=np.float32), np.array(labels, dtype=np.int64)
        vertices, labels = filter_vertices(vertices, labels, ignore_under=10, drop_under=1)
        image, vertices = resize_img(image, vertices, self.image_size)
        return image, vertices, labels

    def load_image(self):
        with multiprocessing.Pool() as pool:
            items = pool.map(self.__load_item__, tqdm(self.image_fnames))
            items = [ item for item in items if item is not None]
            for image, vertices, labels in items:
                self.images.append(image)
                self.vertices.append(vertices)
                self.labels.append(labels)

    def __len__(self):
        return len(self.image_fnames)

    def __getitem__(self, idx):
        image = self.images[idx]
        vertices = self.vertices[idx]
        labels = self.labels[idx]


        image, vertices = adjust_height(image, vertices)
        image, vertices = rotate_img(image, vertices, angle_range=10)

        transform = random.choice(self.transforms)
        transformed = transform(image=image, word_bboxes=vertices.reshape(-1,4,2))
        image = transformed['image']
        vertices = transformed['word_bboxes']


        word_bboxes = np.reshape(vertices, (-1, 4, 2))
        roi_mask = generate_roi_mask(image, vertices, labels)

        return image, word_bboxes, roi_mask
      
class SceneTextDatasetNoAug(Dataset):
    def __init__(self, root_dir, split='train', image_size=1024, crop_size=512, color_jitter=True,
                    normalize=True, augmentation=True):
        self.input_size = crop_size
        with open(osp.join(root_dir, 'ufo/{}.json'.format(split)), 'r') as f:
            anno = json.load(f)

        self.anno = anno
        self.image_fnames = sorted(anno['images'].keys())
        self.image_dir = osp.join(root_dir, 'images')

        self.image_size, self.crop_size = image_size, crop_size

        self.images = []
        self.vertices = []
        self.labels = []

        self.load_image()

    def __len__(self):
        return len(self.image_fnames)
    
    def __load_item__(self, image_fname):
        image_fpath = osp.join(self.image_dir, image_fname)
        image = cv2.imread(image_fpath,cv2.IMREAD_COLOR)
        try:
            image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
        except:
            print('FAIL to load :',image_fpath)
            return None
        vertices, labels = [], []
        for word_info in self.anno['images'][image_fname]['words'].values():
            vertices.append(np.array(word_info['points']).flatten())
            labels.append(int(not word_info['illegibility']))
        vertices, labels = np.array(vertices, dtype=np.float32), np.array(labels, dtype=np.int64)
        vertices, labels = filter_vertices(vertices, labels, ignore_under=10, drop_under=1)
        image, vertices = resize_img(image, vertices, self.image_size)
        image, vertices = self.__convert__(image, vertices, self.input_size)
        return image, vertices, labels

    def load_image(self):
        with multiprocessing.Pool() as pool:
            items = pool.map(self.__load_item__, tqdm(self.image_fnames))
            items = [ item for item in items if item is not None]
            for image, vertices, labels in items:
                self.images.append(image)
                self.vertices.append(vertices)
                self.labels.append(labels)

    def __convert__(self, img, vertices, size):
        h, w = img.shape[:2]
        ratio = size / max(h, w)
        new_vertices = vertices * ratio
        if w > h:
            dst_size = (size, int(h * ratio))
            img = cv2.resize(img, dst_size, interpolation=cv2.INTER_AREA)
            img = cv2.copyMakeBorder(img, 0, abs(dst_size[0] - dst_size[1]), 0, 0, cv2.BORDER_CONSTANT, value=0)
        else:
            dst_size = (int(w * ratio), size)
            img = cv2.resize(img,(int(w * ratio), size),interpolation=cv2.INTER_AREA)
            img = cv2.copyMakeBorder(img, 0, 0, 0, abs(dst_size[0] - dst_size[1]), cv2.BORDER_CONSTANT, value=0)

        return img, new_vertices

    def __getitem__(self, idx):
        image = self.images[idx]
        vertices = self.vertices[idx]
        labels = self.labels[idx]
        
        transform = A.Compose([A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))])
        image = transform(image=image)['image']

        word_bboxes = np.reshape(vertices, (-1, 4, 2))
        roi_mask = generate_roi_mask(image, vertices, labels)

        return image, word_bboxes, roi_mask