yolov8实例分割预测每个物体不同的灰度图

预测代码

点击查看代码

import os
from ultralytics import YOLO
import numpy as np
import cv2
def generate_colors_by_classes_size(classes_size=100):
    # colors = np.zeros((classes_size, 3), dtype=np.uint8)
    colors = np.zeros((classes_size, 3), dtype=np.uint16)
    for i in range(classes_size):
        # color=((i+1)* (255 / classes_size)) % 256
        color=((i+1)* (19999 / classes_size)) % 20000
        colors[i] = color
    return colors
def gen_gray_mask(results,colors,outdir='./out_mask'):
     for result in results:
            input_image = cv2.imread(result.path)
            h,w,c=input_image.shape
            # mask= np.zeros((h, w), dtype=np.uint8) 
            mask= np.zeros((h, w), dtype=np.uint16)   
            labels=result.boxes.cls.tolist()
            # 每张图的mask
            i=0
            for mask_xy in result.masks.xy:
                pixel_xy0 = result.masks.xy[0]
                pixel_xy = mask_xy
                points = np.array(pixel_xy, np.int32)
                if(points.size==0):
                    continue
                # print(points)
                try:
                    # cv2.drawContours(input_image, [points], -1, (0, 255, 0), 2) tuple(arr.tolist())
                    label_index=int(labels[i])
                    clolor=tuple(colors[i].tolist()) #每个实例不同颜色
                    # clolor=tuple(colors[label_index].tolist())#按照类别着色
                    # cv2.drawContours(input_image, [points], -1, clolor, 1)
                    cv2.fillPoly(mask, [points], clolor)
                    i=i+1
                except Exception as e:
                    print(e)
                    break
            basename=os.path.basename(result.path)
            outname=os.path.join(outdir,basename)
            cv2.imwrite(outname, mask)
         
if __name__=='__main__':

        # names=["Primary_Particle_Z","Primary_Particle_X","Primary_Particle_B","Primary_Particle_K","Primary_Particle_C","Primary_Particle5",
        #        "6","7","8","9","10",
        #        "11","12","13","14","15",
        #        "16","17","18","19","20",
        #        "21","22","23","24","25",
        #        "26","27","28","29","30",
        #        "31","32","33","34","35"
        #        ]
        names=["Micropowder","SQ","Primary_Particle_B","Primary_Particle_K","Primary_Particle_C","Primary_Particle5",
               "6","7","8","9","10",
               "11","12","13","14","15",
               "16","17","18","19","20",
               "21","22","23","24","25",
               "26","27","28","29","30",
               "31","32","33","34","35"
               ]

        # colors=generate_random_colors()
        num_class=9999
        colors=generate_colors_by_classes_size(num_class)
        clolor=tuple(colors[0].tolist())
        model = YOLO(r'D:\gzj\PycharmWorkSpace\ultralytics-main\ultralytics-main\yolov8s_wf_big_small4pic.pt')
        results=model.predict(r'D:\gzj\PycharmWorkSpace\ultralytics-main\ultralytics-main\Test\img\001_1.jpg', save=True,save_txt=True, imgsz=640,line_width=2)
        outpath='./out_mask0-1'
        os.makedirs(outpath,exist_ok=True)
        gen_gray_mask(results=results,colors=colors,outdir=outpath) 
        print("done")

原图

灰度图