first commit

Author: zohar_fzh

.gitignore

@@ -0,0 +1,4 @@
+*.engine
+*.pt
+*.onnx
+__pycache__

README.md

@@ -0,0 +1 @@
+# dr_open_track

model/export_engine.sh

@@ -0,0 +1 @@
+yolo export model=yolov8n.pt format=engine half=True simplify=True

src/GStreamerWrapper/GStreamerWrapper.py

@@ -0,0 +1,69 @@
+import gi
+gi.require_version('Gst', '1.0')
+from gi.repository import Gst
+
+import numpy as np
+
+import threading
+import time
+
+class GStreamerWrapper:
+    def __init__(self):
+        ##############
+        #init pipeline
+        ##############
+        Gst.init(None)
+
+        # pull from a device and convert with nvvidconv
+        # self.pipeline_str = "v4l2src device=/dev/video0 ! video/x-raw,format=YUY2,width=640,height=480,framerate=30/1 ! nvvidconv ! video/x-raw,format=BGRx ! videoconvert ! video/x-raw,format=BGR ! appsink name=sink"
+        
+        # pull from a device and convert with videoconvert
+        # self.pipeline_str = "v4l2src device=/dev/video0 ! video/x-raw,format=YUY2,width=640,height=480,framerate=30/1 ! videoconvert ! video/x-raw,format=BGR ! appsink name=sink"
+        
+        # pull from a rtsp address and decode with nvv4l2decoder
+        self.pipeline_str = "rtspsrc location=rtsp://192.168.1.120:8554/test ! rtph264depay ! h264parse config-interval=1 ! nvv4l2decoder ! video/x-raw(memory:NVMM) ! nvvidconv ! video/x-raw,format=BGRx ! videoscale ! video/x-raw,width=640,height=360 ! videoconvert ! video/x-raw,format=BGR ! queue max-size-buffers=5 max-size-bytes=0 max-size-time=50000000 leaky=downstream ! appsink name=sink sync=false max-buffers=1 drop=1"
+        
+        # init pipeline
+        self.pipeline = Gst.parse_launch(self.pipeline_str)
+        self.pipeline.set_state(Gst.State.PLAYING)
+        self.sink = self.pipeline.get_by_name('sink')
+        self.sink.set_property("emit-signals", True)
+
+        ###############
+        #get frame size
+        ###############
+        sample = self.sink.emit('pull-sample')
+        caps = sample.get_caps()
+        structure = caps.get_structure(0)
+        self.width = structure.get_value('width')
+        self.height = structure.get_value('height')
+
+        self.break_flag = False
+        self.get_frame_thread = threading.Thread(target=self.GetFrameThreadFunc)
+        self.lock = threading.Lock()
+        self.get_frame_thread.start()
+
+    def __del__(self):
+        self.pipeline.set_state(Gst.State.NULL)
+
+    def StopThread(self):
+        self.break_flag = True
+        self.get_frame_thread.join()
+
+    def GetFrameThreadFunc(self):        
+        while(self.break_flag != True):
+            sample = self.sink.emit('pull-sample')
+            buffer = sample.get_buffer()
+            success, map_info =  buffer.map(Gst.MapFlags.READ)
+            if not success:
+                return None
+            arr = np.frombuffer(map_info.data, dtype=np.uint8)
+            buffer.unmap(map_info)
+            with self.lock:
+                self.frame = arr.reshape((self.height, self.width, 3))
+
+            time.sleep(0.01)
+
+    def GetFrame(self):
+        with self.lock:
+            return self.frame

src/RobotController/FpsCounter/FpsCounter.py

@@ -0,0 +1,19 @@
+import time
+
+class FpsCounter:
+    def __init__(self):
+        self.start_time = time.time()
+        self.count = 0
+        self.fps = 0
+
+    def Count(self):
+        self.count += 1
+        # second
+        t_diff = time.time() - self.start_time
+        if(t_diff) > 1:
+            self.fps = self.count / t_diff
+            self.start_time = time.time()
+            self.count = 0
+
+    def GetFps(self):
+        return self.fps

src/RobotController/ROSTransfer/ROS1Transfer.py

@@ -0,0 +1,15 @@
+
+import rospy
+from geometry_msgs.msg import Twist
+
+class ROS1Transfer:
+    def __init__(self):
+        rospy.init_node('track_twist_publisher')
+        self.cmd_vel_pub = rospy.Publisher('/cmd_vel', Twist, queue_size=1)
+
+    def SendCmdVel(self, linear_velocity, radian_velocity):
+        twist_cmd = Twist()
+        twist_cmd.linear.x = linear_velocity
+        twist_cmd.angular.z = radian_velocity
+        self.cmd_vel_pub.publish(twist_cmd)
+        

src/RobotController/ROSTransfer/ROS2Transfer.py

@@ -0,0 +1,15 @@
+import rclpy
+from rclpy.node import Node
+from geometry_msgs.msg import Twist
+
+class ROS2Transfer(Node):
+    def __init__(self):
+        rclpy.init()
+        super().__init__('track_twist_publisher')
+        self.cmd_vel_pub = self.create_publisher(Twist, 'cmd_vel', 1)
+
+    def SendCmdVel(self, linear_velocity, radian_velocity):
+        twist_cmd = Twist()
+        twist_cmd.linear.x = linear_velocity
+        twist_cmd.angular.z = radian_velocity
+        self.cmd_vel_pub.publish(twist_cmd)

src/RobotController/ROSTransfer/TransferConstants.py

@@ -0,0 +1,3 @@
+
+kMaxRadianVelocity = 1.2
+kMaxLinerVelocity = 1.0

src/RobotController/RobotController.py

@@ -0,0 +1,156 @@
+
+import cv2
+
+from . FpsCounter import FpsCounter
+from . YoloWrapper import YoloWrapper
+from . ROSTransfer import TransferConstants
+
+kUseRos1Transfer = True
+
+if kUseRos1Transfer:
+    from . ROSTransfer import ROS1Transfer
+else:
+    from . ROSTransfer import ROS2Transfer
+
+kDefaultTrackId = 0
+kDefaultTrackMode = False
+
+class RobotController(object):
+    def __init__(self):
+        self.fps_counter = FpsCounter.FpsCounter()
+        self.yolo_wrapper = YoloWrapper.YoloWrapper()
+        if kUseRos1Transfer:
+            self.ros1_transfer = ROS1Transfer.ROS1Transfer()
+        else:
+            self.ros2_transfer = ROS2Transfer.ROS2Transfer()
+    
+        self.is_tracking = kDefaultTrackMode
+        self.target_id = kDefaultTrackId
+        self.id_str = ""
+
+        self.last_linear_velocity = 0.0
+
+    def SetTargetId(self, id):
+        self.target_id = id
+
+    def GetTargetId(self):
+        return self.target_id
+
+    def SetIsTracking(self, state):
+        self.is_tracking = state
+
+    def GetIsTracking(self):
+        return self.is_tracking
+
+    def FindTarget(self, boxes):
+        for box in boxes:
+            if box.id != None:
+                if box.id.item() == self.GetTargetId():
+                    return box
+    
+    def InputAndProcess(self, frame):
+        key = cv2.waitKey(1)
+        if(self.GetIsTracking() == False):
+            if key >= ord('0') and key <= ord('9'):  # 大键盘输入
+                self.id_str += str((int(key - ord('0'))))
+            elif key == ord('\b'):  # 退格
+                self.id_str = self.id_str[:-1]
+            elif key == 10 or key == 13 or key == 141:  # 回车
+                if(self.id_str == ""):
+                    self.SetTargetId(0)
+                else:
+                    self.SetTargetId(int(self.id_str))
+                self.id_str = ""
+
+            if(self.GetTargetId() != kDefaultTrackId):
+                self.SetIsTracking(True)
+
+        else:
+            if key == 10 or key == 13 or key == 141:  # 回车
+                self.id_str == ""
+                self.target_id = 0
+                self.SetIsTracking(False)
+        frame = cv2.putText(frame, self.id_str, (20, 150), cv2.FONT_HERSHEY_PLAIN, 2, [255, 0, 0], 2)
+
+    def TrackAndDraw(self, frame, box):
+        shape = frame.shape
+        frame = cv2.UMat(frame)
+
+        self.fps_counter.Count()
+        frame = cv2.putText(frame, "fps {:.02f}".format(self.fps_counter.GetFps()), (10, 20),
+                    cv2.FONT_HERSHEY_PLAIN, 2, [0, 128, 0], 2)
+        self.InputAndProcess(frame)
+
+        if(box != None):            
+            x1 = int(box.xyxy[0][0].item())
+            y1 = int(box.xyxy[0][1].item())
+            x2 = int(box.xyxy[0][2].item())
+            y2 = int(box.xyxy[0][3].item())
+            center=((x1+x2)//2, (y1+y2)//2)
+            cv2.circle(frame,center,2,[0,0,255],-1) # 画出选框中心点
+
+            #cal radian_velocity
+            angle=(center[0]/shape[1]-0.5)*2        # range -1 to 1
+            radian_velocity = -angle*TransferConstants.kMaxRadianVelocity
+
+            #cal linear_velocity
+            top=y1/shape[0]
+            if top < 0.25:
+                kp_linear_x = 5
+                kd_linear_x = 0.2
+                linear_velocity = top * kp_linear_x - self.last_linear_velocity * kd_linear_x
+            else:
+                linear_velocity=TransferConstants.kMaxLinerVelocity
+
+            #draw
+            cv2.putText(frame,"Selected ID {:}".format(self.GetTargetId()),(20,125), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 2)
+            cv2.putText(frame,"Press \"Enter\" to reset ID",(20,150), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 2)
+            label = '{}{:d}'.format("", self.GetTargetId())
+            t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2 , 2)[0]
+            cv2.rectangle(frame, (x1, y1), (x2, y2), [255,128,128], 2)
+            cv2.rectangle(frame,(x1, y1),(x1+t_size[0]+3,y1+t_size[1]+4), [255,128,128],-1)
+            cv2.putText(frame,label,(x1,y1+t_size[1]+4), cv2.FONT_HERSHEY_PLAIN, 2, [255,255,255], 2)
+            cv2.putText(frame,"{:.02f} m/s".format(linear_velocity),(0,250), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 3)
+            cv2.putText(frame,"{:.02f} rad/s".format(radian_velocity),(0,300), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 3)
+
+            #pub cmdvel
+            if kUseRos1Transfer:
+                self.ros1_transfer.SendCmdVel(linear_velocity, radian_velocity)
+            else:
+                self.ros2_transfer.SendCmdVel(linear_velocity, radian_velocity)
+            self.last_linear_velocity = linear_velocity
+            return frame
+        else:
+            cv2.putText(frame,"Miss Person".format(self.GetTargetId()),(20,100), cv2.FONT_HERSHEY_PLAIN, 2, [0,0,255], 3)
+            cv2.putText(frame,"Press \"Enter\" to reset ID",(20,125), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 2)
+            cv2.putText(frame,"{:.02f} m/s".format(0),(0,250), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 3)
+            cv2.putText(frame,"{:.02f} rad/s".format(0),(0,300), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 3)   
+            return frame
+
+    def NonTrackAndDraw(self, frame):
+        frame = cv2.UMat(frame)
+        self.fps_counter.Count()
+        frame = cv2.putText(frame, "fps {:.02f}".format(self.fps_counter.GetFps()), (10, 20),
+                    cv2.FONT_HERSHEY_PLAIN, 2, [0, 128, 0], 2)
+        self.InputAndProcess(frame)
+
+        cv2.putText(frame,"Stop",(20,100), cv2.FONT_HERSHEY_PLAIN, 2, [0,0,255], 3)
+        cv2.putText(frame,"Enter the object ID:",(20,125), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 2)
+        cv2.putText(frame,"{:.02f} m/s".format(0),(0,250), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 3)
+        cv2.putText(frame,"{:.02f} rad/s".format(0),(0,300), cv2.FONT_HERSHEY_PLAIN, 2, [255,0,0], 3)    
+        return frame
+
+    def Run(self, frame):
+        results = self.yolo_wrapper.Track(frame)
+        if(len(results)>0):
+            if(self.GetIsTracking()):
+                box = self.FindTarget(results[0].boxes)                
+                return self.TrackAndDraw(frame, box)
+            else:
+                frame = results[0].plot()
+                return self.NonTrackAndDraw(frame)
+
+
+
+
+

src/RobotController/YoloWrapper/CocoTypeId.py

@@ -0,0 +1,84 @@
+#coco type id according to: https://docs.ultralytics.com/datasets/detect/coco/#applications
+
+kPerson = 0
+kBicycle = 1
+kMotorcycle = 2
+kCar = 3
+kAirplane = 4
+kBus = 5
+kTrain = 6
+kTruck = 7
+kBoat = 8
+kTrafficLight = 9
+kFireHydrant = 10
+kStopSign = 11
+kParkingMeter = 12
+kBench = 13
+kBird = 14
+kCat = 15
+kDog = 16
+kHorse = 17
+kSheep = 18
+kCow = 19
+kElephant = 20
+kBear = 21
+kZebra = 22
+kGiraffe = 23
+kBackpack = 24
+kUmbrella = 25
+kHandbag = 26
+kTie = 27
+kSuitcase = 28
+kFrisbee = 29
+kSkis = 30
+kSnowboard = 31
+kSportsBall = 32
+kKite = 33
+kBaseballBat = 34
+kBaseballGlove = 35
+kSkateboard = 36
+kSurfboard = 37
+kTennisRacket = 38
+kBottle = 39
+kWineGlass = 40
+kCup = 41
+kFork = 42
+kKnife = 43
+kSpoon = 44
+kBowl = 45
+kBanana = 46
+kApple = 47
+kSandwich = 48
+kOrange = 49
+kBroccoli = 50
+kCarrot = 51
+kHotDog = 52
+kPizza = 53
+kDonut = 54
+kCake = 55
+kChair = 56
+kCouch = 57
+kPottedPlant = 58
+kBed = 59
+kDiningTable = 60
+kToilet = 61
+kTv = 62
+kLaptop = 63
+kMouse = 64
+kRemote = 65
+kKeyboard = 66
+kCellPhone = 67
+kMicrowave = 68
+kOven = 69
+kToaster =70 
+kSink = 71
+kRefrigerator = 72
+kBook = 73
+kClock = 74
+kVase = 75
+kScissors = 76
+kTeddyBear = 77
+kHairDrier = 78
+kToothbrush = 79
+
+