cv-analysis-service/cv_analysis/utils/structures.py

from json import dumps

import numpy as np
from funcy import identity


class Rectangle:
    def __init__(self, x1=None, y1=None, w=None, h=None, x2=None, y2=None, indent=4, format="xywh", discrete=True):
        make_discrete = int if discrete else identity
        try:
            self.x1 = make_discrete(x1)
            self.y1 = make_discrete(y1)
            self.w = make_discrete(w) if w else make_discrete(x2 - x1)
            self.h = make_discrete(h) if h else make_discrete(y2 - y1)
            self.x2 = make_discrete(x2) if x2 else self.x1 + self.w
            self.y2 = make_discrete(y2) if y2 else self.y1 + self.h
            assert np.isclose(self.x1 + self.w, self.x2)
            assert np.isclose(self.y1 + self.h, self.y2)
            self.indent = indent
            self.format = format
        except Exception as err:
            raise Exception("x1, y1, (w|x2), and (h|y2) must be defined.") from err

    def json_xywh(self):
        return {"x": self.x1, "y": self.y1, "width": self.w, "height": self.h}

    def json_xyxy(self):
        return {"x1": self.x1, "y1": self.y1, "x2": self.x2, "y2": self.y2}

    def json_full(self):
        return {"x1": self.x1, "y1": self.y1, "x2": self.x2, "y2": self.y2, "width": self.w, "height": self.h}

    def json(self):
        json_func = {"xywh": self.json_xywh, "xyxy": self.json_xyxy}.get(self.format, self.json_full)
        return json_func()

    def xyxy(self):
        return self.x1, self.y1, self.x2, self.y2

    def xywh(self):
        return self.x1, self.y1, self.w, self.h

    @classmethod
    def from_xyxy(cls, xyxy_tuple, discrete=True):
        x1, y1, x2, y2 = xyxy_tuple
        return cls(x1=x1, y1=y1, x2=x2, y2=y2, discrete=discrete)

    @classmethod
    def from_xywh(cls, xywh_tuple, discrete=True):
        x, y, w, h = xywh_tuple
        return cls(x1=x, y1=y, w=w, h=h, discrete=discrete)

    def __str__(self):
        return dumps(self.json(), indent=self.indent)

    def __repr__(self):
        return str(self.json())

    def __iter__(self):
        return list(self.json().values()).__iter__()


class Contour:
    def __init__(self):
        pass