|
@@ -8,12 +8,11 @@ from abc import ABC, abstractmethod, abstractproperty
|
|
from . import utils
|
|
from . import utils
|
|
|
|
|
|
class Parts(object):
|
|
class Parts(object):
|
|
|
|
+
|
|
def __init__(self, image, part_annotations, rescale_size):
|
|
def __init__(self, image, part_annotations, rescale_size):
|
|
super(Parts, self).__init__()
|
|
super(Parts, self).__init__()
|
|
- annots = utils.rescale_parts(image, part_annotations, rescale_size)
|
|
|
|
|
|
|
|
- self._parts = [BasePart.new(a) for a in annots]
|
|
|
|
- self.rescale_size = rescale_size
|
|
|
|
|
|
+ self._parts = [BasePart.new(image, a, rescale_size) for a in part_annotations]
|
|
|
|
|
|
def __getitem__(self, i):
|
|
def __getitem__(self, i):
|
|
return self._parts[i]
|
|
return self._parts[i]
|
|
@@ -68,24 +67,31 @@ class Parts(object):
|
|
|
|
|
|
|
|
|
|
class BasePart(ABC):
|
|
class BasePart(ABC):
|
|
- def __init__(self, annotation):
|
|
|
|
- super(BasePart, self).__init__()
|
|
|
|
- self.read_annotation(annotation)
|
|
|
|
|
|
|
|
@staticmethod
|
|
@staticmethod
|
|
- def new(annotation):
|
|
|
|
|
|
+ def new(image, annotation, rescale_size=-1):
|
|
if len(annotation) == 4:
|
|
if len(annotation) == 4:
|
|
- return LocationPart(annotation)
|
|
|
|
|
|
+ return LocationPart(image, annotation, rescale_size)
|
|
elif len(annotation) == 5:
|
|
elif len(annotation) == 5:
|
|
- return BBoxPart(annotation)
|
|
|
|
|
|
+ return BBoxPart(image, annotation, rescale_size)
|
|
else:
|
|
else:
|
|
raise ValueError("Unknown part annotation format: {}".format(annotation))
|
|
raise ValueError("Unknown part annotation format: {}".format(annotation))
|
|
|
|
|
|
|
|
|
|
@abstractmethod
|
|
@abstractmethod
|
|
- def read_annotation(self, annotation):
|
|
|
|
|
|
+ def __init__(self, *args, **kwargs):
|
|
raise NotImplementedError
|
|
raise NotImplementedError
|
|
|
|
|
|
|
|
+ def rescale(self, image, annotation, rescale_size):
|
|
|
|
+ if rescale_size is not None and rescale_size > 0:
|
|
|
|
+ h, w, c = utils.dimensions(image)
|
|
|
|
+ scale = np.array([w, h]) / rescale_size
|
|
|
|
+ xy = annotation[1:3]
|
|
|
|
+ xy = xy * scale
|
|
|
|
+ annotation[1:3] = xy
|
|
|
|
+
|
|
|
|
+ return annotation
|
|
|
|
+
|
|
@property
|
|
@property
|
|
def is_visible(self):
|
|
def is_visible(self):
|
|
return self._is_visible
|
|
return self._is_visible
|
|
@@ -98,20 +104,32 @@ class BasePart(ABC):
|
|
def xy(self):
|
|
def xy(self):
|
|
return np.array([self.x, self.y])
|
|
return np.array([self.x, self.y])
|
|
|
|
|
|
- @abstractmethod
|
|
|
|
- def crop(self, *args, **kwargs):
|
|
|
|
- raise NotImplementedError
|
|
|
|
|
|
+ def crop(self, im, w, h, padding_mode="edge", is_location=True):
|
|
|
|
+
|
|
|
|
+ x, y = self.xy
|
|
|
|
+ pad_h, pad_w = h // 2, w // 2
|
|
|
|
+
|
|
|
|
+ padded_im = np.pad(im, [(pad_h, pad_h), (pad_w, pad_w), [0,0]], mode=padding_mode)
|
|
|
|
+ x0, y0 = x + pad_w, y + pad_h
|
|
|
|
+
|
|
|
|
+ if is_location:
|
|
|
|
+ x0, y0 = x0 - w // 2, y0 - h // 2
|
|
|
|
+
|
|
|
|
+ return padded_im[y0:y0+h, x0:x0+w]
|
|
|
|
+
|
|
|
|
|
|
def plot(self, **kwargs):
|
|
def plot(self, **kwargs):
|
|
return
|
|
return
|
|
|
|
|
|
-
|
|
|
|
class LocationPart(BasePart):
|
|
class LocationPart(BasePart):
|
|
|
|
+ DEFAULT_RATIO = np.sqrt(49 / 400) # 0.35
|
|
|
|
+
|
|
|
|
+ def __init__(self, image, annotation, rescale_size):
|
|
|
|
+ annotation = self.rescale(image, annotation, rescale_size)
|
|
|
|
|
|
- def read_annotation(self, annotation):
|
|
|
|
# here x,y are the center of the part
|
|
# here x,y are the center of the part
|
|
self._id, self.x, self.y, self.is_visible = annotation
|
|
self._id, self.x, self.y, self.is_visible = annotation
|
|
- self._ratio = None
|
|
|
|
|
|
+ self._ratio = LocationPart.DEFAULT_RATIO
|
|
|
|
|
|
def crop(self, image, ratio=None, padding_mode="edge", *args, **kwargs):
|
|
def crop(self, image, ratio=None, padding_mode="edge", *args, **kwargs):
|
|
ratio = ratio or self._ratio
|
|
ratio = ratio or self._ratio
|
|
@@ -121,11 +139,10 @@ class LocationPart(BasePart):
|
|
if not self.is_visible:
|
|
if not self.is_visible:
|
|
return np.zeros((h, w, c), dtype=np.uint8)
|
|
return np.zeros((h, w, c), dtype=np.uint8)
|
|
else:
|
|
else:
|
|
- return utils.crop(image, self.xy, w, h,
|
|
|
|
|
|
+ return super(LocationPart, self).crop(image, w, h,
|
|
padding_mode, is_location=True)
|
|
padding_mode, is_location=True)
|
|
|
|
|
|
-
|
|
|
|
- def reveal(self, im, ratio, *args, **kwargs):
|
|
|
|
|
|
+ def reveal(self, im, ratio=None, *args, **kwargs):
|
|
_h, _w, c = utils.dimensions(im)
|
|
_h, _w, c = utils.dimensions(im)
|
|
w, h = int(_w * ratio), int(_h * ratio)
|
|
w, h = int(_w * ratio), int(_h * ratio)
|
|
x,y = self.xy
|
|
x,y = self.xy
|
|
@@ -146,13 +163,27 @@ class LocationPart(BasePart):
|
|
|
|
|
|
class BBoxPart(BasePart):
|
|
class BBoxPart(BasePart):
|
|
|
|
|
|
- def read_annotation(self, annotation):
|
|
|
|
|
|
+ def __init__(self, image, annotation, rescale_size):
|
|
|
|
+
|
|
|
|
+ annotation = self.rescale(image, annotation, rescale_size)
|
|
# here x,y are top left corner of the part
|
|
# here x,y are top left corner of the part
|
|
self._id, self.x, self.y, self.w, self.h = annotation
|
|
self._id, self.x, self.y, self.w, self.h = annotation
|
|
- self._is_visible = True
|
|
|
|
|
|
+ self.is_visible = True
|
|
|
|
+
|
|
|
|
+ def rescale(self, image, annotation, rescale_size):
|
|
|
|
+ if rescale_size is not None and rescale_size > 0:
|
|
|
|
+ annotation = super(BBoxPart, self).rescale(image, annotation, rescale_size)
|
|
|
|
+ h, w, c = utils.dimensions(image)
|
|
|
|
+ scale = np.array([w, h]) / rescale_size
|
|
|
|
+ wh = annotation[3:5]
|
|
|
|
+ wh = wh * scale
|
|
|
|
+ annotation[3:5] = wh
|
|
|
|
+
|
|
|
|
+ return annotation
|
|
|
|
+
|
|
|
|
|
|
def crop(self, image, padding_mode="edge", *args, **kwargs):
|
|
def crop(self, image, padding_mode="edge", *args, **kwargs):
|
|
- return utils.crop(image, self.xy, self.w, self.h,
|
|
|
|
|
|
+ return super(BBoxPart, self).crop(image, self.w, self.h,
|
|
padding_mode, is_location=False)
|
|
padding_mode, is_location=False)
|
|
|
|
|
|
def reveal(self, im, ratio, *args, **kwargs):
|
|
def reveal(self, im, ratio, *args, **kwargs):
|