UPDATE openvino

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

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()