一,2d图像mask输出boundingbox
代码实现:输入原图像地址,mask图像地址,输出一个yolov5格式的txt文件,一个不同分类不同颜色的原图像框选结果。
将2d的多分类图像mask输出boundingbox。
import cv2 import matplotlib.pyplot as plt import numpy as np import os import shutil def mask_find_bboxs(mask): retval, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8) # connectivity参数的默认值为8 # print(retval,labels,centroids) stats = stats[stats[:,4].argsort()] return stats[:-1] # 排除最外层的连通图
上面的函数实现了对连通域进行划分,对结果进行输出。想了解搜一下cv2.connectedComponentsWithStats就行。
mask_ori = cv2.imread(r'.label_process_results esults_compareTSE_10_label.png', cv2.COLOR_BGR2GRAY) mask_rgb = cv2.imread(r'.label_process_results esults_compareTSE_10.png') ret, mask = cv2.threshold(mask_ori, 20, 255, cv2.THRESH_BINARY) # 30是最小分类,处理成灰度图才能进mask_find_bboxs函数 bboxs = mask_find_bboxs(mask) # 框选目标的参数 width = len(mask[1, :]) # 宽 height = len(mask[:, 1]) # 高 file = open(r". estyolo.txt", "w", encoding='utf-8') ''' 参数介绍b[0],b[1]:左上角坐标x,y x变大向右平移y变大向下平移 b[2] :矩形宽度 b[3] :矩形高度 转换成yolo。txt格式 类别(0开始) 方框中心点比例(x,y) width宽占比例 height 高占比例 ''' classification = 10 # 10表示出错 for b in bboxs: pixel_catch = 0 for i in range(3, 8): for j in range(3, 8): pixel_catch = max(mask_ori[int((b[1] + (b[3] * (j / 10)))), int((b[0] + (b[2] * (i / 10))))], pixel_catch) yolo_para_width = b[2] / width yolo_para_height = b[3] / height yolo_para_x_center = (b[0] / width) + (yolo_para_width / 2) yolo_para_y_center = (b[1] / height) + (yolo_para_height / 2) print(pixel_catch) if pixel_catch == 40: cv2.rectangle(mask_rgb, (b[0], b[1]), (b[0] + b[2], b[1] + b[3]), (255, 255, 0), 2) classification = 2 if pixel_catch == 80: cv2.rectangle(mask_rgb, (b[0], b[1]), (b[0] + b[2], b[1] + b[3]), (255, 255, 0), 2) classification = 3 if pixel_catch == 120: cv2.rectangle(mask_rgb, (b[0], b[1]), (b[0]+b[2], b[1]+b[3]), (0, 0, 255), 2) classification = 0 if pixel_catch == 160: cv2.rectangle(mask_rgb, (b[0], b[1]), (b[0] + b[2], b[1] + b[3]), (255, 255, 0), 2) classification = 1 if pixel_catch == 200: cv2.rectangle(mask_rgb, (b[0], b[1]), (b[0] + b[2], b[1] + b[3]), (255, 255, 0), 2) classification = 4 if pixel_catch == 240: cv2.rectangle(mask_rgb, (b[0], b[1]), (b[0] + b[2], b[1] + b[3]), (255, 255, 0), 2) classification = 5 file.write(str(classification) + ' ' + str(yolo_para_x_center) + ' ' + str(yolo_para_y_center) + ' ' + str( yolo_para_width) + ' ' + str(yolo_para_height) + ' ') cv2.imwrite('./000_new.jpg', mask_rgb) file.close()
该代码通过输入原图像和mask图像,先把通过上面的函数把每一个连通域找到,再框选出图像,根据框选结果计算出yolo需要的参数进行输出。