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")

#遙控調(diào)用要關(guān)閉debug模式。debug模式會顯示識別視頻。
parser.add_argument("--debug", type=bool, ? default="True")

args = parser.parse_args()

#上層貨架抓取，機(jī)械臂的控制位置，可以通過注釋其他代碼進(jìn)行調(diào)試
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]

#下層貨架抓取，機(jī)械臂的控制位置，可以通過注釋其他代碼進(jìn)行調(diào)試
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):
? ?#初始化函數(shù)
? ?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,
? ? ? ?)
? ?
? ?#計算物體的坐標(biāo)和尺寸
? ?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 ?

? ?#計算物體的坐標(biāo)和尺寸
? ?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()
? ? ? ?#輸入數(shù)據(jù)處理
? ? ? ?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]

? ? ? ?
? ? ? ?#獲得識別結(jié)果
? ? ? ?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
? ?
? ?#小車右轉(zhuǎn)
? ?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
? ?
? ?#小車左轉(zhuǎn)
? ?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

? ?#機(jī)械臂放置積木到右側(cè)
? ?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)

? ?#旋轉(zhuǎn)圖像
? ?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())


? ?#發(fā)送速度指令給小車，進(jìn)行物體跟隨
? ?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) ?

? ?#機(jī)械臂到達(dá)準(zhǔn)備抓取的位置
? ?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)
? ?
? ?#機(jī)械臂進(jìn)行抓取
? ?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()

? ?#主函數(shù)
? ?def run(self): ?

? ? ? ?self.mc.set_color(0,0,255)#blue, arm is busy
? ? ? ?
? ? ? ?#機(jī)械臂到達(dá)準(zhǔn)備抓取的位置
? ? ? ?self.ready_arm_pose()

? ? ? ?self.is_find = False
? ? ? ?#尋找指定的物體
? ? ? ?self.find_target_image() ? ?

? ? ?
? ? ? ?#time.sleep(10)#debug
? ? ? ?
? ? ? ?#跟隨物體
? ? ? ?self.follow_obj() ? ?
? ? ? ?
? ? ? ?#機(jī)械臂進(jìn)行抓取
? ? ? ?self.move_and_pick()
? ? ? ?
? ? ? ?#機(jī)械臂放置積木到右側(cè)
? ? ? ?self.place2right()

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


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