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