| import numpy as np |
| import cv2 |
| from enum import IntEnum |
|
|
|
|
| class SegIEPolyType(IntEnum): |
| EXCLUDE = 0 |
| INCLUDE = 1 |
| |
|
|
|
|
| class SegIEPoly(): |
| def __init__(self, type=None, pts=None, **kwargs): |
| self.type = type |
| |
| if pts is None: |
| pts = np.empty( (0,2), dtype=np.float32 ) |
| else: |
| pts = np.float32(pts) |
| self.pts = pts |
| self.n_max = self.n = len(pts) |
|
|
| def dump(self): |
| return {'type': int(self.type), |
| 'pts' : self.get_pts(), |
| } |
| |
| def identical(self, b): |
| if self.n != b.n: |
| return False |
| return (self.pts[0:self.n] == b.pts[0:b.n]).all() |
| |
| def get_type(self): |
| return self.type |
|
|
| def add_pt(self, x, y): |
| self.pts = np.append(self.pts[0:self.n], [ ( float(x), float(y) ) ], axis=0).astype(np.float32) |
| self.n_max = self.n = self.n + 1 |
|
|
| def undo(self): |
| self.n = max(0, self.n-1) |
| return self.n |
|
|
| def redo(self): |
| self.n = min(len(self.pts), self.n+1) |
| return self.n |
|
|
| def redo_clip(self): |
| self.pts = self.pts[0:self.n] |
| self.n_max = self.n |
|
|
| def insert_pt(self, n, pt): |
| if n < 0 or n > self.n: |
| raise ValueError("insert_pt out of range") |
| self.pts = np.concatenate( (self.pts[0:n], pt[None,...].astype(np.float32), self.pts[n:]), axis=0) |
| self.n_max = self.n = self.n+1 |
| |
| def remove_pt(self, n): |
| if n < 0 or n >= self.n: |
| raise ValueError("remove_pt out of range") |
| self.pts = np.concatenate( (self.pts[0:n], self.pts[n+1:]), axis=0) |
| self.n_max = self.n = self.n-1 |
| |
| def get_last_point(self): |
| return self.pts[self.n-1].copy() |
|
|
| def get_pts(self): |
| return self.pts[0:self.n].copy() |
| |
| def get_pts_count(self): |
| return self.n |
|
|
| def set_point(self, id, pt): |
| self.pts[id] = pt |
| |
| def set_points(self, pts): |
| self.pts = np.array(pts) |
| self.n_max = self.n = len(pts) |
| |
| def mult_points(self, val): |
| self.pts *= val |
| |
| |
|
|
| class SegIEPolys(): |
| def __init__(self): |
| self.polys = [] |
|
|
| def identical(self, b): |
| polys_len = len(self.polys) |
| o_polys_len = len(b.polys) |
| if polys_len != o_polys_len: |
| return False |
| |
| return all ([ a_poly.identical(b_poly) for a_poly, b_poly in zip(self.polys, b.polys) ]) |
| |
| def add_poly(self, ie_poly_type): |
| poly = SegIEPoly(ie_poly_type) |
| self.polys.append (poly) |
| return poly |
|
|
| def remove_poly(self, poly): |
| if poly in self.polys: |
| self.polys.remove(poly) |
|
|
| def has_polys(self): |
| return len(self.polys) != 0 |
| |
| def get_poly(self, id): |
| return self.polys[id] |
|
|
| def get_polys(self): |
| return self.polys |
| |
| def get_pts_count(self): |
| return sum([poly.get_pts_count() for poly in self.polys]) |
| |
| def sort(self): |
| poly_by_type = { SegIEPolyType.EXCLUDE : [], SegIEPolyType.INCLUDE : [] } |
|
|
| for poly in self.polys: |
| poly_by_type[poly.type].append(poly) |
| |
| self.polys = poly_by_type[SegIEPolyType.INCLUDE] + poly_by_type[SegIEPolyType.EXCLUDE] |
|
|
| def __iter__(self): |
| for poly in self.polys: |
| yield poly |
|
|
| def overlay_mask(self, mask): |
| h,w,c = mask.shape |
| white = (1,)*c |
| black = (0,)*c |
| for poly in self.polys: |
| pts = poly.get_pts().astype(np.int32) |
| if len(pts) != 0: |
| cv2.fillPoly(mask, [pts], white if poly.type == SegIEPolyType.INCLUDE else black ) |
|
|
| def dump(self): |
| return {'polys' : [ poly.dump() for poly in self.polys ] } |
| |
| def mult_points(self, val): |
| for poly in self.polys: |
| poly.mult_points(val) |
| |
| @staticmethod |
| def load(data=None): |
| ie_polys = SegIEPolys() |
| if data is not None: |
| if isinstance(data, list): |
| |
| ie_polys.polys = [ SegIEPoly(type=type, pts=pts) for (type, pts) in data ] |
| elif isinstance(data, dict): |
| ie_polys.polys = [ SegIEPoly(**poly_cfg) for poly_cfg in data['polys'] ] |
| |
| ie_polys.sort() |
| |
| return ie_polys |
