38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
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
|