|
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 |
|
|
|
|
|
|
|
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 |
