diff --git a/app/service/design/items/pipelines/segmentation.py b/app/service/design/items/pipelines/segmentation.py index ac38b34..2937df6 100644 --- a/app/service/design/items/pipelines/segmentation.py +++ b/app/service/design/items/pipelines/segmentation.py @@ -23,12 +23,18 @@ class Segmentation(object): # 转换颜色空间为 RGB(OpenCV 默认是 BGR) image_rgb = cv2.cvtColor(seg_mask, cv2.COLOR_BGR2RGB) - front_color = np.array([255, 0, 0], dtype=np.uint8) - back_color = np.array([0, 255, 0], dtype=np.uint8) + # 定义红色和绿色的颜色范围 + # 红色范围: 下界 [R-10, G-10, B-10], 上界 [R+10, G+10, B+10] + red_lower = np.array([50, 0, 0], dtype=np.uint8) + red_upper = np.array([255, 50, 50], dtype=np.uint8) + + # 绿色范围: 下界 [R-10, G-10, B-10], 上界 [R+10, G+10, B+10] + green_lower = np.array([0, 50, 0], dtype=np.uint8) + green_upper = np.array([50, 255, 50], dtype=np.uint8) # 创建红色和绿色掩码 - result['front_mask'] = cv2.inRange(image_rgb, front_color, front_color) - result['back_mask'] = cv2.inRange(image_rgb, back_color, back_color) + result['front_mask'] = cv2.inRange(image_rgb, red_lower, red_upper) + result['back_mask'] = cv2.inRange(image_rgb, green_lower, green_upper) result['mask'] = result['front_mask'] + result['back_mask'] else: # 本地查询seg 缓存是否存在