import cv2
import numpy as np


class ContourDetection:
    def __call__(self, result):
        Contour = self.get_contours(result['image'])
        Mask = np.zeros(result['image'].shape[:2], np.uint8)
        if len(Contour):
            Max_contour = Contour[0]
            Epsilon = 0.001 * cv2.arcLength(Max_contour, True)
            Approx = cv2.approxPolyDP(Max_contour, Epsilon, True)
            cv2.drawContours(Mask, [Approx], -1, 255, -1)
        else:
            Mask = np.ones(result['image'].shape[:2], np.uint8) * 255
        # TODO 修复部分图片出现透明的情况 下版本上线
        # img2gray = cv2.cvtColor(result['image'], cv2.COLOR_BGR2GRAY)
        # ret, Mask = cv2.threshold(img2gray, 126, 255, cv2.THRESH_BINARY)
        # Mask = cv2.bitwise_not(Mask)
        if result['pre_mask'] is None:
            result['mask'] = Mask
        else:
            result['mask'] = cv2.bitwise_and(Mask, result['pre_mask'])
        result['front_mask'] = result['mask']
        result['back_mask'] = result['mask']
        return result

    @staticmethod
    def get_contours(image):
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        Edge = cv2.Canny(gray, 10, 150)
        kernel = np.ones((5, 5), np.uint8)
        Edge = cv2.dilate(Edge, kernel=kernel, iterations=1)
        Edge = cv2.erode(Edge, kernel=kernel, iterations=1)
        Contour, _ = cv2.findContours(Edge, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        Contour = sorted(Contour, key=cv2.contourArea, reverse=True)
        return Contour