xiongyi/AutoControlSystem-git
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

UPDATE openvino

Browse Source
This commit is contained in:
HJW
2024-09-11 13:50:53 +08:00
parent e58417c806
commit 7fab3b3ffa
8 changed files with 10074 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files

333
Vision/camera_coordinate_dete.py
View File

@ -1,11 +1,11 @@
#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：AutoControlSystem-master
@Project ：AutoControlSystem-master
@File ：camera_coordinate_dete.py
@IDE ：PyCharm
@IDE ：PyCharm
@Author ：hjw
@Date ：2024/8/27 14:24
@Date ：2024/8/27 14:24
'''
import numpy as np
@ -13,189 +13,202 @@ import cv2
import open3d as o3d
from Vision.tool.CameraRVC import camera
from Vision.yolo.yolov8_pt_seg import yolov8_segment
from Vision.yolo.yolov8_openvino import yolov8_segment_openvino
from Vision.tool.utils import find_position
from Vision.tool.utils import class_names
import os
class Detection:
def __init__(self):
model_path = './pt_model/best.pt'
device = 'cpu'
self.camera_rvc = camera()
self.model = yolov8_segment()
self.model.load_model(model_path, device)
# if self.camera_rvc.caminit_isok == True:
# print("RVC X Camera is opened!")
# else:
# print("RVC X Camera is not opened!")
def __init__(self, use_openvino_model = False):
print(os.getcwd())
self.use_openvino_model = use_openvino_model
if self.use_openvino_model == False:
model_path = ''.join([os.getcwd(), '/model/pt/best.pt'])
device = 'cpu'
self.model = yolov8_segment()
self.model.load_model(model_path, device)
else:
model_path = ''.join([os.getcwd(), '/model/openvino/bset.xml'])
device = 'CPU'
self.model = yolov8_segment_openvino(model_path, device, conf_thres=0.3, iou_thres=0.3)
def get_position(self, Point_isVision=False, first_depth=False):
img_path = ''.join([os.getcwd(), '/model/data/test0910.png'])
point_path = ''.join([os.getcwd(), '/model/data/test0910.xyz'])
self.img = cv2.imread(img_path)
self.point = np.loadtxt(point_path).reshape((1080, 1440, 3))
def get_position(self, Point_isVision=True):
""
'''
:param api:
Point_isVision： 点云可视化
first_depth： 按Z轴深度返回坐标及向量
:param api: None
:return: ret , img, (x,y,z), (nx, ny, nz)
'''
ret, img, pm = self.camera_rvc.get_img_and_point_map() # 拍照，获取图像及
if self.camera_rvc.caminit_isok == True:
if ret == 1:
#ret, img, pm = self.camera_rvc.get_img_and_point_map() # 拍照，获取图像及
ret = 1
img = self.img
pm = self.point
if ret == 1:
if self.use_openvino_model == False:
flag, det_cpu, dst_img, masks, category_names = self.model.model_inference(img, 0)
if flag == 1:
xyz = []
nx_ny_nz = []
RegionalArea = []
Depth_Z = []
uv = []
seg_point = []
if Point_isVision==True:
pm2 = pm.copy()
pm2 = pm2.reshape(-1, 3)
pm2 = pm2[~np.isnan(pm2).all(axis=-1), :]
pm2[:, 2] = pm2[:, 2] + 0.25
pcd2 = o3d.geometry.PointCloud()
pcd2.points = o3d.utility.Vector3dVector(pm2)
# o3d.visualization.draw_geometries([pcd2])
else:
flag, det_cpu, scores, masks, category_names = self.model.segment_objects(img)
dst_img = img
if flag == 1:
xyz = []
nx_ny_nz = []
RegionalArea = []
Depth_Z = []
uv = []
seg_point = []
if Point_isVision == True:
pm2 = pm.copy()
pm2 = pm2.reshape(-1, 3)
pm2 = pm2[~np.isnan(pm2).all(axis=-1), :]
pm2[:, 2] = pm2[:, 2] + 0.25
pcd2 = o3d.geometry.PointCloud()
pcd2.points = o3d.utility.Vector3dVector(pm2)
# o3d.visualization.draw_geometries([pcd2])
for i, item in enumerate(det_cpu):
for i, item in enumerate(det_cpu):
# 画box
box_x1, box_y1, box_x2, box_y2 = item[0:4].astype(np.int32)
# 画box
box_x1, box_y1, box_x2, box_y2 = item[0:4].astype(np.int32)
if self.use_openvino_model == False:
label = category_names[int(item[5])]
rand_color = (0, 255, 255)
score = item[4]
org = (int((box_x1 + box_x2) / 2), int((box_y1 + box_y2) / 2))
x_center = int((box_x1 + box_x2) / 2)
y_center = int((box_y1 + box_y2) / 2)
text = '{}|{:.2f}'.format(label, score)
cv2.putText(img, text, org=org, fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.8,
color=rand_color,
thickness=2)
# 画mask
# mask = masks[i].cpu().numpy().astype(int)
else:
label = class_names[int(item[4])]
rand_color = (0, 255, 255)
score = item[4]
org = (int((box_x1 + box_x2) / 2), int((box_y1 + box_y2) / 2))
x_center = int((box_x1 + box_x2) / 2)
y_center = int((box_y1 + box_y2) / 2)
text = '{}|{:.2f}'.format(label, score)
cv2.putText(img, text, org=org, fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.8,
color=rand_color,
thickness=2)
# 画mask
# mask = masks[i].cpu().numpy().astype(int)
if self.use_openvino_model == False:
mask = masks[i].cpu().data.numpy().astype(int)
# mask = masks[i].numpy().astype(int)
bbox_image = img[box_y1:box_y2, box_x1:box_x2]
h, w = box_y2 - box_y1, box_x2 - box_x1
mask_colored = np.zeros((h, w, 3), dtype=np.uint8)
mask_colored[np.where(mask)] = rand_color
##################################
imgray = cv2.cvtColor(mask_colored, cv2.COLOR_BGR2GRAY)
# cv2.imshow('mask',imgray)
# cv2.waitKey(1)
# 2、二进制图像
ret, binary = cv2.threshold(imgray, 10, 255, 0)
# 阈值 二进制图像
# cv2.imshow('bin',binary)
# cv2.waitKey(1)
contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
# all_piont_list = contours_in(contours)
# print(len(all_piont_list))
max_contour = None
max_perimeter = 0
for contour in contours: # 排除小分割区域或干扰区域
perimeter = cv2.arcLength(contour, True)
if perimeter > max_perimeter:
max_perimeter = perimeter
max_contour = contour
'''
提取区域范围内的（x, y）
'''
mask_inside = np.zeros(binary.shape, np.uint8)
cv2.drawContours(mask_inside, [max_contour], 0, 255, -1)
pixel_point2 = cv2.findNonZero(mask_inside)
# result = np.zeros_like(color_image)
select_point = []
for i in range(pixel_point2.shape[0]):
select_point.append(pm[pixel_point2[i][0][1]+box_y1, pixel_point2[i][0][0]+box_x1])
select_point = np.array(select_point)
pm_seg = select_point.reshape(-1, 3)
pm_seg = pm_seg[~np.isnan(pm_seg).all(axis=-1), :] # 剔除 nan
# cv2.imshow('result', piont_result)
if pm_seg.size < 300:
print("分割出的点云数量较少，无法拟合平面！")
continue
else:
mask = masks[i].astype(int)
mask = mask[box_y1:box_y2, box_x1:box_x2]
'''
拟合平面，计算法向量
'''
# mask = masks[i].numpy().astype(int)
h, w = box_y2 - box_y1, box_x2 - box_x1
mask_colored = np.zeros((h, w, 3), dtype=np.uint8)
mask_colored[np.where(mask)] = rand_color
##################################
imgray = cv2.cvtColor(mask_colored, cv2.COLOR_BGR2GRAY)
# cv2.imshow('mask',imgray)
# cv2.waitKey(1)
# 2、二进制图像
ret, binary = cv2.threshold(imgray, 10, 255, 0)
# 阈值 二进制图像
# cv2.imshow('bin',binary)
# cv2.waitKey(1)
contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
# all_piont_list = contours_in(contours)
# print(len(all_piont_list))
max_contour = None
max_perimeter = 0
for contour in contours: # 排除小分割区域或干扰区域
perimeter = cv2.arcLength(contour, True)
if perimeter > max_perimeter:
max_perimeter = perimeter
max_contour = contour
'''
提取区域范围内的（x, y）
'''
mask_inside = np.zeros(binary.shape, np.uint8)
cv2.drawContours(mask_inside, [max_contour], 0, 255, -1)
pixel_point2 = cv2.findNonZero(mask_inside)
# result = np.zeros_like(color_image)
select_point = []
for i in range(pixel_point2.shape[0]):
select_point.append(pm[pixel_point2[i][0][1] + box_y1, pixel_point2[i][0][0] + box_x1])
select_point = np.array(select_point)
pm_seg = select_point.reshape(-1, 3)
pm_seg = pm_seg[~np.isnan(pm_seg).all(axis=-1), :] # 剔除 nan
# cv2.imshow('result', piont_result)
'''
拟合平面，计算法向量
'''
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(pm_seg)
plane_model, inliers = pcd.segment_plane(distance_threshold=0.01,
ransac_n=5,
num_iterations=5000)
[a, b, c, d] = plane_model
# print(f"Plane equation: {a:.2f}x + {b:.2f}y + {c:.2f}z + {d:.2f} = 0")
# inlier_cloud = pcd.select_by_index(inliers) # 点云可视化
# inlier_cloud.paint_uniform_color([1.0, 0, 0])
# outlier_cloud = pcd.select_by_index(inliers, invert=True)
# outlier_cloud.paint_uniform_color([0, 1, 0])
# o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud])
'''
拟合最小外接矩形，计算矩形中心
'''
rect = cv2.minAreaRect(max_contour)
# cv2.boxPoints可以将轮廓点转换为四个角点坐标
box = cv2.boxPoints(rect)
# 这一步不影响后面的画图，但是可以保证四个角点坐标为顺时针
startidx = box.sum(axis=1).argmin()
box = np.roll(box, 4 - startidx, 0)
# 在原图上画出预测的外接矩形
box = box.reshape((-1, 1, 2)).astype(np.int32)
box = box + [[[box_x1, box_y1]], [[box_x1, box_y1]], [[box_x1, box_y1]], [[box_x1, box_y1]]]
x_rotation_center = int((box[0][0][0] + box[1][0][0] + box[2][0][0] + box[3][0][0]) / 4)
y_rotation_center = int((box[0][0][1] + box[1][0][1] + box[2][0][1] + box[3][0][1]) / 4)
point_x, point_y, point_z = pm[y_rotation_center, x_rotation_center]
cv2.circle(img, (x_rotation_center, y_rotation_center), 4, (255, 255, 255), 5) # 标出中心点
if np.isnan(point_x): # 点云值为无效值
continue
else:
xyz.append([point_x * 1000, point_y * 1000, point_z * 1000])
Depth_Z.append(point_z * 1000)
nx_ny_nz.append([a, b, c])
RegionalArea.append(cv2.contourArea(max_contour))
uv.append([x_rotation_center, y_rotation_center])
seg_point.append(pm_seg)
cv2.polylines(img, [box], True, (0, 255, 0), 2)
_idx = find_position(Depth_Z, RegionalArea, 100, True)
if _idx == None:
return 1, img, None, None
else:
cv2.circle(img, (uv[_idx][0], uv[_idx][1]), 30, (0, 0, 255), 20) # 标出中心点
if Point_isVision == True:
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(pm_seg)
pcd.points = o3d.utility.Vector3dVector(seg_point[_idx])
plane_model, inliers = pcd.segment_plane(distance_threshold=0.01,
ransac_n=5,
num_iterations=5000)
[a, b, c, d] = plane_model
#print(f"Plane equation: {a:.2f}x + {b:.2f}y + {c:.2f}z + {d:.2f} = 0")
# inlier_cloud = pcd.select_by_index(inliers) # 点云可视化
# inlier_cloud.paint_uniform_color([1.0, 0, 0])
# outlier_cloud = pcd.select_by_index(inliers, invert=True)
# outlier_cloud.paint_uniform_color([0, 1, 0])
# o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud])
'''
拟合最小外接矩形，计算矩形中心
'''
rect = cv2.minAreaRect(max_contour)
# cv2.boxPoints可以将轮廓点转换为四个角点坐标
box = cv2.boxPoints(rect)
# 这一步不影响后面的画图，但是可以保证四个角点坐标为顺时针
startidx = box.sum(axis=1).argmin()
box = np.roll(box, 4 - startidx, 0)
# 在原图上画出预测的外接矩形
box = box.reshape((-1, 1, 2)).astype(np.int32)
box = box + [[[box_x1, box_y1]], [[box_x1, box_y1]], [[box_x1, box_y1]], [[box_x1, box_y1]]]
x_rotation_center = int((box[0][0][0] + box[1][0][0] + box[2][0][0] + box[3][0][0]) / 4)
y_rotation_center = int((box[0][0][1] + box[1][0][1] + box[2][0][1] + box[3][0][1]) / 4)
point_x, point_y, point_z = pm[y_rotation_center, x_rotation_center]
cv2.circle(img, (x_rotation_center, y_rotation_center), 4, (255, 255, 255), 5) # 标出中心点
if np.isnan(point_x): # 点云值为无效值
continue
else:
xyz.append([point_x*1000, point_y*1000, point_z*1000])
Depth_Z.append(point_z*1000)
nx_ny_nz.append([a, b, c])
RegionalArea.append(cv2.contourArea(max_contour))
uv.append([x_rotation_center, y_rotation_center])
seg_point.append(pm_seg)
cv2.polylines(img, [box], True, (0, 255, 0), 2)
_idx = find_position(Depth_Z, RegionalArea, 100, first_depth)
if _idx == None:
return 1, img, None, None
else:
cv2.circle(img, (uv[_idx][0], uv[_idx][1]), 30, (255, 0, 0), 20) # 标出中心点
if Point_isVision==True:
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(seg_point[_idx])
plane_model, inliers = pcd.segment_plane(distance_threshold=0.01,
ransac_n=5,
num_iterations=5000)
inlier_cloud = pcd.select_by_index(inliers) # 点云可视化
inlier_cloud.paint_uniform_color([1.0, 0, 0])
outlier_cloud = pcd.select_by_index(inliers, invert=True)
outlier_cloud.paint_uniform_color([0, 1, 0])
o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud, pcd2])
return 1, img, xyz[_idx], nx_ny_nz[_idx]
else:
return 1, img, None, None
inlier_cloud = pcd.select_by_index(inliers) # 点云可视化
inlier_cloud.paint_uniform_color([1.0, 0, 0])
outlier_cloud = pcd.select_by_index(inliers, invert=True)
outlier_cloud.paint_uniform_color([0, 1, 0])
o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud, pcd2])
return 1, img, xyz[_idx], nx_ny_nz[_idx]
else:
print("RVC X Camera capture failed!")
return 0, None, None, None
return 1, img, None, None
else:
print("RVC X Camera is not opened!")
print("RVC X Camera capture failed!")
return 0, None, None, None
def release(self):
self.camera_rvc.release()
self.model.clear()

60
Vision/camera_coordinate_dete_img.py
View File

@ -1,11 +1,11 @@
#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：AutoControlSystem-master
@Project ：AutoControlSystem-master
@File ：camera_coordinate_dete.py
@IDE ：PyCharm
@IDE ：PyCharm
@Author ：hjw
@Date ：2024/8/27 14:24
@Date ：2024/8/27 14:24
'''
import numpy as np
@ -13,17 +13,27 @@ import cv2
import open3d as o3d
from Vision.tool.CameraRVC import camera
from Vision.yolo.yolov8_pt_seg import yolov8_segment
from Vision.yolo.yolov8_openvino import yolov8_segment_openvino
from Vision.tool.utils import find_position
from Vision.tool.utils import class_names
import os
class Detection:
def __init__(self):
model_path = './pt_model/best.pt'
device = 'cpu'
img_path = './pt_model/test0910.png'
point_path = './pt_model/test0910.xyz'
self.model = yolov8_segment()
self.model.load_model(model_path, device)
def __init__(self, use_openvino_model = False):
self.use_openvino_model = use_openvino_model
if self.use_openvino_model == False:
model_path = ''.join([os.getcwd(), '/model/pt/best.pt'])
device = 'cpu'
self.model = yolov8_segment()
self.model.load_model(model_path, device)
else:
model_path = ''.join([os.getcwd(), '/model/openvino/best.xml'])
device = 'CPU'
self.model = yolov8_segment_openvino(model_path, device, conf_thres=0.3, iou_thres=0.3)
img_path = ''.join([os.getcwd(), '/model/data/test0910.png'])
point_path = ''.join([os.getcwd(), '/model/data/test0910.xyz'])
self.img = cv2.imread(img_path)
self.point = np.loadtxt(point_path).reshape((1080, 1440, 3))
@ -39,9 +49,12 @@ class Detection:
ret = 1
img = self.img
pm = self.point
if ret == 1:
flag, det_cpu, dst_img, masks, category_names = self.model.model_inference(img, 0)
if self.use_openvino_model == False:
flag, det_cpu, dst_img, masks, category_names = self.model.model_inference(img, 0)
else:
flag, det_cpu, scores, masks, category_names = self.model.segment_objects(img)
dst_img = img
if flag == 1:
xyz = []
nx_ny_nz = []
@ -62,7 +75,10 @@ class Detection:
# 画box
box_x1, box_y1, box_x2, box_y2 = item[0:4].astype(np.int32)
label = category_names[int(item[5])]
if self.use_openvino_model == False:
label = category_names[int(item[5])]
else:
label = class_names[int(item[4])]
rand_color = (0, 255, 255)
score = item[4]
org = (int((box_x1 + box_x2) / 2), int((box_y1 + box_y2) / 2))
@ -74,9 +90,13 @@ class Detection:
thickness=2)
# 画mask
# mask = masks[i].cpu().numpy().astype(int)
mask = masks[i].cpu().data.numpy().astype(int)
if self.use_openvino_model == False:
mask = masks[i].cpu().data.numpy().astype(int)
else:
mask = masks[i].astype(int)
mask = mask[box_y1:box_y2, box_x1:box_x2]
# mask = masks[i].numpy().astype(int)
bbox_image = img[box_y1:box_y2, box_x1:box_x2]
h, w = box_y2 - box_y1, box_x2 - box_x1
mask_colored = np.zeros((h, w, 3), dtype=np.uint8)
mask_colored[np.where(mask)] = rand_color
@ -113,9 +133,7 @@ class Detection:
pm_seg = select_point.reshape(-1, 3)
pm_seg = pm_seg[~np.isnan(pm_seg).all(axis=-1), :] # 剔除 nan
# cv2.imshow('result', piont_result)
if pm_seg.size < 300:
print("分割出的点云数量较少，无法拟合平面！")
continue
'''
拟合平面，计算法向量
'''
@ -167,7 +185,7 @@ class Detection:
if _idx == None:
return 1, img, None, None
else:
cv2.circle(img, (uv[_idx][0], uv[_idx][1]), 30, (255, 0, 0), 20) # 标出中心点
cv2.circle(img, (uv[_idx][0], uv[_idx][1]), 30, (0, 0, 255), 20) # 标出中心点
if Point_isVision == True:
pcd = o3d.geometry.PointCloud()
@ -185,6 +203,10 @@ class Detection:
else:
return 1, img, None, None
else:
print("RVC X Camera capture failed!")
return 0, None, None, None
def release(self):
self.model.clear()

BIN
Vision/model/openvino/best.bin Normal file
View File

Binary file not shown.

9457
Vision/model/openvino/best.xml Normal file
View File

File diff suppressed because it is too large Load Diff

14
Vision/model/openvino/metadata.yaml Normal file
View File

@ -0,0 +1,14 @@
description: Ultralytics YOLOv8n-seg model trained on D:\work\ultralytics-main\ultralytics\cfg\datasets\coco8-seg-dy.yaml
author: Ultralytics
date: '2024-09-10T13:48:59.901222'
version: 8.2.86
license: AGPL-3.0 License (https://ultralytics.com/license)
docs: https://docs.ultralytics.com
stride: 32
task: segment
batch: 1
imgsz:
- 640
- 640
names:
0: box

0
Vision/pt_model/best.pt → Vision/model/pt/best.pt
View File

160
Vision/tool/utils.py
View File

@ -1,13 +1,16 @@
#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：AutoControlSystem-master
@Project ：AutoControlSystem-master
@File ：utils.py
@IDE ：PyCharm
@IDE ：PyCharm
@Author ：hjw
@Date ：2024/8/29 15:07
@Date ：2024/8/29 15:07
'''
import numpy as np
import cv2
def find_position(Depth_Z, RegionalArea, RegionalArea_Threshold, first_depth=True):
if first_depth == True:
sorted_id = sorted(range(len(Depth_Z)), key=lambda k: Depth_Z[k], reverse=False)
@ -25,6 +28,157 @@ def find_position(Depth_Z, RegionalArea, RegionalArea_Threshold, first_depth=Tru
if RegionalArea1[i] > RegionalArea_Threshold:
return sorted_id[i]
class_names = ['box', 'other']
# Create a list of colors for each class where each color is a tuple of 3 integer values
rng = np.random.default_rng(3)
colors = rng.uniform(0, 255, size=(len(class_names), 3))
def nms(boxes, scores, iou_threshold):
# Sort by score
sorted_indices = np.argsort(scores)[::-1]
keep_boxes = []
while sorted_indices.size > 0:
# Pick the last box
box_id = sorted_indices[0]
keep_boxes.append(box_id)
# Compute IoU of the picked box with the rest
ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
# Remove boxes with IoU over the threshold
keep_indices = np.where(ious < iou_threshold)[0]
# print(keep_indices.shape, sorted_indices.shape)
sorted_indices = sorted_indices[keep_indices + 1]
return keep_boxes
def compute_iou(box, boxes):
# Compute xmin, ymin, xmax, ymax for both boxes
xmin = np.maximum(box[0], boxes[:, 0])
ymin = np.maximum(box[1], boxes[:, 1])
xmax = np.minimum(box[2], boxes[:, 2])
ymax = np.minimum(box[3], boxes[:, 3])
# Compute intersection area
intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)
# Compute union area
box_area = (box[2] - box[0]) * (box[3] - box[1])
boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
union_area = box_area + boxes_area - intersection_area
# Compute IoU
iou = intersection_area / union_area
return iou
def xywh2xyxy(x):
# Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2)
y = np.copy(x)
y[..., 0] = x[..., 0] - x[..., 2] / 2
y[..., 1] = x[..., 1] - x[..., 3] / 2
y[..., 2] = x[..., 0] + x[..., 2] / 2
y[..., 3] = x[..., 1] + x[..., 3] / 2
return y
def sigmoid(x):
return 1 / (1 + np.exp(-x))
def draw_detections(image, boxes, scores, class_ids, mask_alpha=0.3, mask_maps=None):
img_height, img_width = image.shape[:2]
size = min([img_height, img_width]) * 0.0006
text_thickness = int(min([img_height, img_width]) * 0.001)
mask_img = draw_masks(image, boxes, class_ids, mask_alpha, mask_maps)
# Draw bounding boxes and labels of detections
for box, score, class_id in zip(boxes, scores, class_ids):
color = colors[class_id]
x1, y1, x2, y2 = box.astype(int)
# Draw rectangle
cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, 2)
label = class_names[class_id]
caption = f'{label} {int(score * 100)}%'
(tw, th), _ = cv2.getTextSize(text=caption, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=size, thickness=text_thickness)
th = int(th * 1.2)
cv2.rectangle(mask_img, (x1, y1),
(x1 + tw, y1 - th), color, -1)
cv2.putText(mask_img, caption, (x1, y1),
cv2.FONT_HERSHEY_SIMPLEX, size, (255, 255, 255), text_thickness, cv2.LINE_AA)
return mask_img
def draw_masks(image, boxes, class_ids, mask_alpha=0.3, mask_maps=None):
mask_img = image.copy()
# Draw bounding boxes and labels of detections
for i, (box, class_id) in enumerate(zip(boxes, class_ids)):
color = colors[class_id]
x1, y1, x2, y2 = box.astype(int)
# Draw fill mask image
if mask_maps is None:
cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, -1)
else:
crop_mask = mask_maps[i][y1:y2, x1:x2, np.newaxis]
crop_mask_img = mask_img[y1:y2, x1:x2]
crop_mask_img = crop_mask_img * (1 - crop_mask) + crop_mask * color
mask_img[y1:y2, x1:x2] = crop_mask_img
return cv2.addWeighted(mask_img, mask_alpha, image, 1 - mask_alpha, 0)
def draw_comparison(img1, img2, name1, name2, fontsize=2.6, text_thickness=3):
(tw, th), _ = cv2.getTextSize(text=name1, fontFace=cv2.FONT_HERSHEY_DUPLEX,
fontScale=fontsize, thickness=text_thickness)
x1 = img1.shape[1] // 3
y1 = th
offset = th // 5
cv2.rectangle(img1, (x1 - offset * 2, y1 + offset),
(x1 + tw + offset * 2, y1 - th - offset), (0, 115, 255), -1)
cv2.putText(img1, name1,
(x1, y1),
cv2.FONT_HERSHEY_DUPLEX, fontsize,
(255, 255, 255), text_thickness)
(tw, th), _ = cv2.getTextSize(text=name2, fontFace=cv2.FONT_HERSHEY_DUPLEX,
fontScale=fontsize, thickness=text_thickness)
x1 = img2.shape[1] // 3
y1 = th
offset = th // 5
cv2.rectangle(img2, (x1 - offset * 2, y1 + offset),
(x1 + tw + offset * 2, y1 - th - offset), (94, 23, 235), -1)
cv2.putText(img2, name2,
(x1, y1),
cv2.FONT_HERSHEY_DUPLEX, fontsize,
(255, 255, 255), text_thickness)
combined_img = cv2.hconcat([img1, img2])
if combined_img.shape[1] > 3840:
combined_img = cv2.resize(combined_img, (3840, 2160))
return combined_img

232
Vision/yolo/yolov8_openvino.py Normal file
View File

@ -0,0 +1,232 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
# @Time : 2024/9/10 16:33
# @Author : hjw
# @File : yolov8_openvino.py
'''
import math
import time
import cv2
import numpy as np
from openvino.runtime import Core
from Vision.tool.utils import xywh2xyxy, nms, draw_detections, sigmoid
class yolov8_segment_openvino:
def __init__(self, path, device_name="CPU", conf_thres=0.7, iou_thres=0.5, num_masks=32):
self.conf_threshold = conf_thres
self.iou_threshold = iou_thres
self.num_masks = num_masks
# Initialize model
self.initialize_model(path, device_name)
def __call__(self, image):
return self.segment_objects(image)
def initialize_model(self, path, device_name):
self.core = Core()
self.net = self.core.compile_model(path, device_name=device_name)
self.ir = self.net.create_infer_request()
input_shape = self.net.inputs[0].shape
self.input_height = input_shape[2]
self.input_width = input_shape[3]
def segment_objects(self, image):
input_tensor = self.prepare_input(image)
# Perform inference on the image
outputs = self.ir.infer(input_tensor)
self.boxes, self.scores, self.class_ids, mask_pred = self.process_box_output(outputs[0])
self.mask_maps = self.process_mask_output(mask_pred, outputs[1])
self.boxes = self.boxes.tolist()
self.scores = self.scores.tolist()
for t in range(len(self.scores)):
self.boxes[t].append(self.scores[t])
self.boxes = np.array(self.boxes)
if len(self.boxes) == 0:
return 0, None, None, None, None
else:
return 1, self.boxes, self.scores, self.mask_maps, self.class_ids
def prepare_input(self, image):
self.img_height, self.img_width = image.shape[:2]
input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Resize input image
input_img = cv2.resize(input_img, (self.input_width, self.input_height))
# Scale input pixel values to 0 to 1
input_img = input_img / 255.0
input_img = input_img.transpose(2, 0, 1)
input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32)
return input_tensor
def process_box_output(self, box_output):
predictions = np.squeeze(box_output).T
num_classes = box_output.shape[1] - self.num_masks - 4
# Filter out object confidence scores below threshold
scores = np.max(predictions[:, 4:4+num_classes], axis=1)
predictions = predictions[scores > self.conf_threshold, :]
scores = scores[scores > self.conf_threshold]
if len(scores) == 0:
return [], [], [], np.array([])
box_predictions = predictions[..., :num_classes+4]
mask_predictions = predictions[..., num_classes+4:]
# Get the class with the highest confidence
class_ids = np.argmax(box_predictions[:, 4:], axis=1)
# Get bounding boxes for each object
boxes = self.extract_boxes(box_predictions)
# Apply non-maxima suppression to suppress weak, overlapping bounding boxes
indices = nms(boxes, scores, self.iou_threshold)
return boxes[indices], scores[indices], class_ids[indices], mask_predictions[indices]
def process_mask_output(self, mask_predictions, mask_output):
if mask_predictions.shape[0] == 0:
return []
mask_output = np.squeeze(mask_output)
# Calculate the mask maps for each box
num_mask, mask_height, mask_width = mask_output.shape # CHW
masks = sigmoid(mask_predictions @ mask_output.reshape((num_mask, -1)))
masks = masks.reshape((-1, mask_height, mask_width))
# Downscale the boxes to match the mask size
scale_boxes = self.rescale_boxes(self.boxes,
(self.img_height, self.img_width),
(mask_height, mask_width))
# For every box/mask pair, get the mask map
mask_maps = np.zeros((len(scale_boxes), self.img_height, self.img_width))
blur_size = (int(self.img_width / mask_width), int(self.img_height / mask_height))
for i in range(len(scale_boxes)):
scale_x1 = int(math.floor(scale_boxes[i][0]))
scale_y1 = int(math.floor(scale_boxes[i][1]))
scale_x2 = int(math.ceil(scale_boxes[i][2]))
scale_y2 = int(math.ceil(scale_boxes[i][3]))
x1 = int(math.floor(self.boxes[i][0]))
y1 = int(math.floor(self.boxes[i][1]))
x2 = int(math.ceil(self.boxes[i][2]))
y2 = int(math.ceil(self.boxes[i][3]))
scale_crop_mask = masks[i][scale_y1:scale_y2, scale_x1:scale_x2]
crop_mask = cv2.resize(scale_crop_mask,
(x2 - x1, y2 - y1),
interpolation=cv2.INTER_CUBIC)
crop_mask = cv2.blur(crop_mask, blur_size)
crop_mask = (crop_mask > 0.5).astype(np.uint8)
mask_maps[i, y1:y2, x1:x2] = crop_mask
return mask_maps
def extract_boxes(self, box_predictions):
# Extract boxes from predictions
boxes = box_predictions[:, :4]
# Scale boxes to original image dimensions
boxes = self.rescale_boxes(boxes,
(self.input_height, self.input_width),
(self.img_height, self.img_width))
# Convert boxes to xyxy format
boxes = xywh2xyxy(boxes)
# Check the boxes are within the image
boxes[:, 0] = np.clip(boxes[:, 0], 0, self.img_width)
boxes[:, 1] = np.clip(boxes[:, 1], 0, self.img_height)
boxes[:, 2] = np.clip(boxes[:, 2], 0, self.img_width)
boxes[:, 3] = np.clip(boxes[:, 3], 0, self.img_height)
return boxes
def draw_detections(self, image, draw_scores=True, mask_alpha=0.4):
return draw_detections(image, self.boxes, self.scores,
self.class_ids, mask_alpha)
def draw_masks(self, image, draw_scores=True, mask_alpha=0.5):
return draw_detections(image, self.boxes, self.scores,
self.class_ids, mask_alpha, mask_maps=self.mask_maps)
@staticmethod
def rescale_boxes(boxes, input_shape, image_shape):
# Rescale boxes to original image dimensions
input_shape = np.array([input_shape[1], input_shape[0], input_shape[1], input_shape[0]])
boxes = np.divide(boxes, input_shape, dtype=np.float32)
boxes *= np.array([image_shape[1], image_shape[0], image_shape[1], image_shape[0]])
return boxes
# frame = cv2.imread('../1.png')
# # frame = cv2.resize(frame,(640,640))
# model_path = ("../runs/segment/train2/weights/last-0903_openvino_model/last-0903.xml")
# #model_path = ("../yolov8n-seg_openvino_model/yolov8n-seg.xml")
# device_name = "GPU"
# yoloseg = yolov8_segment_openvino(model_path, device_name, conf_thres=0.3, iou_thres=0.3)
#
# start = time.time()
# boxes, scores, class_ids, masks = yoloseg(frame)
# # postprocess and draw masks
# combined_img = yoloseg.draw_masks(frame)
# end = time.time()
# # show FPS
# print(end- start)
# fps = (1 / (end - start))
# fps_label = "Throughput: %.2f FPS" % fps
# cv2.putText(combined_img, fps_label, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
# cv2.imwrite('rs2.jpg', combined_img)
# # show ALL
# cv2.imshow("YOLOv8 Segmentation OpenVINO inference Demo", combined_img)
#
# cv2.waitKey(0)
# # Initialize the VideoCapture
# cap = cv2.VideoCapture("store-aisle-detection.mp4")
# img = cv2.imread('../1.png')
# # Initialize YOLOv5 Instance Segmentator
# model_path = "yolov8n-seg.xml"
# device_name = "GPU"
# yoloseg = YOLOSeg(model_path, device_name, conf_thres=0.3, iou_thres=0.3)
# while cap.isOpened():
# # Read frame from the video
# ret, frame = cap.read()
# if not ret:
# break
# # Update object localizer
# start = time.time()
# boxes, scores, class_ids, masks = yoloseg(frame)
# # postprocess and draw masks
# combined_img = yoloseg.draw_masks(frame)
# end = time.time()
# # show FPS
# fps = (1 / (end - start))
# fps_label = "Throughput: %.2f FPS" % fps
# cv2.putText(combined_img, fps_label, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
# # show ALL
# cv2.imshow("YOLOv8 Segmentation OpenVINO inference Demo", combined_img)
#
# # Press Any key stop
# if cv2.waitKey(1) > -1:
# print("finished by user")
# break
#
# cap.release()
# cv2.destroyAllWindows()