logo
Browse Source

Support input one box

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 2 years ago
parent
commit
c1b5fd29f9
  1. 22
      image_crop_cv2.py

22
image_crop_cv2.py

@ -31,20 +31,28 @@ log = logging.getLogger()
class ImageCropCV2(PyOperator): class ImageCropCV2(PyOperator):
def __init__(self, clamp = False): def __init__(self, clamp = False):
self.clamp = clamp self.clamp = clamp
self.h = None
self.w = None
@staticmethod @staticmethod
def _clamp(x, minimum, maximum): def _clamp(x, minimum, maximum):
return max(minimum, min(x, maximum)) return max(minimum, min(x, maximum))
def get_box(self, img, box):
x1, y1, x2, y2 = box
x1 = ImageCropCV2._clamp(x1, 0, self.w)
x2 = ImageCropCV2._clamp(x2, 0, self.w)
y1 = ImageCropCV2._clamp(y1, 0, self.h)
y2 = ImageCropCV2._clamp(y2, 0, self.h)
return img[y1:y2,x1:x2,:]
def __call__(self, img: np.ndarray, bboxes: List[Tuple]): def __call__(self, img: np.ndarray, bboxes: List[Tuple]):
h, w, _ = img.shape
self.h, self.w, _ = img.shape
res = [] res = []
try:
res.append(self.get_box(img, bboxes))
except:
for box in bboxes: for box in bboxes:
x1, y1, x2, y2 = box
x1 = ImageCropCV2._clamp(x1, 0, w)
x2 = ImageCropCV2._clamp(x2, 0, w)
y1 = ImageCropCV2._clamp(y1, 0, h)
y2 = ImageCropCV2._clamp(y2, 0, h)
res.append(img[y1:y2,x1:x2,:])
res.append(self.get_box(img, box))
return res return res

Loading…
Cancel
Save