interpreter.py

from typing import NamedTuple, List, Callable
import sys
import re
import numpy as np
import torch
from numpy.linalg import norm
from itertools import product, groupby
from PIL import Image

def ccw(A, B, C):
    return (C.y - A.y) * (B.x - A.x) > (B.y - A.y) * (C.x - A.x)


def intersect(A, B, C, D):
    """Do line segments AB and CD intersect?"""
    return ccw(A, C, D) != ccw(B, C, D) and ccw(A, B, C) != ccw(A, B, D)


class Box(NamedTuple):
    x: int
    y: int
    w: int = 0
    h: int = 0

    @property
    def left(self):
        return self.x

    @property
    def right(self):
        return self.x + self.w

    @property
    def top(self):
        return self.y

    @property
    def bottom(self):
        return self.y + self.h

    @property
    def center(self):
        return Box(self.x + self.w // 2, self.y + self.h // 2)

    def corners(self):
        yield Box(self.x, self.y)
        yield Box(self.x + self.w, self.y)
        yield Box(self.x + self.w, self.y + self.h)
        yield Box(self.x, self.y + self.h)

    @property
    def area(self):
        return self.w * self.h

    def intersect(self, other: "Box") -> "Box":
        x1 = max(self.x, other.x)
        x2 = max(x1, min(self.x+self.w, other.x+other.w))
        y1 = max(self.y, other.y)
        y2 = max(y1, min(self.y+self.h, other.y+other.h))
        return Box(x=x1, y=y1, w=x2-x1, h=y2-y1)

    def min_bounding(self, other: "Box") -> "Box":
        corners = list(self.corners())
        corners.extend(other.corners())
        min_x = min_y = float("inf")
        max_x = max_y = -float("inf")

        for item in corners:
            min_x = min(min_x, item.x)
            min_y = min(min_y, item.y)
            max_x = max(max_x, item.x)
            max_y = max(max_y, item.y)

        return Box(min_x, min_y, max_x - min_x, max_y - min_y)

    def expand(self, growth: float = .1) -> "Box":
        factor = 1 + growth
        w = factor * self.w
        h = factor * self.h
        return Box(min_x - (w - self.w) / 2, min_y - (h - self.h) / 2, w, h)


def iou(box1, box2):
    x1 = max(box1.x, box2.x)
    x2 = max(x1, min(box1.x+box1.w, box2.x+box2.w))
    y1 = max(box1.y, box2.y)
    y2 = max(y1, min(box1.y+box1.h, box2.y+box2.h))
    intersection = Box(x=x1, y=y1, w=x2-x1, h=y2-y1)
    intersection_area = intersection.area
    union_area = box1.area+box2.area-intersection_area
    return intersection_area / union_area


def all_equal(iterable):
    """Are all elements the same?"""
    g = groupby(iterable)
    return next(g, True) and not next(g, False)


class spatial:
    """A decorator that converts a predicate over boxes to a function that returns a tensor over all boxes."""

    def __init__(self, arity: int = 2, enforce_antisymmetry: bool = False):
        self.arity = arity
        self.enforce_antisymmetry = enforce_antisymmetry  # Zero out any entries where two boxes are the same.

    def __call__(self, predicate: Callable[[Box], float]) -> Callable[["Environment"], np.ndarray]:
        def _rel(env):
            n_boxes = len(env.boxes)
            tensor = np.empty([n_boxes for _ in range(self.arity)])
            enum_boxes = list(enumerate(env.boxes))
            for pairs in product(*[enum_boxes for _ in range(self.arity)]):
                indices, boxes = zip(*pairs)
                if self.enforce_antisymmetry and len(set(indices)) < len(indices):
                    tensor[indices] = 0.
                else:
                    tensor[indices] = predicate(*boxes)
            return tensor
        return _rel


class Environment:
    def __init__(self, image: Image, boxes: List[Box], executor: "Executor" = None, image_name: str = None):
        self.image = image
        self.boxes = boxes
        self.executor = executor  
        self.image_name = image_name

    def uniform(self) -> np.ndarray:
        n_boxes = len(self.boxes)
        return 1 / n_boxes * np.ones(n_boxes)

    def filter(self,
               caption: str,
               temperature: float = 1.,
               area_threshold: float = 0.0,
               softmax: bool = False,
               expand: float = None
              ) -> np.ndarray:
        """Return a new distribution reflecting the likelihood that `caption` describes the content of each box."""
        area_filtered_dist = torch.from_numpy(self.filter_area(area_threshold)).to(self.executor.device)
        candidate_indices = [i for i in range(len(self.boxes)) if float(area_filtered_dist[i]) > 0.0]
        boxes = [self.boxes[i] for i in candidate_indices]
        if len(boxes) == 0:
            boxes = self.boxes
            candidate_indices = list(range(len(boxes)))
        if expand is not None:
            boxes = [box.expand(expand) for box in boxes]
        result_partial = self.executor(caption, self.image, boxes)
        result_partial = result_partial.float()
        if not softmax:
            result_partial = (result_partial-result_partial.mean()) / (result_partial.std() + 1e-9)
            result_partial = (temperature * result_partial).sigmoid()
            result = torch.zeros((len(self.boxes))).to(result_partial.device)
            result[candidate_indices] = result_partial
        else:
            result = torch.zeros((len(self.boxes))).to(result_partial.device)
            result[candidate_indices] = result_partial.softmax(dim=-1)
        return result.cpu().numpy()

    def filter_area(self, area_threshold: float) -> np.ndarray:
        """Return a new distribution in which all boxes whose area as a fraction of the image is less than the threshold."""
        image_area = self.image.width*self.image.height
        return np.array([1 if self.boxes[i].area/image_area > area_threshold else 0 for i in range(len(self.boxes))])

    @spatial()
    def left_of(b1, b2):
        return (b1.right+b1.left) / 2 < (b2.right+b2.left) / 2

    @spatial()
    def right_of(b1, b2):
        return (b1.right+b1.left) / 2 > (b2.right+b2.left) / 2

    @spatial()
    def above(b1, b2):
        return (b1.bottom+b1.top) < (b2.bottom+b2.top)

    @spatial()
    def below(b1, b2):
        return (b1.bottom+b1.top) > (b2.bottom+b2.top)

    @spatial()
    def bigger_than(b1, b2):
        return b1.area > b2.area

    @spatial()
    def smaller_than(b1, b2):
        return b1.area < b2.area

    @spatial(enforce_antisymmetry=False)
    def within(box1, box2):
        """Return percent of box1 inside box2."""
        intersection = box1.intersect(box2)
        return intersection.area / box1.area

    @spatial(arity=3, enforce_antisymmetry=True)
    def between(box1, box2, box3):
        """How much of box1 lies in min bounding box over box2 and box3?"""
        min_bounding = box2.min_bounding(box3)
        intersect = box1.intersect(min_bounding)
        return intersect.area / box1.area