多个roi生成二值图:
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)}")