File size: 1,370 Bytes
d02e83e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import cv2
import numpy as np
def filter_mask(output, image):
image_h, image_w = image.shape[:2]
predict_mask = (output.squeeze().cpu().numpy().round())
predict_mask = predict_mask.astype('uint8')*255
predict_mask = cv2.resize( predict_mask, (image_w, image_h) )
ret, thresh = cv2.threshold(predict_mask, 127, 255, 0)
contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
points = contours[0]
rect = cv2.minAreaRect(points)
box = cv2.boxPoints(rect)
box = np.int0(box)
img_vis = cv2.drawContours(image.copy(),[box],0,(255,0,0), 6)
(cX, cY), (w, h), angle = rect
if w<h:
angle -= 90
#(cX, cY) = (image_w // 2, image_h // 2)
M = cv2.getRotationMatrix2D((cX, cY), angle, 1.0)
rotated = cv2.warpAffine(image, M, (round(image_w), round(image_h)))
ones = np.ones(shape=(len(points), 1))
points = np.squeeze(points)
points_ones = np.hstack([points, ones])
points_rotate = M.dot(points_ones.T).T
(cX, cY), (w, h), angle = cv2.minAreaRect(points_rotate.round().astype(int))
if angle < 45 :
crop_img = rotated[round(cY-h//2):round(cY+h//2), round(cX - w//2):round(cX + w//2)]
else:
crop_img = rotated[round(cY-w//2):round(cY+w//2), round(cX - h//2):round(cX + h//2)]
return crop_img, img_vis |