design batch 代码整理

app/service/design_batch/pipeline/contour_detection.py

@@ -0,0 +1,37 @@
+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