Yolov8 根据需要自定义颜色

1、找到系统文件 plotting.py,加一个类 Color

class Colors:
    # Ultralytics color palette https://ultralytics.com/
    def __init__(self):
        # hex = matplotlib.colors.TABLEAU_COLORS.values()
        hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
                '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
        self.n = len(self.palette)
        self.pose_palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255],
                                      [153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255],
                                      [255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102],
                                      [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255]],
                                     dtype=np.uint8)

    def __call__(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  # rgb order (PIL)
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))

class Color:
    def __init__(self):
        self.red = (255, 0, 0) # 深红色
        self.green = (0, 255, 0) # 绿色

    def __call__(self, i, bgr=False):
        if i == 1:
            return self.red if not bgr else self.red[::-1]
        elif i == 0:
            return self.green if not bgr else self.green[::-1]
        else:
            raise ValueError("i 暂时只能取 0 或 1")

colors = Colors()  # create instance for 'from utils.plots import colors'
color = Color()

2、找到 results.py ,先导入函数

from ultralytics.yolo.utils.plotting import Annotator, colors, save_one_box, color # 我这个在第17行,主要是加一个 color

# 然后是在212行,更改如下
annotator.masks(pred_masks.data, colors=[color(1, True)], im_gpu=img_gpu) # 红色
# 最后是219行,更改如下
annotator.box_label(d.xyxy.squeeze(), label, color=color(0, True)) # 绿色
# 可能由于版本不一样导致行数不一样,需要自己找一下

3、最后运行 model.predict() 即可
posted @ 2023-05-11 08:29  ~逍遥子~  阅读(1677)  评论(1编辑  收藏  举报