torch_lydorn/torchvision/transforms/rasterize.py

import math
import sys
import time

import skimage.morphology
import skimage.io
from PIL import Image, ImageDraw, ImageFilter
import numpy as np
import shapely.geometry
import shapely.affinity
from lydorn_utils import print_utils
from scipy.ndimage.morphology import distance_transform_edt
import cv2 as cv

from functools import partial

import torch_lydorn.torchvision


class Rasterize(object):
    """Rasterize polygons"""

    def __init__(self, fill=True, edges=True, vertices=True, line_width=3, antialiasing=False, return_distances=False,
                 return_sizes=False):
        self.fill = fill
        self.edges = edges
        self.vertices = vertices
        self.line_width = line_width
        self.antialiasing = antialiasing

        if not return_distances and not return_sizes:
            self.raster_func = partial(draw_polygons, fill=self.fill, edges=self.edges, vertices=self.vertices,
                                       line_width=self.line_width, antialiasing=self.antialiasing)
        elif return_distances and return_sizes:
            self.raster_func = partial(compute_raster_distances_sizes, fill=self.fill, edges=self.edges, vertices=self.vertices,
                                       line_width=self.line_width, antialiasing=self.antialiasing)
        else:
            raise NotImplementedError

    def __call__(self, image, polygons):
        """
        If distances is True, also returns distances image
        (sum of distance to closest and second-closest annotation for each pixel).
        Same for sizes (size of annotation the pixel belongs to).

        """
        size = (image.shape[0], image.shape[1])
        out = self.raster_func(polygons, size)
        return out


def compute_raster_distances_sizes(polygons, shape, fill=True, edges=True, vertices=True, line_width=3, antialiasing=False):
    """
    Returns:
         - distances: sum of distance to closest and second-closest annotation for each pixel.
         - size_weights: relative size (normalized by image area) of annotation the pixel belongs to.
    """
    assert type(polygons) == list, "polygons should be a list"

    # Filter out zero-area polygons
    polygons = [polygon for polygon in polygons if 0 < polygon.area]

    # tic = time.time()

    channel_count = fill + edges + vertices
    polygons_raster = np.zeros((*shape, channel_count), dtype=np.uint8)
    distance_maps = np.ones((*shape, len(polygons)))  # Init with max value (distances are normed)
    sizes = np.ones(shape)  # Init with max value (sizes are normed)
    image_area = shape[0] * shape[1]
    for i, polygon in enumerate(polygons):
        minx, miny, maxx, maxy = polygon.bounds
        mini = max(0, math.floor(miny) - 2*line_width)
        minj = max(0, math.floor(minx) - 2*line_width)
        maxi = min(polygons_raster.shape[0], math.ceil(maxy) + 2*line_width)
        maxj = min(polygons_raster.shape[1], math.ceil(maxx) + 2*line_width)
        bbox_shape = (maxi - mini, maxj - minj)
        bbox_polygon = shapely.affinity.translate(polygon, xoff=-minj, yoff=-mini)
        bbox_raster = draw_polygons([bbox_polygon], bbox_shape, fill, edges, vertices, line_width, antialiasing)
        polygons_raster[mini:maxi, minj:maxj] = np.maximum(polygons_raster[mini:maxi, minj:maxj], bbox_raster)
        bbox_mask = 0 < np.sum(bbox_raster, axis=2)  # Polygon interior + edge + vertex
        if bbox_mask.max():  # Make sure mask is not empty
            polygon_mask = np.zeros(shape, dtype=np.bool)
            polygon_mask[mini:maxi, minj:maxj] = bbox_mask
            polygon_dist = cv.distanceTransform(1 - polygon_mask.astype(np.uint8), distanceType=cv.DIST_L2, maskSize=cv.DIST_MASK_5,
                                        dstType=cv.CV_64F)
            polygon_dist /= (polygon_mask.shape[0] + polygon_mask.shape[1])  # Normalize dist
            distance_maps[:, :, i] = polygon_dist

            selem = skimage.morphology.disk(line_width)
            bbox_dilated_mask = skimage.morphology.binary_dilation(bbox_mask, selem=selem)
            sizes[mini:maxi, minj:maxj][bbox_dilated_mask] = polygon.area / image_area

    polygons_raster = np.clip(polygons_raster, 0, 255)
    # skimage.io.imsave("polygons_raster.png", polygons_raster)

    if edges:
        edge_channels = -1 + fill + edges
        # Remove border edges because they correspond to cut buildings:
        polygons_raster[:line_width, :, edge_channels] = 0
        polygons_raster[-line_width:, :, edge_channels] = 0
        polygons_raster[:, :line_width, edge_channels] = 0
        polygons_raster[:, -line_width:, edge_channels] = 0

    distances = compute_distances(distance_maps)
    # skimage.io.imsave("distances.png", distances)

    distances = distances.astype(np.float16)
    sizes = sizes.astype(np.float16)

    # toc = time.time()
    # print(f"Rasterize {len(polygons)} polygons: {toc - tic}s")

    return polygons_raster, distances, sizes


def compute_distances(distance_maps):
    distance_maps.sort(axis=2)
    distance_maps = distance_maps[:, :, :2]
    distances = np.sum(distance_maps, axis=2)
    return distances


def draw_polygons(polygons, shape, fill=True, edges=True, vertices=True, line_width=3, antialiasing=False):
    assert type(polygons) == list, "polygons should be a list"
    assert type(polygons[0]) == shapely.geometry.Polygon, "polygon should be a shapely.geometry.Polygon"

    if antialiasing:
        draw_shape = (2 * shape[0], 2 * shape[1])
        polygons = [shapely.affinity.scale(polygon, xfact=2.0, yfact=2.0, origin=(0, 0)) for polygon in polygons]
        line_width *= 2
    else:
        draw_shape = shape
    # Channels
    fill_channel_index = 0  # Always first channel
    edges_channel_index = fill  # If fill == True, take second channel. If not then take first
    vertices_channel_index = fill + edges  # Same principle as above
    channel_count = fill + edges + vertices
    im_draw_list = []
    for channel_index in range(channel_count):
        im = Image.new("L", (draw_shape[1], draw_shape[0]))
        im_px_access = im.load()
        draw = ImageDraw.Draw(im)
        im_draw_list.append((im, draw))

    for polygon in polygons:
        if fill:
            draw = im_draw_list[fill_channel_index][1]
            draw.polygon(polygon.exterior.coords, fill=255)
            for interior in polygon.interiors:
                draw.polygon(interior.coords, fill=0)
        if edges:
            draw = im_draw_list[edges_channel_index][1]
            draw.line(polygon.exterior.coords, fill=255, width=line_width)
            for interior in polygon.interiors:
                draw.line(interior.coords, fill=255, width=line_width)
        if vertices:
            draw = im_draw_list[vertices_channel_index][1]
            for vertex in polygon.exterior.coords:
                torch_lydorn.torchvision.transforms.functional.draw_circle(draw, vertex, line_width / 2, fill=255)
            for interior in polygon.interiors:
                for vertex in interior.coords:
                    torch_lydorn.torchvision.transforms.functional.draw_circle(draw, vertex, line_width / 2, fill=255)

    im_list = []
    if antialiasing:
        # resize images:
        for im_draw in im_draw_list:
            resize_shape = (shape[1], shape[0])
            im_list.append(im_draw[0].resize(resize_shape, Image.BILINEAR))
    else:
        for im_draw in im_draw_list:
            im_list.append(im_draw[0])

    # Convert image to numpy array with the right number of channels
    array_list = [np.array(im) for im in im_list]
    array = np.stack(array_list, axis=-1)
    return array


def _rasterize_coco(image, polygons):
    import pycocotools.mask as cocomask

    image_size = image.shape[:2]
    mask = np.zeros(image_size)
    for polygon in polygons:
        rle = cocomask.frPyObjects([np.array(polygon.exterior.coords).reshape(-1)], image_size[0], image_size[1])
        m = cocomask.decode(rle)

        for i in range(m.shape[-1]):
            mi = m[:, :, i]
            mi = mi.reshape(image_size)
            mask += mi
    return mask


def _test():
    import skimage.io

    rasterize = Rasterize(fill=True, edges=False, vertices=False, line_width=2, antialiasing=True, return_distances=True, return_sizes=True)

    image = np.zeros((300, 300))
    polygons = [
        shapely.geometry.Polygon([
            [10.5, 10.5],
            [100, 10],
            [100, 150],
            [10, 100],
            [10, 10],
        ]),
        shapely.geometry.Polygon([
            [10+150, 10],
            [100+150, 10],
            [100+150, 100],
            [10+150, 100],
            [10+150, 10],
        ]),
    ]
    polygons_raster, distances, size_weights = rasterize(image, polygons)

    skimage.io.imsave('rasterize.polygons_raster.png', polygons_raster)
    skimage.io.imsave('rasterize.distances.png', distances)
    skimage.io.imsave('rasterize.size_weights.png', size_weights)

    # Rasterize with pycocotools
    coco_mask = _rasterize_coco(image, polygons)
    skimage.io.imsave('rasterize.coco_mask.png', coco_mask)


if __name__ == "__main__":
    _test()