update vision

This commit is contained in:
hjw
2024-09-11 14:27:58 +08:00
parent 2b20732119
commit 6bc52d70f5
3 changed files with 13 additions and 9 deletions

View File

@ -15,11 +15,12 @@ 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, use_openvino_model = False):
def __init__(self, use_openvino_model = True):
self.use_openvino_model = use_openvino_model
if self.use_openvino_model==False:
model_path = ''.join([os.getcwd(), '/model/pt/best.pt'])
@ -28,14 +29,14 @@ class Detection:
self.model = yolov8_segment()
self.model.load_model(model_path, device)
else:
model_path = ''.join([os.getcwd(), '/model/openvino/last-0903.xml'])
model_path = ''.join([os.getcwd(), '/model/openvino/best.xml'])
device = 'CPU'
self.camera_rvc = camera()
self.model = yolov8_segment_openvino(model_path, device, conf_thres=0.3, iou_thres=0.3)
def get_position(self, Point_isVision=True):
def get_position(self, Point_isVision=False):
""
'''
:param api: None
@ -68,7 +69,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))