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