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)
帮我总结此代码实现的内容
理解ndarray.shape属性：图像维度解析
ndarray.shape属性返回图像的维度信息，包括行数、列数和通道数。例如，对于一个n行m列、k通道的图像，shape会返回(n, m, k)。当图像有四个通道，如包含alpha通道时，第三个数值会变为4。这有助于理解图像数据的结构。
扫一扫
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
