DEMO

demo

import CV2 as cv
import os
import numpy as np
import time
from pymycobot.mycobot import MyCobot
from opencv_yolo import yolo
from VideoCapture import FastVideoCapture
from GrabParams import grabParams
import math
import rospy
from geometry_msgs.msg import Twist



import basic
import argparse

parser = argparse.ArgumentParser(description='manual to this script')

#0：'apple', 1：'clock', 2：'banana'
parser.add_argument("--cls", type=int, default=-1)

#pose: 'top', 'bottom'
parser.add_argument("--pose", type=str, default="top")

#遥控调用要关闭debug模式。debug模式会显示识别视频。
parser.add_argument("--debug", type=bool,   default="True")

args = parser.parse_args()

#上层货架抓取，机械臂的控制位置，可以通过注释其他代码进行调试
coords_top_ready = [100.3, -63.8, 330, -91.49, 145, -90.53]
coords_top_ready_ok = [100.3, -63.8, 360, -91.49, 145, -90.53]
coords_top_grap = [170.3, -63.8, 360, -91.49, 145, -90.53]
coords_top_grap_ok = [80.3, -63.8, 360, -91.49, 145, -90.53]

#下层货架抓取，机械臂的控制位置，可以通过注释其他代码进行调试
coords_bottom_ready = [100.3, -63.8, 230, -91.49, 145, -90.53]
coords_bottom_ready_ok = [100.3, -63.8, 260, -91.49, 145, -90.53]
coords_bottom_grap = [170.3, -63.8, 260, -91.49, 145, -90.53]
coords_bottom_grap_ok = [80.3, -63.8, 260, -91.49, 145, -90.53]


class ssdemo(object):
   #初始化函数
   def __init__(self):
       super(ssdemo, self).__init__()

       rospy.init_node('ssdemo', anonymous=True)
       self.pub = rospy.Publisher('cmd_vel', Twist, queue_size=1)
       self.rate = rospy.Rate(20) # 20hz
       

       # Creating a Camera Object
       self.cap = FastVideoCapture(grabParams.cap_num)
     
       self.c_x, self.c_y = grabParams.IMG_SIZE/2, grabParams.IMG_SIZE/2


       self.miss_count = 0

       self.mc = MyCobot(grabParams.usb_dev, grabParams.baudrate)
       self.mc.power_on()

       self.yolo = yolo()

       # Get ArUco marker dict that can be detected.
       self.aruco_dict = cv.aruco.Dictionary_get(cv.aruco.DICT_6X6_250)
       # Get ArUco marker params.
       self.aruco_params = cv.aruco.DetectorParameters_create()
       self.calibrationParams = cv.FileStorage("calibrationFileName.xml", cv.FILE_STORAGE_READ)
       # Get distance coefficient.
       self.dist_coeffs = self.calibrationParams.getNode("distCoeffs").mat()

       height = grabParams.IMG_SIZE
       focal_length = width = grabParams.IMG_SIZE
       self.center = [width / 2, height / 2]
       # Calculate the camera matrix.
       self.camera_matrix = np.array(
           [
               [focal_length, 0, self.center[0]],
               [0, focal_length, self.center[1]],
               [0, 0, 1],
           ],
           dtype=np.float32,
       )
   
   #计算物体的坐标和尺寸
   def get_position(self, corner):
       x = corner[0][0] + corner[1][0] + corner[2][0] + corner[3][0]
       y = corner[0][1] + corner[1][1] + corner[2][1] + corner[3][1]
       x = x/4.0
       y = y/4.0
       x_size_p = abs(x - corner[0][0])*2
       y_size_p = abs(y - corner[0][1])*2
       return (-(x - self.c_x)), (-(y - self.c_y)), x_size_p, y_size_p  

   #计算物体的坐标和尺寸
   def get_position_size(self, box):
       left, top, right, bottom = box
       # print("box",box)
       x = left + right
       y = bottom + top
       x = x*0.5
       y = y*0.5
       x_size_p = abs(left - right)
       y_size_p = abs(top - bottom)
       return (-(x - self.c_x)), (-(y - self.c_y)), x_size_p, y_size_p

   #视频显示
   def show_image(self, img):
       if args.debug:
           cv.imshow("figure", img)
           cv.waitKey(50)
   
   #物体识别
   def obj_detect(self, img):
       self.is_find = False
       img_ori = img
       img_ori = self.transform_frame(img)
       img = self.transform_frame_128(img)

       #加载模型
       net = cv.dnn.readNetFromONNX("comp.onnx")
       
       t1 = time.time()
       #输入数据处理
       blob = cv.dnn.blobFromImage(img, 1 / 255.0, (128, 128), [0, 0, 0], swapRB=True, crop=False)
       net.setInput(blob)
       
       #推理
       outputs = net.forward(net.getUnconnectedOutLayersNames())[0]

       
       #获得识别结果
       boxes, classes, scores = self.yolo.yolov5_post_process_simple(outputs)
       t2 = time.time()

       best_result = (0,0,0)
       
       #物体边框处理
       if boxes is not None:
           boxes = boxes*5
           for box, score, cl in zip(boxes, scores, classes):
               self.image_info = self.get_position_size(box)
               wph = self.image_info[2]/self.image_info[3]
               # print(cl, wph, self.image_info)

               if (wph > 0.8 and wph <1.2):
                   if (args.pose == "top" and args.cls == cl) or (args.pose == "bottom" and self.image_info[1] < 0):
                   #if (args.pose == "top" ) or (args.pose == "bottom" and self.image_info[1] < 0):
                       temp_result = (box, score, cl)                    
                       if score > best_result[1]:
                           best_result = temp_result
                           self.target_image_info = self.get_position_size(box)
                           self.is_find = True
       
       #画物体边框
       if self.is_find:                
           box, score, cl = best_result
           self.yolo.draw_single(img_ori, box, score, cl)
           self.mc.set_color(255,0,0)#red, find object

       #视频显示
       self.show_image(img_ori)
 

   #图像处理，适配物体识别
   def transform_frame(self, frame):
       frame, ratio, (dw, dh) = self.yolo.letterbox(frame, (grabParams.IMG_SIZE, grabParams.IMG_SIZE))

       return frame
       
   #图像处理，适配物体识别
   def transform_frame_128(self, frame):
       frame, ratio, (dw, dh) = self.yolo.letterbox(frame, (128, 128))

       return frame

   #小车后退
   def moveback(self):
       print("backward...")
       count = 25
       move_cmd = Twist()
       while count > 0:
           move_cmd.linear.x = -0.05
           self.pub.publish(move_cmd)
           self.rate.sleep()
           count -= 1

   def moveforward(self):
       print("forward...")
       count = 2
       move_cmd = Twist()
       while count > 0:
           move_cmd.linear.x = 0.05
           self.pub.publish(move_cmd)
           self.rate.sleep()
           count -= 1
   
   #小车右转
   def rotate_to_right(self):
       print("rotate_to_right...")
       count = 20
       move_cmd = Twist()
       while count > 0:
           move_cmd.angular.z = -0.2
           self.pub.publish(move_cmd)
           self.rate.sleep()
           count -= 1
   
   #小车左转
   def rotate_to_left(self):
       print("rotate_to_left...")
       count = 20
       move_cmd = Twist()
       while count > 0:
           move_cmd.angular.z = 0.2
           self.pub.publish(move_cmd)
           self.rate.sleep()            
           count -= 1

   #机械臂放置积木到右侧
   def place2right(self):
       
       angles = [-87.27, 0, 0, 0, 0, 0]
       self.mc.send_angles(angles,25)
       time.sleep(3)
       angles = [-87.27, -45.26, 2.28, 1.66, -0.96, 47.02]
       self.mc.send_angles(angles,25)
       time.sleep(3)

       # open
       basic.grap(False)

       angles = [0, 0, 0, 0, 0, 0]
       self.mc.send_angles(angles,25)

   #夹爪闭合
   def pick(self):
       basic.grap(True)

   #旋转图像
   def rotate_image(self, image, rotate_angle):
       rows,cols  = image.shape[:2]
       angle=rotate_angle
       center = ( cols/2,rows/2)
       heightNew=int(cols*abs(math.sin(math.radians(angle)))+rows*abs(math.cos(math.radians(angle))))
       widthNew=int(rows*abs(math.sin(math.radians(angle)))+cols*abs(math.cos(math.radians(angle))))


       M = cv.getRotationMatrix2D(center,angle,1)
       # print(M)
       M[0,2] +=(widthNew-cols)/2  
       M[1,2] +=(heightNew-rows)/2
       # print(M)

       # print(widthNew,heightNew)
       rotated_image  = cv.warpAffine(image,M,(widthNew,heightNew))
       return rotated_image

   #二维码检测
   def aruco_correct_detect(self):
       img = self.cap.read()    
       img = self.transform_frame(img)      
       img = self.rotate_image(img, -90)
       self.show_image(img)
       
       # transfrom the img to model of gray
       gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
       # Detect ArUco marker.
       corners, ids, rejectImaPoint = cv.aruco.detectMarkers(
           gray, self.aruco_dict, parameters=self.aruco_params
       )

       # print corners, len(corners)

       id_aruco = -1
       if len(corners) > 0:                
           if ids is not None:
               # get informations of aruco
               ret = cv.aruco.estimatePoseSingleMarkers(
                   corners, 0.021, self.camera_matrix, self.dist_coeffs
               )
               # rvec:rotation offset,tvec:translation deviator
               (rvec, tvec) = (ret[0], ret[1])
               (rvec - tvec).any()
               
               # print corners

               for i in range(rvec.shape[0]):
                   # draw the aruco on img
                   cv.aruco.drawDetectedMarkers(img, corners)
                   cv.aruco.drawAxis(
                       img,
                       self.camera_matrix,
                       self.dist_coeffs,
                       rvec[i, :, :],
                       tvec[i, :, :],
                       0.03,
                   )
                   self.show_image(img)
                   # print(corners)
                   real_x, real_y, xsize, ysize = detect.get_position(corners[i][0])
                   self.aruco_info = (real_x, real_y, xsize, ysize)
                   id_aruco = ids[i][0]                    
                   print("aruco info: ", id_aruco, real_x, real_y, xsize, ysize)

                   if not self.is_find_correct_aruco:
                       dx_image_aruco = abs(real_x - self.target_image_info[0])
                       dy_image_aruco = abs(real_y - self.target_image_info[1])
                       print("dx_dy_image_aruco: ", dx_image_aruco, dy_image_aruco)
                       if dx_image_aruco < 100 and dy_image_aruco < 200:
                           self.is_find_correct_aruco = True
                           self.id_correct_aruco = id_aruco #find the id of aruco
                   elif id_aruco == self.id_correct_aruco:
                       dy_image_aruco = abs(real_y - self.target_image_info[1])
                       if dy_image_aruco >= 200:
                           id_aruco = -1 # not correct aruco, maybe it's on top floor
                       break
           return id_aruco

                   

   #跟随二维码
   def follow_aruco(self):          
       self.is_find_correct_aruco = False
       self.is_follow_aruco_done = False

       while not self.is_follow_aruco_done:
           id_aruco = self.aruco_correct_detect()            
           if self.is_find_correct_aruco and id_aruco == self.id_correct_aruco:
               # print("send_cmd_vel")
               self.send_cmd_vel(self.aruco_info)
           else:
               move_cmd = Twist()
               self.pub.publish(Twist())

   #跟随物体
   def follow_obj(self):
       self.is_follow_obj_done = False

       while not self.is_follow_obj_done:
           img = self.cap.read()            
           img = self.rotate_image(img, -90)
           self.obj_detect(img)          
           if self.is_find:
               # print("send_cmd_vel")
               self.send_cmd_vel(self.target_image_info)
           else:
               move_cmd = Twist()
               self.pub.publish(Twist())


   #发送速度指令给小车，进行物体跟随
   def send_cmd_vel(self, info):

       x, y, xsize, ysize = info
       print(info)

       move_cmd = Twist()      

       if xsize > 130 or ysize > 130:
           self.pub.publish(Twist())
           self.is_follow_obj_done = True
       elif xsize > 80 or ysize > 80:
           move_cmd.linear.x = 0.05
           if abs(x) > 1:
               move_cmd.angular.z = x/self.cap.getWidth()
           self.pub.publish(move_cmd)
       else:
           move_cmd.linear.x = 0.15
           if abs(x) > 1:
               move_cmd.angular.z = x/self.cap.getWidth()
           self.pub.publish(move_cmd)
       self.rate.sleep()
       print(move_cmd.linear.x, move_cmd.angular.z)

   #寻找指定的物体
   def find_target_image(self):        
       while not self.is_find:  
           img = self.cap.read()            
           img = self.rotate_image(img, -90)
           self.obj_detect(img)  

   #机械臂到达准备抓取的位置
   def ready_arm_pose(self):
       basic.grap(False)
       if args.pose == "top":
           self.mc.send_coords(coords_top_ready,20,0)
           time.sleep(5)
       else:
           self.mc.send_coords(coords_top_ready,20,0)
           time.sleep(3)
           self.mc.send_coords(coords_bottom_ready,20,0)
           time.sleep(2)
   
   #机械臂进行抓取
   def move_and_pick(self):
       if args.pose == "top":
           self.mc.send_coords(coords_top_ready_ok,20,0)
           time.sleep(2)
           self.moveforward()
           self.mc.send_coords(coords_top_grap,20,0)
           time.sleep(3)
           basic.grap(True)#闭合夹爪

           self.mc.send_coords(coords_top_grap_ok,20,0)
           time.sleep(3)
       else:
           self.mc.send_coords(coords_bottom_ready_ok,20,0)
           time.sleep(2)
           self.moveforward()
           self.mc.send_coords(coords_bottom_grap,20,0)
           time.sleep(3)
           basic.grap(True)#闭合夹爪

           self.mc.send_coords(coords_bottom_grap_ok,20,0)
           time.sleep(3)
       self.moveback()

   #主函数
   def run(self):  

       self.mc.set_color(0,0,255)#blue, arm is busy
       
       #机械臂到达准备抓取的位置
       self.ready_arm_pose()

       self.is_find = False
       #寻找指定的物体
       self.find_target_image()    

     
       #time.sleep(10)#debug
       
       #跟随物体
       self.follow_obj()    
       
       #机械臂进行抓取
       self.move_and_pick()
       
       #机械臂放置积木到右侧
       self.place2right()

       cv.destroyAllWindows()
       self.mc.set_color(0,255,0)#green, arm is free


if __name__ == "__main__":    
   detect = ssdemo()
   detect.run()