#!/usr/bin/env python3
-- coding: utf-8 --
import rospy
import cv2
import atexit
import numpy as np
import torch
from ultralytics import YOLO
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from detection_msgs.msg import BoundingBox, BoundingBoxes, Label
class Detector(object):
def init(self):
self.class_names = {
0: '0',
1: '1',
2: '2',
3: '3',
4: '4',
5: '5',
6: '6',
7: '7',
8: '8',
9: '9',
10: '10',
11: '11'
}
self.task = 'detect'
self.weights = '/home/user/catkin_ws/src/ERP42-main/yolo/ultralytics/cfg/datasets/runs/detect/pothole/weights/best.pt' # 첫 번째 모델의 가중치 파일 경로
self.device = 0
self.half = True
self.img_size = 416
self.conf_thres = 0.5
self.iou_thres = 0.7
self.img_topic = "/usb_cam/image_raw"
self.model1 = YOLO(self.weights)
# 두 번째 모델 설정
self.weights2 = '/home/user/catkin_ws/src/ERP42-main/yolo/ultralytics/cfg/datasets/runs/detect/manhole/weights/best.pt' # 두 번째 모델의 가중치 파일 경로
self.model2 = YOLO(self.weights2)
self.model1 = YOLO(self.weights)#, self.task)
self.bridge = CvBridge()
self.img_sub = rospy.Subscriber(self.img_topic, Image, self.img_callback)
self.BoundingBox_pub = rospy.Publisher("/Bounding_box", BoundingBox, queue_size=5)
self.pred_pub = rospy.Publisher("/Bounding_boxs", BoundingBoxes, queue_size=5)
self.process_time_old = 0.0
self.process_times = []
atexit.register(self.print_avg_process_time)
print("ready...")
def img_callback(self, data):
img = self.bridge.imgmsg_to_cv2(data, 'bgr8')
# 첫 번째 모델로 검출
results1 = self.model1(img)
xyxy1 = results1[0].boxes.xyxy.to('cpu').numpy()
clss1 = results1[0].boxes.cls
conf1 = results1[0].boxes.conf
# 두 번째 모델로 검출
results2 = self.model2(img)
xyxy2 = results2[0].boxes.xyxy.to('cpu').numpy()
clss2 = results2[0].boxes.cls
conf2 = results2[0].boxes.conf
# 두 모델의 결과를 합치기
xyxy = np.concatenate((xyxy1, xyxy2), axis=0)
clss = np.concatenate((clss1.cpu(), clss2.cpu()), axis=0)
conf = np.concatenate((conf1.cpu(), conf2.cpu()), axis=0)
# 이후 처리는 합쳐진 결과에 대해 진행
bounding_boxes = BoundingBoxes()
bounding_boxes.header = data.header
bounding_boxes.image_header = data.header
for box, label, confidence in zip(xyxy, clss, conf):
bounding_box = BoundingBox()
x1, y1, x2, y2 = box
bounding_box.Class = self.class_names[label.item()]
bounding_box.probability = confidence
bounding_box.xmin = int(x1)
bounding_box.ymin = int(y1)
bounding_box.xmax = int(x2)
bounding_box.ymax = int(y2)
bounding_boxes.bounding_boxes.append(bounding_box)
if self.BoundingBox_pub is not None:
self.BoundingBox_pub.publish(bounding_box)
x1 = int(x1)
y1 = int(y1)
x2 = int(x2)
y2 = int(y2)
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
class_name = self.class_names[label.item()]
label_text = f'{class_name}: {confidence.item():.2f}'
(text_width, text_height), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
text_x = x1
text_y = y1 - 10 if y1 - 10 > 10 else y1 + 20
cv2.rectangle(img, (text_x, text_y), (text_x + text_width, text_y - text_height), (0, 255, 0), cv2.FILLED)
cv2.putText(img, label_text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
# 최종 결과를 발행하고 화면에 표시
self.pred_pub.publish(bounding_boxes)
cv2.imshow('Image with Bounding Boxes and Labels', img)
cv2.waitKey(1)
def print_avg_process_time(self):
avg_process_time = sum(self.process_times) / len(self.process_times)
#print(f"Average model process time: {avg_process_time} ms")
if name == 'main':
rospy.init_node("detector")
detector = Detector()
rospy.spin()
# 0: '0',
# 1: '1',
# 2: '10',
# 3: '11',
# 4: '2',
# 5: '3',
# 6: '4',
# 7: '5',
# 8: '6',
# 9: '7',
# 10: '8',
# 11: '9' 안녕하세요 라즈베리파이4를 mobaxterm에서 linux버전으로 실행시키고 있는데 이코드를 실행시키면 화면이 나오자마자 첫장면만 나오고 그 이후에 멈춰있습니다 이코드를 어떻게 수정해야되나요?
를 한 번 눌러주시길 부탁드려요!