aidlab/AiDA_Python
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat

Browse Source 
fix   design mask 红绿判断修正
This commit is contained in:
zhouchengrong
2024-09-04 15:05:05 +08:00
parent f5fc6b0c68
commit cfa2cd1987
2 changed files with 8 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
app/service/design/items/pipelines/segmentation.py
View File

@@ -19,22 +19,17 @@ class Segmentation(object):
def __call__(self, result):
if "seg_mask_url" in result.keys() and result['seg_mask_url'] != "":
seg_mask = oss_get_image(bucket=result['seg_mask_url'].split('/')[0], object_name=result['seg_mask_url'][result['seg_mask_url'].find('/') + 1:], data_type="cv2")
seg_mask = cv2.resize(seg_mask, (result['img_shape'][1], result['img_shape'][0]))
seg_mask = cv2.resize(seg_mask, (result['img_shape'][1], result['img_shape'][0]), interpolation=cv2.INTER_NEAREST)
# 转换颜色空间为 RGB（OpenCV 默认是 BGR）
image_rgb = cv2.cvtColor(seg_mask, cv2.COLOR_BGR2RGB)
# 定义红色和绿色的颜色范围
# 红色范围: 下界 [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)
r, g, b = cv2.split(image_rgb)
red_mask = r > g
green_mask = g > r
# 创建红色和绿色掩码
result['front_mask'] = cv2.inRange(image_rgb, red_lower, red_upper)
result['back_mask'] = cv2.inRange(image_rgb, green_lower, green_upper)
result['front_mask'] = np.array(red_mask, dtype=np.uint8) * 255
result['back_mask'] = np.array(green_mask, dtype=np.uint8) * 255
result['mask'] = result['front_mask'] + result['back_mask']
else:
# 本地查询seg 缓存是否存在