多个roi生成二值图:
python
import cv2
import numpy as np
import json
import os
class ROIDrawer:
def __init__(self, image_o, label="beizi_5"):
self.drawing = False
self.ix, self.iy = -1, -1
self.rois = [] # 存储多个ROI
self.image_o = image_o
self.image = self.image_o.copy()
self.temp_image = self.image.copy()
self.ok = False
self.label = label # 目标标签
def draw_crosshair(self, event, x, y, flags, param):
self.temp_image = self.image.copy() # 每次更新临时图像
# 关键修改:无论是否有已框选的ROI,只要不在绘制中就显示十字星
if not self.drawing: # 只有当不处于拖拽框选状态时,才显示十字星
cv2.line(self.temp_image, (x, 0), (x, self.temp_image.shape[0]), (0, 255, 0), 1)
cv2.line(self.temp_image, (0, y), (self.temp_image.shape[1], y), (0, 255, 0), 1)
# 鼠标按下:开始框选
if event == cv2.EVENT_LBUTTONDOWN:
self.drawing = True
self.ix, self.iy = x, y
# 鼠标移动:实时绘制矩形(此时不显示十字星,因为drawing=True)
elif event == cv2.EVENT_MOUSEMOVE:
if self.drawing:
# 绘制当前正在拖拽的矩形
cv2.rectangle(self.temp_image, (self.ix, self.iy), (x, y), (255, 0, 0), 2)
# 鼠标左键释放:确认当前ROI(继续框选)
elif event == cv2.EVENT_LBUTTONUP:
self.drawing = False # 结束绘制,恢复十字星显示
# 计算规范的坐标(确保x1 < x2, y1 < y2)
x1, y1 = min(self.ix, x), min(self.iy, y)
x2, y2 = max(self.ix, x), max(self.iy, y)
# 绘制最终矩形到原始图像
cv2.rectangle(self.image, (x1, y1), (x2, y2), (255, 0, 0), 2)
# 保存ROI坐标
self.rois.append([[x1, y1], [x2, y2]])
print(f"已添加ROI: {[x1, y1]} - {[x2, y2]} (共{len(self.rois)}个)")
# 鼠标右键释放:完成框选
elif event == cv2.EVENT_RBUTTONUP:
self.drawing = False # 结束绘制,恢复十字星显示
x1, y1 = min(self.ix, x), min(self.iy, y)
x2, y2 = max(self.ix, x), max(self.iy, y)
cv2.rectangle(self.image, (x1, y1), (x2, y2), (255, 0, 0), 2)
self.rois.append([[x1, y1], [x2, y2]])
print(f"已添加ROI: {[x1, y1]} - {[x2, y2]} (共{len(self.rois)}个)")
self.ok = True
def save_json(self, image_path, output_json_path=None):
if not self.rois:
print("没有框选任何目标,不保存JSON")
return
if not output_json_path:
img_dir, img_name = os.path.split(image_path)
img_base = os.path.splitext(img_name)[0]
output_json_path = os.path.join(img_dir, f"{img_base}.json")
# 构建JSON结构
json_data = {
"version": "1.0.0",
"flags": {},
"shapes": [],
"imagePath": image_path,
"imageHeight": self.image_o.shape[0],
"imageWidth": self.image_o.shape[1]
}
for points in self.rois:
shape = {
"label": self.label,
"shape_type": "rectangle",
"points": points,
"description": "",
"flags": {}
}
json_data["shapes"].append(shape)
with open(output_json_path, 'w', encoding='utf-8') as f:
json.dump(json_data, f, ensure_ascii=False, indent=4)
print(f"JSON已保存至: {output_json_path}")
def save_roi_as_black_white_png(self, image_path):
"""将ROI区域保存为黑白PNG图像"""
if not self.rois:
print("没有框选任何目标,不保存PNG")
return
# 创建全黑图像(与原始图像相同大小)
h, w = self.image_o.shape[:2]
bw_image = np.zeros((h, w), dtype=np.uint8) # 单通道黑色图像
# 在黑色图像上绘制白色矩形(ROI区域)
for points in self.rois:
[x1, y1], [x2, y2] = points
# 确保坐标在图像范围内
x1, y1 = max(0, x1), max(0, y1)
x2, y2 = min(w, x2), min(h, y2)
# 绘制白色矩形区域(填充)
bw_image[y1:y2, x1:x2] = 255
# 保存PNG文件
img_dir, img_name = os.path.split(image_path)
img_base = os.path.splitext(img_name)[0]
output_png_path = os.path.join(img_dir, f"{img_base}_roi_mask.png")
# 保存为PNG格式
cv2.imwrite(output_png_path, bw_image)
print(f"黑白PNG掩码已保存至: {output_png_path}")
print(f"图像大小: {w}x{h}, 白色区域数量: {len(self.rois)}")
# 可选:显示保存的图像
self.display_bw_image(bw_image)
return bw_image
def display_bw_image(self, bw_image):
"""显示黑白图像"""
# 放大显示以便观察
scale_factor = 800 / max(bw_image.shape)
display_h = int(bw_image.shape[0] * scale_factor)
display_w = int(bw_image.shape[1] * scale_factor)
display_img = cv2.resize(bw_image, (display_w, display_h), interpolation=cv2.INTER_NEAREST)
# 创建彩色版本用于显示(蓝色表示白色区域)
colored_display = cv2.cvtColor(display_img, cv2.COLOR_GRAY2BGR)
colored_display[display_img == 255] = [255, 0, 0] # 将白色区域显示为蓝色
cv2.imshow('ROI Mask Preview (Blue=White Area)', colored_display)
cv2.waitKey(3000) # 显示3秒
cv2.destroyAllWindows()
def save_roi_as_transparent_png(self, image_path):
"""将ROI区域保存为带透明通道的PNG(白色区域不透明,黑色区域透明)"""
if not self.rois:
print("没有框选任何目标,不保存PNG")
return
# 创建RGBA图像
h, w = self.image_o.shape[:2]
rgba_image = np.zeros((h, w, 4), dtype=np.uint8) # 4通道:B,G,R,A
# 在RGBA图像上绘制白色矩形
for points in self.rois:
[x1, y1], [x2, y2] = points
x1, y1 = max(0, x1), max(0, y1)
x2, y2 = min(w, x2), min(h, y2)
# 设置白色区域为不透明
rgba_image[y1:y2, x1:x2, 0:3] = 255 # BGR通道为白色
rgba_image[y1:y2, x1:x2, 3] = 255 # Alpha通道为255(不透明)
# 黑色区域设置为透明
# (np.zeros已经初始化所有通道为0,包括alpha通道,所以黑色区域是透明的)
# 保存为PNG文件
img_dir, img_name = os.path.split(image_path)
img_base = os.path.splitext(img_name)[0]
output_png_path = os.path.join(img_dir, f"{img_base}_roi_transparent.png")
# 使用cv2保存(注意颜色通道顺序)
cv2.imwrite(output_png_path, rgba_image)
print(f"透明PNG掩码已保存至: {output_png_path}")
return rgba_image
def save_roi_separate_pngs(self, image_path):
"""为每个ROI单独保存为黑白PNG"""
if not self.rois:
print("没有框选任何目标,不保存PNG")
return
h, w = self.image_o.shape[:2]
img_dir, img_name = os.path.split(image_path)
img_base = os.path.splitext(img_name)[0]
for i, points in enumerate(self.rois):
# 创建全黑图像
bw_image = np.zeros((h, w), dtype=np.uint8)
[x1, y1], [x2, y2] = points
x1, y1 = max(0, x1), max(0, y1)
x2, y2 = min(w, x2), min(h, y2)
# 绘制白色矩形
bw_image[y1:y2, x1:x2] = 255
# 保存单个ROI
output_png_path = os.path.join(img_dir, f"{img_base}_roi_{i + 1:02d}.png")
cv2.imwrite(output_png_path, bw_image)
print(f"ROI {i + 1} 已保存至: {output_png_path}")
def run(self, image_path, output_json=None, save_format="black_white"):
"""
运行ROI绘制工具
参数:
image_path: 图像路径
output_json: JSON输出路径(可选)
save_format: 保存格式,可选值:
"black_white" - 黑白PNG(默认)
"transparent" - 透明PNG
"separate" - 每个ROI单独保存
"all" - 保存所有格式
"""
cv2.namedWindow('Draw ROI')
cv2.setMouseCallback('Draw ROI', self.draw_crosshair)
print("=" * 50)
print("ROI绘制工具")
print("=" * 50)
print("操作说明:")
print("1. 左键拖拽框选目标(松开后继续框选下一个)")
print("2. 右键拖拽框选最后一个目标(松开后结束框选)")
print("3. 按Esc键取消操作,按Enter键保存并退出")
print("=" * 50)
while True:
cv2.imshow('Draw ROI', self.temp_image)
key = cv2.waitKey(1) & 0xFF
if key == 27: # Esc键:取消操作
print("已取消操作")
self.rois = []
break
elif key == 13: # Enter键:保存并退出
break
elif self.ok: # 右键结束框选
break
cv2.destroyAllWindows()
if self.rois:
# 保存JSON
self.save_json(image_path, output_json)
# 根据选择的格式保存PNG
if save_format == "black_white" or save_format == "all":
self.save_roi_as_black_white_png(image_path)
if save_format == "transparent" or save_format == "all":
self.save_roi_as_transparent_png(image_path)
if save_format == "separate" or save_format == "all":
self.save_roi_separate_pngs(image_path)
print(f"\n✅ 处理完成!共框选 {len(self.rois)} 个ROI区域")
else:
print("⚠️ 没有框选任何ROI区域")
return self.rois
if __name__ == '__main__':
# 示例用法
image_path = r"D:\project\seg\RobustVideoMatting-master\output_rvm.png"
image_o = cv2.imread(image_path)
if image_o is None:
print(f"无法读取图像: {image_path}")
exit(0)
# 创建ROI绘制器
roi_drawer = ROIDrawer(image_o, label="penzi")
# 运行并保存为黑白PNG(默认)
selected_rois = roi_drawer.run(
image_path=image_path,
output_json="./output.json",
save_format="black_white" # 或 "transparent", "separate", "all"
)
print(f"最终框选的ROI数量: {len(selected_rois)}")