512 lines
17 KiB
Python
512 lines
17 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
"""
|
||
|
|
\author Ma Yixiao
|
||
|
|
\date Dec. 2024
|
||
|
|
"""
|
||
|
|
from __future__ import division
|
||
|
|
from __future__ import print_function
|
||
|
|
|
||
|
|
import actionlib
|
||
|
|
import PIL
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
import matplotlib
|
||
|
|
import skimage.io
|
||
|
|
import time
|
||
|
|
import logging
|
||
|
|
import threading
|
||
|
|
import copy
|
||
|
|
import cv2
|
||
|
|
from skimage.transform import resize
|
||
|
|
import message_filters
|
||
|
|
|
||
|
|
from cv_bridge import CvBridge, CvBridgeError
|
||
|
|
import math
|
||
|
|
import random
|
||
|
|
import sys
|
||
|
|
|
||
|
|
from maskrcnn_ros.msg import *
|
||
|
|
import maskrcnn_ros.msg
|
||
|
|
|
||
|
|
import ros_numpy
|
||
|
|
import numpy as np
|
||
|
|
import roslib.packages
|
||
|
|
import rospy
|
||
|
|
from sensor_msgs.msg import Image
|
||
|
|
from ultralytics import YOLO
|
||
|
|
from vision_msgs.msg import Detection2D, Detection2DArray, ObjectHypothesisWithPose
|
||
|
|
from ultralytics_ros.msg import YoloResult
|
||
|
|
#myx : add for Time of each request
|
||
|
|
from std_msgs.msg import Float32
|
||
|
|
|
||
|
|
import colorsys
|
||
|
|
import random
|
||
|
|
|
||
|
|
import os
|
||
|
|
#os.environ["CUDA_VISIBLE_DEVICES"]="1"
|
||
|
|
#myx : gpu
|
||
|
|
import torch
|
||
|
|
|
||
|
|
print(torch.__version__)
|
||
|
|
print(torch.cuda.is_available())
|
||
|
|
|
||
|
|
if torch.cuda.is_available():
|
||
|
|
print(torch.cuda.get_device_name(0))
|
||
|
|
print(torch.version.cuda)
|
||
|
|
|
||
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||
|
|
print(device)
|
||
|
|
|
||
|
|
logging.basicConfig(filename='yolov8_action_server.log',
|
||
|
|
format='%(asctime)s %(levelname)s:%(message)s',
|
||
|
|
level=logging.INFO,
|
||
|
|
datefmt='%m/%d/%Y %I:%M:%S %p')
|
||
|
|
|
||
|
|
|
||
|
|
ROOT_DIR = os.path.abspath(
|
||
|
|
"/root/catkin_ws/src/ultralytics_ros/include/ultralytics/")
|
||
|
|
print("ROOT_DIR: ", ROOT_DIR)
|
||
|
|
|
||
|
|
sys.path.append(ROOT_DIR) # To find local version of the library
|
||
|
|
|
||
|
|
sys.path.append(os.path.join(ROOT_DIR, "samples/coco/"))
|
||
|
|
|
||
|
|
|
||
|
|
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
|
||
|
|
IMAGE_DIR = os.path.join(ROOT_DIR, "images")
|
||
|
|
|
||
|
|
class_names = [
|
||
|
|
'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
|
||
|
|
'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign',
|
||
|
|
'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
||
|
|
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag',
|
||
|
|
'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite',
|
||
|
|
'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
||
|
|
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon',
|
||
|
|
'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
|
||
|
|
'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant',
|
||
|
|
'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
|
||
|
|
'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
|
||
|
|
'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
||
|
|
'hair drier', 'toothbrush'
|
||
|
|
]
|
||
|
|
# Totoal time consuming
|
||
|
|
total_time = float(0.0)
|
||
|
|
# Total number of images
|
||
|
|
total_number = int(0)
|
||
|
|
|
||
|
|
# To sync semantic segmentation thread and action server thread
|
||
|
|
cn_task = threading.Condition()
|
||
|
|
cn_ready = threading.Condition()
|
||
|
|
|
||
|
|
|
||
|
|
batch_size = 2
|
||
|
|
#myx global results
|
||
|
|
def random_colors(N, bright=True):
|
||
|
|
"""
|
||
|
|
Generate random colors.
|
||
|
|
To get visually distinct colors, generate them in HSV space then
|
||
|
|
convert to RGB.
|
||
|
|
"""
|
||
|
|
brightness = 1.0 if bright else 0.7
|
||
|
|
hsv = [(float(i) / N, 1, brightness) for i in range(N)]
|
||
|
|
colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))
|
||
|
|
random.shuffle(colors)
|
||
|
|
return colors
|
||
|
|
|
||
|
|
def apply_mask(image, mask, color, alpha=0.5):
|
||
|
|
"""Apply the given mask to the image.
|
||
|
|
"""
|
||
|
|
for c in range(3):
|
||
|
|
image[:, :, c] = np.where(
|
||
|
|
mask == 1, image[:, :, c] * (1 - alpha) + alpha * color[c] * 255,
|
||
|
|
image[:, :, c])
|
||
|
|
return image
|
||
|
|
|
||
|
|
def ros_semantic_result(image,
|
||
|
|
boxes,
|
||
|
|
masks,
|
||
|
|
class_ids,
|
||
|
|
class_names,
|
||
|
|
scores=None,
|
||
|
|
show_mask=True,
|
||
|
|
show_bbox=True,
|
||
|
|
colors=None,
|
||
|
|
show_opencv=None):
|
||
|
|
|
||
|
|
N = boxes.shape[0]
|
||
|
|
if not N:
|
||
|
|
print("\n*** No instances to display *** \n")
|
||
|
|
else:
|
||
|
|
# print("boxes.shape[0] == masks.shape[-1] == class_ids.shape[0]",boxes.shape[0],len(masks),class_ids.shape[0])
|
||
|
|
assert boxes.shape[0] == len(masks) == class_ids.shape[0]
|
||
|
|
|
||
|
|
# Generate random colors
|
||
|
|
colors = colors or random_colors(N)
|
||
|
|
|
||
|
|
height, width, channels = image.shape
|
||
|
|
print("h*w*c*dtype= ", height, width, channels, image.dtype.type)
|
||
|
|
if image is not None:
|
||
|
|
masked_image = image.astype(np.uint8).copy()
|
||
|
|
else:
|
||
|
|
masked_image = np.zeros(image.shape, dtype=np.uint8)
|
||
|
|
|
||
|
|
for i in range(N):
|
||
|
|
#myx : only people
|
||
|
|
# if class_ids[i] != 0:
|
||
|
|
# continue
|
||
|
|
|
||
|
|
color = colors[i]
|
||
|
|
|
||
|
|
# Bounding box
|
||
|
|
if not np.any(boxes[i]):
|
||
|
|
# Skip this instance. Has no bbox. Likely lost in image cropping.
|
||
|
|
continue
|
||
|
|
# y1, x1, y2, x2 = boxes[i].xyxy
|
||
|
|
# box = boxes[i].numpy()
|
||
|
|
#box = boxes[i].xyxy
|
||
|
|
# myx :
|
||
|
|
|
||
|
|
# x1 = box[0][0]
|
||
|
|
# y1 = box[0][1]
|
||
|
|
# x2 = box[0][2]
|
||
|
|
# y2 = box[0][3]
|
||
|
|
# # print("x1 y1 x2 y2",x1,y1,x2,y2)
|
||
|
|
|
||
|
|
# x1 = int(x1.item())
|
||
|
|
# y1 = int(y1.item())
|
||
|
|
# x2 = int(x2.item())
|
||
|
|
# y2 = int(y2.item())
|
||
|
|
|
||
|
|
|
||
|
|
# if show_bbox:
|
||
|
|
# masked_image = cv2.rectangle(image, (x1, y1), (x2, y2), (255,20,147), 1)
|
||
|
|
class_id = class_ids[i]
|
||
|
|
# score = scores[i] if scores is not None else None
|
||
|
|
# label = class_names[class_id]
|
||
|
|
# print("class ID: ", str(class_id), "label: ", str(label), " score: ", str(score))
|
||
|
|
mask = masks[i]
|
||
|
|
if show_mask:
|
||
|
|
masked_image = apply_mask(masked_image, mask, color)
|
||
|
|
#has the different color and boxes!!
|
||
|
|
cv2.imwrite('/root/Dataset/output!!!!!!!!!!!!@@@@@.jpg', masked_image)
|
||
|
|
|
||
|
|
# cv2.namedWindow("Result")
|
||
|
|
if show_opencv:
|
||
|
|
if N and np.any(masked_image):
|
||
|
|
cv2.imshow("Result", masked_image)
|
||
|
|
cv2.waitKey(1)
|
||
|
|
|
||
|
|
return masked_image
|
||
|
|
|
||
|
|
def save_image(image,i,id):
|
||
|
|
filename = os.path.join("/root/Dataset/", str(i) + "_" + str(id) + ".jpg")
|
||
|
|
cv2.imwrite(filename, image)
|
||
|
|
def save_mask_image(image,i,id):
|
||
|
|
filename = os.path.join("/root/Dataset/", str(i) + "___" + str(id) + ".jpg")
|
||
|
|
cv2.imwrite(filename, image)
|
||
|
|
|
||
|
|
class YoloV8(object):
|
||
|
|
def __init__(self):
|
||
|
|
print('Init YOLO')
|
||
|
|
|
||
|
|
yolo_model = rospy.get_param("~yolo_model", "yolov8n-seg.pt")
|
||
|
|
self.classes = rospy.get_param("~classes", class_names)
|
||
|
|
path = roslib.packages.get_pkg_dir("ultralytics_ros")
|
||
|
|
self.model = YOLO(f"{path}/models/{yolo_model}")
|
||
|
|
self.model.fuse()
|
||
|
|
self.masked_image__pub = rospy.Publisher("yolo_image", Image, queue_size=10)
|
||
|
|
self.use_segmentation = yolo_model.endswith("-seg.pt")
|
||
|
|
print("=============Initialized YOLO==============")
|
||
|
|
|
||
|
|
def segment(self, image):
|
||
|
|
timer_start = rospy.Time.now()
|
||
|
|
results = self.model.predict(
|
||
|
|
source=image,
|
||
|
|
device=0,
|
||
|
|
retina_masks=True,
|
||
|
|
)
|
||
|
|
# results = self.model(image)
|
||
|
|
|
||
|
|
segment_time = (rospy.Time.now() - timer_start).to_sec() * 1000
|
||
|
|
print("predict time: %f ms \n" % segment_time)
|
||
|
|
print("res.len" ,len(results))
|
||
|
|
|
||
|
|
#maybe the results is NULL
|
||
|
|
timer_start = rospy.Time.now()
|
||
|
|
#r = results[0]
|
||
|
|
if results[0].boxes.cls.numel() != 0:
|
||
|
|
boxes = results[0].boxes
|
||
|
|
# print("222222222222222222222222222",boxes.shape[0])
|
||
|
|
classes = results[0].boxes.cls
|
||
|
|
confidence_score = results[0].boxes.conf
|
||
|
|
class_ids = []
|
||
|
|
scores = []
|
||
|
|
masks = []
|
||
|
|
|
||
|
|
for cls, conf in zip(classes, confidence_score):
|
||
|
|
id = int(cls)
|
||
|
|
class_ids.append(id)
|
||
|
|
score = float(conf)
|
||
|
|
scores.append(score)
|
||
|
|
|
||
|
|
for mask_tensor in results[0].masks:
|
||
|
|
mask_numpy = (
|
||
|
|
np.squeeze(mask_tensor.data.to("cpu").detach().numpy()).astype(
|
||
|
|
np.uint8
|
||
|
|
)
|
||
|
|
)
|
||
|
|
masks.append(mask_numpy)
|
||
|
|
|
||
|
|
#masks = results[0].masks
|
||
|
|
scores = np.array(scores)
|
||
|
|
class_ids = np.array(class_ids)
|
||
|
|
# masks_count = masks.shape[-1]
|
||
|
|
|
||
|
|
#print("!!!!!!!!!!!!!masks_count = ",results[0].masks.shape[-1])
|
||
|
|
# print("masks_count = ",len(masks))
|
||
|
|
|
||
|
|
self.masked_image_ = ros_semantic_result(image[0],
|
||
|
|
boxes,
|
||
|
|
masks,
|
||
|
|
class_ids,
|
||
|
|
class_names,
|
||
|
|
scores,
|
||
|
|
show_opencv=False)
|
||
|
|
|
||
|
|
else:
|
||
|
|
print("..............111...............")
|
||
|
|
self.masked_image_ = np.zeros(image[0].shape, dtype=np.uint8)
|
||
|
|
# self.masked_image_ = image.astype(np.uint8).copy()
|
||
|
|
|
||
|
|
segment_time = (rospy.Time.now() - timer_start).to_sec() * 1000
|
||
|
|
print("Visualize time: %f ms \n" % segment_time)
|
||
|
|
|
||
|
|
return self.masked_image_, results
|
||
|
|
|
||
|
|
def publish_result(self, image, stamp):
|
||
|
|
if image is None or stamp is None:
|
||
|
|
print("Image invalid")
|
||
|
|
return
|
||
|
|
msg_img = CvBridge().cv2_to_imgmsg(image, encoding='passthrough')
|
||
|
|
msg_img.header.stamp = stamp
|
||
|
|
self.masked_image__pub.publish(msg_img)
|
||
|
|
|
||
|
|
def worker(yolov8, bPublish_result=False):
|
||
|
|
global bStartYOLO
|
||
|
|
global color_image
|
||
|
|
global stamp
|
||
|
|
global masked_image
|
||
|
|
global result
|
||
|
|
# global color_image_np
|
||
|
|
# global masked_image_msg
|
||
|
|
# global masked_msg
|
||
|
|
|
||
|
|
yolov8 = YoloV8()
|
||
|
|
bStartYOLO = True
|
||
|
|
color_image = []
|
||
|
|
#color_image_np = []
|
||
|
|
stamp = []
|
||
|
|
while bStartYOLO:
|
||
|
|
with cn_task:
|
||
|
|
cn_task.wait()
|
||
|
|
print("New task comming")
|
||
|
|
timer_start = rospy.Time.now()
|
||
|
|
print("color_image&&&&&&&&&&&&&&&&& len &&&&&&&&&&&",len(color_image))
|
||
|
|
masked_image, result = yolov8.segment(color_image)
|
||
|
|
|
||
|
|
segment_time = (rospy.Time.now() - timer_start).to_sec() * 1000
|
||
|
|
print("Yolo segment time for cuurent image: %f ms \n" %
|
||
|
|
segment_time)
|
||
|
|
with cn_ready:
|
||
|
|
cn_ready.notifyAll()
|
||
|
|
if bPublish_result:
|
||
|
|
yolov8.publish_result(masked_image, stamp)
|
||
|
|
|
||
|
|
print("Exit YoloV8 thread")
|
||
|
|
|
||
|
|
class SemanticActionServer(object):
|
||
|
|
_feedback = maskrcnn_ros.msg.batchFeedback()
|
||
|
|
_result = maskrcnn_ros.msg.batchResult()
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
print("Initialize Action Server")
|
||
|
|
|
||
|
|
# Get action server name
|
||
|
|
self._action_name = rospy.get_param('/semantic/action_name',
|
||
|
|
'/yolov8_action_server')
|
||
|
|
print("Action name: ", self._action_name)
|
||
|
|
# Start Action server
|
||
|
|
self._as = actionlib.SimpleActionServer(
|
||
|
|
self._action_name,
|
||
|
|
maskrcnn_ros.msg.batchAction,
|
||
|
|
execute_cb=self.execute_cb,
|
||
|
|
auto_start=False)
|
||
|
|
self._as.start()
|
||
|
|
print("YoloV8 action server start...")
|
||
|
|
#myx : Time of each request
|
||
|
|
#self.time_pub = rospy.Publisher('/time_of_request', Float32, queue_size = 1)
|
||
|
|
|
||
|
|
def execute_cb(self, goal):
|
||
|
|
global total_time
|
||
|
|
global total_number
|
||
|
|
global color_image
|
||
|
|
global stamp
|
||
|
|
global masked_image
|
||
|
|
# global masked_image_msg
|
||
|
|
# global masked_msg
|
||
|
|
global result
|
||
|
|
# global bSaveResult
|
||
|
|
|
||
|
|
# clear the old data
|
||
|
|
color_image = []
|
||
|
|
#color_image_np = []
|
||
|
|
stamp = []
|
||
|
|
|
||
|
|
if not self._as.is_active():
|
||
|
|
logging.debug("[Error] YoloV8 action server cannot active")
|
||
|
|
return
|
||
|
|
|
||
|
|
time_each_loop_start = rospy.Time.now()
|
||
|
|
|
||
|
|
print("----------------------------------")
|
||
|
|
print("ID: %d" % goal.id)
|
||
|
|
|
||
|
|
for i in range(batch_size):
|
||
|
|
try:
|
||
|
|
color_image.append(
|
||
|
|
CvBridge().imgmsg_to_cv2(goal.image[i], 'bgr8'))
|
||
|
|
print("Number of images:", len(goal.image))
|
||
|
|
# if goal.id >= 136 :
|
||
|
|
# save_image(CvBridge().imgmsg_to_cv2(goal.image[i], 'bgr8') ,i , goal.id)
|
||
|
|
stamp.append(goal.image[i].header.stamp)
|
||
|
|
|
||
|
|
except CvBridgeError as e:
|
||
|
|
print(e)
|
||
|
|
return
|
||
|
|
|
||
|
|
|
||
|
|
timer_start = rospy.Time.now()
|
||
|
|
# perform segmenting
|
||
|
|
if np.any(color_image[0]):
|
||
|
|
with cn_task:
|
||
|
|
print("Inform new task")
|
||
|
|
cn_task.notifyAll()
|
||
|
|
|
||
|
|
with cn_ready:
|
||
|
|
cn_ready.wait()
|
||
|
|
print("semantic result ready")
|
||
|
|
|
||
|
|
segment_time = (rospy.Time.now() - timer_start).to_sec() * 1000
|
||
|
|
print("Yolov8 segment time: %f ms \n" % segment_time)
|
||
|
|
# calculate average time consuming
|
||
|
|
total_time = float(total_time) + float(segment_time)
|
||
|
|
total_number = total_number + 1
|
||
|
|
if int(total_number) > 0:
|
||
|
|
average = total_time / float(total_number)
|
||
|
|
print("Average time: %f ms" % average)
|
||
|
|
|
||
|
|
# object_num = []
|
||
|
|
labelMsgs = []
|
||
|
|
# scoreMsgs = []
|
||
|
|
|
||
|
|
for i in range(batch_size):
|
||
|
|
if result[i].boxes.cls.numel() != 0:
|
||
|
|
|
||
|
|
box_es = result[i].boxes
|
||
|
|
# masks = result[i].masks
|
||
|
|
masks = []
|
||
|
|
for mask_tensor in result[i].masks:
|
||
|
|
mask_numpy = (
|
||
|
|
np.squeeze(mask_tensor.data.to("cpu").detach().numpy()).astype(
|
||
|
|
np.uint8
|
||
|
|
)
|
||
|
|
)
|
||
|
|
masks.append(mask_numpy)
|
||
|
|
# Total number of objects
|
||
|
|
N = box_es.shape[0]
|
||
|
|
|
||
|
|
|
||
|
|
label_img = np.zeros(masks[0].shape, dtype=np.uint8)
|
||
|
|
|
||
|
|
logging.info('shap of label image: %s', label_img.shape)
|
||
|
|
|
||
|
|
# file = open("/root/Dataset/outputtxt.txt", "w")
|
||
|
|
class_color = 10 # myx every person has the different color
|
||
|
|
person_color = 1
|
||
|
|
for n in range(N):
|
||
|
|
id = int(result[i].boxes.cls[n])
|
||
|
|
#print("..............N = ",N)
|
||
|
|
mask = masks[n] * 255
|
||
|
|
mask_area = np.sum(mask == 255) # 计算掩码的面积，即 mask 中像素值为 255 的个数
|
||
|
|
#print("......................Mask area............................ = ",mask_area)
|
||
|
|
if mask_area < 200:
|
||
|
|
continue
|
||
|
|
#person's id is 0!!! id + 10 ,input all objects' class
|
||
|
|
if id == 0:
|
||
|
|
label_img += (masks[n] * person_color ).astype(np.uint8)
|
||
|
|
person_color += 1
|
||
|
|
else:
|
||
|
|
label_img += (masks[n] * class_color ).astype(np.uint8)
|
||
|
|
class_color += 10 #the first object ID is 10, second is 20 ...
|
||
|
|
cv2.imwrite('/root/Dataset/output_all_mask.jpg', label_img)
|
||
|
|
msg_label = CvBridge().cv2_to_imgmsg(label_img.astype(np.uint8), encoding='mono8')
|
||
|
|
labelMsgs.append(msg_label)
|
||
|
|
|
||
|
|
else:
|
||
|
|
print(".................22222...................")
|
||
|
|
label_img = np.zeros_like(color_image[0])
|
||
|
|
label_img = cv2.cvtColor(label_img, cv2.COLOR_BGR2GRAY)
|
||
|
|
msg_label = CvBridge().cv2_to_imgmsg(label_img.astype(np.uint8), encoding='mono8')
|
||
|
|
labelMsgs.append(msg_label)
|
||
|
|
|
||
|
|
|
||
|
|
self._result.id = goal.id
|
||
|
|
self._result.label = labelMsgs # all objects' mask
|
||
|
|
|
||
|
|
self._feedback.complete = True
|
||
|
|
self._as.set_succeeded(self._result)
|
||
|
|
self._as.publish_feedback(self._feedback)
|
||
|
|
|
||
|
|
time_each_loop = (rospy.Time.now() -
|
||
|
|
time_each_loop_start).to_sec() * 1000
|
||
|
|
print("Time of each request: %f ms \n" % time_each_loop)
|
||
|
|
#myx :
|
||
|
|
#self.time_pub.publish(time_each_loop)
|
||
|
|
|
||
|
|
|
||
|
|
def main(args):
|
||
|
|
global total_time
|
||
|
|
global total_number
|
||
|
|
|
||
|
|
|
||
|
|
rospy.init_node('yolov8_server', anonymous=False)
|
||
|
|
|
||
|
|
th_yolov8 = threading.Thread(name='yolov8',
|
||
|
|
target=worker,
|
||
|
|
args=(
|
||
|
|
YoloV8,
|
||
|
|
False,
|
||
|
|
))
|
||
|
|
th_yolov8.start()
|
||
|
|
|
||
|
|
actionserver = SemanticActionServer()
|
||
|
|
|
||
|
|
print('Setting up YoloV8 Action Server...')
|
||
|
|
|
||
|
|
logging.debug('Waiting for worker threads')
|
||
|
|
bStartYOLO = False
|
||
|
|
main_thread = threading.currentThread()
|
||
|
|
|
||
|
|
try:
|
||
|
|
rospy.spin()
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
print("Shutting down")
|
||
|
|
|
||
|
|
for t in threading.enumerate():
|
||
|
|
if t is not main_thread:
|
||
|
|
t.join()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main(sys.argv)
|