【记录】基于opencv和python+【根据HSV自动选择区分度的颜色】+绘制矩形框+可视化数据集

完成功能：

• 根据HSV，对不同类别自动选亮度足够且区分度明显的颜色
• 绘制矩形框和显示类别
• 有图例

（其中HSV转RGB部分参考：https://blog.csdn.net/rickwy/article/details/86501181）

1、定义DrawLabels类

import cv2
import numpy as np
import math

class DrawLabels:
    def __init__(self,CLASSES):
        self.CLASSES=CLASSES
        self.CLASSES_COLORS={
            CLASSES[j]:self.hsv2rgb(i/len(CLASSES)*360,1.0,1.0) for i,j in enumerate(CLASSES.keys())
        }
        self.font = cv2.FONT_HERSHEY_SIMPLEX
    @staticmethod
    def hsv2rgb(h, s, v):
        h = float(h)
        s = float(s)
        v = float(v)
        h60 = h / 60.0
        h60f = math.floor(h60)
        hi = int(h60f) % 6
        f = h60 - h60f
        p = v * (1 - s)
        q = v * (1 - f * s)
        t = v * (1 - (1 - f) * s)
        r, g, b = 0, 0, 0
        if hi == 0: r, g, b = v, t, p
        elif hi == 1: r, g, b = q, v, p
        elif hi == 2: r, g, b = p, v, t
        elif hi == 3: r, g, b = p, q, v
        elif hi == 4: r, g, b = t, p, v
        elif hi == 5: r, g, b = v, p, q
        r, g, b = int(r * 255), int(g * 255), int(b * 255)
        return r, g, b
    def draw(self,img,anns):
        # anns=[
        #     [label,x_center,y_center,w,h],
        #     ...,
        #     [label,x_center,y_center,w,h],
        # ]
        img=img.copy()
        for ann in anns:
            label, x, y, w, h = ann
            img = cv2.rectangle(img, (int(x-w/2),int(y-h/2)), (int(x+w/2),int(y+h/2)),
                                self.CLASSES_COLORS[label],2)
            img = cv2.putText(img, str(label), (int(x), int(y)), 
                              self.font, 1.2, self.CLASSES_COLORS[label], 2)
        img2 = np.zeros((img.shape[0]+20*(len(self.CLASSES.keys())+1),
                         img.shape[1],img.shape[2]))+128
        img2[:img.shape[0],:img.shape[1],:img.shape[2]]=img
        for i,classname in enumerate(self.CLASSES.keys()):
            img2 = cv2.putText(img2, str(self.CLASSES[classname])+":"+classname, 
                               (0, img.shape[0]+(i+1)*20), 
                self.font, 0.8, self.CLASSES_COLORS[self.CLASSES[classname]], 2)
        return img2

2、如何使用DrawLabels类

（1）初始化

CLASSES={
    "trash_over":0,
    "trash_no_full":1,
    "garbage":2,
}
drawlabels = DrawLabels(CLASSES)

（2）绘制矩形框

img = cv2.imread("/home/zjj/workspace/cvmart_topList/trash/dataset/206/garbage_overflow_171.jpg")
anns=[
    [1, 20.0, 40.0, 40.0, 40.0],
    [0, 30.0, 60.0, 20.0, 30.0],
]
img2 = drawlabels.draw(img,anns)

（3）直接绘制任意框的示例

3、一个可视化yolov5标注的完整示例

• 以yolov5格式的标注文件为例，完整代码为

import cv2
import numpy as np
import math

class DrawLabels:
    def __init__(self,CLASSES):
        self.CLASSES=CLASSES
        self.CLASSES_COLORS={
            CLASSES[j]:self.hsv2rgb(i/len(CLASSES)*360,1.0,1.0) for i,j in enumerate(CLASSES.keys())
        }
        self.font = cv2.FONT_HERSHEY_SIMPLEX
    @staticmethod
    def hsv2rgb(h, s, v):
        h = float(h)
        s = float(s)
        v = float(v)
        h60 = h / 60.0
        h60f = math.floor(h60)
        hi = int(h60f) % 6
        f = h60 - h60f
        p = v * (1 - s)
        q = v * (1 - f * s)
        t = v * (1 - (1 - f) * s)
        r, g, b = 0, 0, 0
        if hi == 0: r, g, b = v, t, p
        elif hi == 1: r, g, b = q, v, p
        elif hi == 2: r, g, b = p, v, t
        elif hi == 3: r, g, b = p, q, v
        elif hi == 4: r, g, b = t, p, v
        elif hi == 5: r, g, b = v, p, q
        r, g, b = int(r * 255), int(g * 255), int(b * 255)
        return r, g, b
    def draw(self,img,anns):
        # anns=[
        #     [label,x_center,y_center,w,h],
        #     ...,
        #     [label,x_center,y_center,w,h],
        # ]
        img=img.copy()
        for ann in anns:
            label, x, y, w, h = ann
            img = cv2.rectangle(img, (int(x-w/2),int(y-h/2)), (int(x+w/2),int(y+h/2)),
                                self.CLASSES_COLORS[label],2)
            img = cv2.putText(img, str(label), (int(x), int(y)), 
                              self.font, 1.2, self.CLASSES_COLORS[label], 2)
        img2 = np.zeros((img.shape[0]+20*(len(self.CLASSES.keys())+1),
                         img.shape[1],img.shape[2]))+128
        img2[:img.shape[0],:img.shape[1],:img.shape[2]]=img
        for i,classname in enumerate(self.CLASSES.keys()):
            img2 = cv2.putText(img2, str(self.CLASSES[classname])+":"+classname, 
                               (0, img.shape[0]+(i+1)*20), 
                self.font, 0.8, self.CLASSES_COLORS[self.CLASSES[classname]], 2)
        return img2
def main():
    CLASSES={
        "trash_over":0,
        "trash_no_full":1,
        "garbage":2,
    }
    drawlabels = DrawLabels(CLASSES)
    img=cv2.imread("/home/zjj/workspace/cvmart_topList/"
                   "trash/dataset/206/garbage_overflow_171.jpg")
    txtpath = "/home/zjj/workspace/cvmart_topList/trash/dataset/206/garbage_overflow_171.txt"
    save_path="/home/zjj/workspace/cvmart_topList/trash/dataset/show/"
    height,weight,c = img.shape
    with open(txtpath) as f:
        anns = f.readlines()
        anns2 = []
        for ann in anns:
            label, x, y, w, h = [float(i) for i in ann.split(" ")]
            label = int(label)
            x *= weight
            w *= weight
            y *= height
            h *= height
            anns2.append([label, x, y, w, h])
    img2 = drawlabels.draw(img,anns2)
    cv2.imwrite(save_path+"garbage_overflow_171.jpg",img2)
if __name__=="__main__":
    main()def main():
    CLASSES={
        "trash_over":0,
        "trash_no_full":1,
        "garbage":2,
    }
    drawlabels = DrawLabels(CLASSES)
    img=cv2.imread("/home/zjj/workspace/cvmart_topList/"
                   "trash/dataset/206/garbage_overflow_171.jpg")
    txtpath = "/home/zjj/workspace/cvmart_topList/trash/dataset/206/garbage_overflow_171.txt"
    save_path="/home/zjj/workspace/cvmart_topList/trash/dataset/show/"
    height,weight,c = img.shape
    with open(txtpath) as f:
        anns = f.readlines()
        anns2 = []
        for ann in anns:
            label, x, y, w, h = [float(i) for i in ann.split(" ")]
            label = int(label)
            x *= weight
            w *= weight
            y *= height
            h *= height
            anns2.append([label, x, y, w, h])
    img2 = drawlabels.draw(img,anns2)
    cv2.imwrite(save_path+"garbage_overflow_171.jpg",img2)
if __name__=="__main__":
    main()