预测代码
点击查看代码
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")
灰度图