模型是应用yolov8n 分割模型训练：
from ultralytics import YOLO
import torch
# Load a model
model = YOLO('yolov8-seg.yaml') # build a new model from YAML
model = YOLO('yolov8n-seg.pt') # load a pretrained model (recommended for training)
model = YOLO('yolov8-seg.yaml').load('yolov8n.pt') # build from YAML and transfer weights
results = model.train(pose=True, data='./datasets/Dataset_A2C_2025-08-05-1/Dataset_A2C_2025-08-05-1.yaml', epochs=200, imgsz=256)
现想实现预测，如下代码，但报错，请帮我修正：
import numpy as np
import cv2
from numpy.array_api import uint8
from skimage.measure import label, regionprops
from openvino.runtime import Core
from PIL import Image
import torch
import openvino as ov
import matplotlib.pyplot as plt
from sympy.codegen.ast import int32
from triton.language import dtype
class YiAtrium():
def __init__(self, device="CPU"):
self.core = Core()
self.model = self.core.read_model("/Work/zhangxin/ultralytics/runs/segment/train/weights/best_openvino_model/best.xml")
self.compiled_model = self.core.compile_model(self.model, device)
# 修改1：使用Opencv库读取数据，适配原有的代码。
# 修改2：输入cvmat图像数据
def preprocess(self, image):
# img = cv2.imread(img_path,cv2.IMREAD_GRAYSCALE)
# img = cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
img_resized = cv2.resize(img, (256,256))
img_ndarray = img_resized
if img_ndarray.ndim == 2:
img_ndarray = img_ndarray[np.newaxis, ...]
img_ndarray = img_ndarray[np.newaxis, ...]
img_ndarray = img_ndarray / 255.0
tensor = torch.from_numpy(img_ndarray)
img_tensor = tensor.permute(2, 0, 1).float().contiguous() # 或 tensor = tensor.permute(2, 0, 1).contiguous()
img_tensor = img_tensor.unsqueeze(0) # 形状变为：(1, 3, 256, 256)
#img_ndarray = img_ndarray / 255.0
#img_tensor = torch.as_tensor(img_ndarray.copy()).float().contiguous()
return img_tensor
@staticmethod
def keep_largest_region(mask):
if mask.max() == 0:
return mask
labeled_mask = label(mask, connectivity=1)
regions = regionprops(labeled_mask)
if not regions:
return mask
largest_region = max(regions, key=lambda r: r.area)
result_mask = np.zeros_like(mask)
result_mask[labeled_mask == largest_region.label] = 1
return result_mask
@staticmethod
def fill_hole(image):
"""填充图像中的空洞"""
# 确保是二值图像
if len(np.unique(image)) > 2:
_, image = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY)
src = image.copy()
mask = np.zeros([src.shape[0] + 2, src.shape[1] + 2], np.uint8)
# 找到背景点
isbreak = False
for i in range(src.shape[0]):
for j in range(src.shape[1]):
if src[i, j] == 0:
seedpoint = (j, i) # 注意坐标顺序 (x,y)
isbreak = True
break
if isbreak:
break
cv2.floodFill(src, mask, seedpoint, 255)
img_floofill_inv = cv2.bitwise_not(src)
im_out = image | img_floofill_inv
return im_out
@staticmethod
def get_edge_points(mask, scalex, scaley):
contours, _ = cv2.findContours(
mask,
cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_NONE
)
edge_points = []
for contour in contours:
epsilon = 0.005 * cv2.arcLength(contour, True)
approx = cv2.approxPolyDP(contour, epsilon, True)
for point in approx:
x, y = point[0]
up_x = int(x * scaley)
up_y = int(y * scalex)
edge_points.append((up_x, up_y))
final_result = {"detection": {}, "segmentation": {}}
point_dict = {"1-LVIntima_pylogon": edge_points.copy()}
final_result["segmentation"] = point_dict
return final_result
# 原始：输入的是图像的路径
# def Yi_Segment(self, image_path):
# input_tensor = self.preprocess(image_path)
# predict_tensor = self.compiled_model(input_tensor)
# output_ndarray = predict_tensor[0].argmax(1).squeeze(0)
# prediction = output_ndarray.astype(np.uint8)
# prediction = prediction * 255
# output_mask = self.fill_hole(self.keep_largest_region(prediction))
# final_output = self.get_edge_points(output_mask)
# return final_output
# 修改：输入cvmat图像数据
def Yi_Segment(self, image):
input_tensor = self.preprocess(image)
predict_tensor = self.compiled_model(input_tensor)
output_ndarray = predict_tensor[0].argmax(1).squeeze(0)
prediction = output_ndarray.astype(np.uint8)
prediction = prediction * 255
output_mask = self.fill_hole(self.keep_largest_region(prediction))
img_h, img_w = image.shape[:2]
scalex = img_h / 256.0
scaley = img_w / 256.0
final_output = self.get_edge_points(output_mask, scalex, scaley)
#final_output = self.get_edge_points(final_output)
return final_output
if __name__ == "__main__":
image_path = "/Work/zhangxin/ultralytics/runs/test/Heart_A2C_0000003_20240416_Comen_cropped_105.bmp" #
segmenter = YiAtrium()
image = cv2.imread(image_path)
final_result = segmenter.Yi_Segment(image)
img = cv2.imread(image_path)
new_img = np.zeros((image.shape[0], image.shape[1]), dtype=uint8)
cv2.drawContours(new_img, [np.asarray(final_result["segmentation"]["1-LVIntima_pylogon"])], -1, [255, 0], thickness=cv2.FILLED)
cv2.imshow("1", new_img)
cv2.imshow(("2"),img)
cv2.waitKey(0)
扫一扫
714
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
