#!/usr/bin/python3 #1
# coding=utf8 #2
# 第7章 ROS机器狗创意课程\4.AI视觉巡线行走(7.ROS Robot Creative Lesson\4.AI Visual Line Follow Walking) #3
import os #4
import sys #5
import cv2 #6
import math #7
import time #8
import rospy #9
import threading #10
import numpy as np #11
from threading import RLock, Timer #12
from ros_robot_controller.msg import RGBState, RGBsState #13
from std_srvs.srv import * #14
from sensor_msgs.msg import Image #15
from object_tracking.srv import * #16
from puppy_control.msg import Velocity, Pose, Gait #17
from large_models.srv import SetString #18

from common import Misc #20


ROS_NODE_NAME = 'visual_patrol_demo' #23

is_shutdown = False #25
start_time = 0 #26
elapsed_time = 0 #27

PuppyMove = {'x':0, 'y':0, 'yaw_rate':0} #29


color_range_list = {} #32

__isRunning = True #34
__target_color = '' #35
enable_running = False #36


line_centerx = -1 # 线条中心坐标(line center coordinates) #39
img_centerx = 320 #40

roi = [ # [ROI, weight] #42
        (240, 280,  0, 640, 0.1),  #43
        (320, 360,  0, 640, 0.2),  #44
        (400, 440,  0, 640, 0.7) #45
       ] #46
roi = [ # [ROI, weight] #47
        (120, 140,  0, 320, 0.1),  #48
        (160, 180,  0, 320, 0.2),  #49
        (200, 220,  0, 320, 0.7) #50
       ] #51

roi_h1 = roi[0][0] #53
roi_h2 = roi[1][0] - roi[0][0] #54
roi_h3 = roi[2][0] - roi[1][0] #55
roi_h_list = [roi_h1, roi_h2, roi_h3] #56

range_rgb = { #58
    'red': (0, 0, 255), #59
    'blue': (255, 0, 0), #60
    'green': (0, 255, 0), #61
    'black': (0, 0, 0), #62
    'white': (255, 255, 255), #63
} #64
draw_color = range_rgb["black"] #65


lock = RLock() #68


# 变量重置(variable reset) #71
def reset(): #72
    global draw_color #73
    global __target_color #74
    global color_range_list #75
    global line_centerx #76
    with lock: #77
        __target_color = 'None' #78
        line_centerx = -1 #79
        draw_color = range_rgb["black"] #80
        color_range_list = rospy.get_param('/lab_config_manager/color_range_list') #81

        PuppyMove['x'] = 0 #83
        PuppyMove['yaw_rate'] = math.radians(0) #84
        PuppyVelocityPub.publish(x=PuppyMove['x'], y=PuppyMove['y'], yaw_rate=PuppyMove['yaw_rate']) #85


# 初始位置(initial position) #88
def initMove(delay=True): #89
    global GaitConfig #90
    PuppyMove['x'] = 0 #91
    PuppyMove['yaw_rate'] = math.radians(0) #92
    PuppyVelocityPub.publish(x=PuppyMove['x'], y=PuppyMove['y'], yaw_rate=PuppyMove['yaw_rate']) #93
    rospy.sleep(0.2) #94
    rospy.ServiceProxy('/puppy_control/go_home', Empty)() #95
    PuppyPosePub.publish(stance_x=PuppyPose['stance_x'], stance_y=PuppyPose['stance_y'], x_shift=PuppyPose['x_shift'] #96
            ,height=PuppyPose['height'], roll=PuppyPose['roll'], pitch=PuppyPose['pitch'], yaw=PuppyPose['yaw'], run_time = 500) #97
    
    rospy.sleep(0.2) #99
    PuppyGaitConfigPub.publish(overlap_time = GaitConfig['overlap_time'], swing_time = GaitConfig['swing_time'] #100
                    , clearance_time = GaitConfig['clearance_time'], z_clearance = GaitConfig['z_clearance']) #101
                    
    with lock: #103
        pass #104
    if delay: #105
        rospy.sleep(0.5) #106


# app初始化调用(app initialization calling) #109
def init(): #110
    print("visual patrol Init") #111
    initMove(True) #112
    reset() #113
        
def move_robot(error): #115
    if abs(error) <= 20: #116
        PuppyMove['x'] = 10 #117
        PuppyMove['yaw_rate'] = math.radians(0) #118
    elif error > 50: #119
        PuppyMove['x'] = 8 #120
        PuppyMove['yaw_rate'] = math.radians(-15) #121
    elif error < -50: #122
        PuppyMove['x'] = 8 #123
        PuppyMove['yaw_rate'] = math.radians(15) #124
    else: time.sleep(0.02) #125
    PuppyVelocityPub.publish(x=PuppyMove['x'], y=PuppyMove['y'], yaw_rate=PuppyMove['yaw_rate']) #126
        
               
        
           
  
# 找出面积最大的轮廓(find out the contour with the maximal area) #132
# 参数为要比较的轮廓的列表(the parameter is the list of contour to be compared) #133
def getAreaMaxContour(contours): #134
    contour_area_temp = 0 #135
    contour_area_max = 0 #136
    area_max_contour = None #137

    for c in contours:  # 历遍所有轮廓(iterate through all contours) #139
        contour_area_temp = math.fabs(cv2.contourArea(c))  # 计算轮廓面积(calculate the contour area) #140
        if contour_area_temp > contour_area_max: #141
            contour_area_max = contour_area_temp #142
            if contour_area_temp > 50:  # 只有在面积大于50时，最大面积的轮廓才是有效的，以过滤干扰(only contours with an area greater than 50 are considered valid, filtering out interference) #143
                area_max_contour = c #144

    return area_max_contour, contour_area_max  # 返回最大的轮廓(return the maximal contour) #146
def run(img): #147
    global draw_color, line_centerx, start_time, elapsed_time, enable_running #148
    size = (320, 240) #149
   
    img_h, img_w = img.shape[:2] #151

    if not __isRunning: #153
        return img #154

    frame_resize = cv2.resize(img, size, interpolation=cv2.INTER_NEAREST) #156
    frame_gb = cv2.GaussianBlur(frame_resize, (3, 3), 3)       #157
    frame_lab = cv2.cvtColor(frame_gb, cv2.COLOR_BGR2LAB)  # 将图像转换到LAB空间(convert the image to LAB space) #158


    centroid_x_sum = 0 #161
    weight_sum = 0 #162
    center_ = [] #163
    n = 0 #164
    #将图像分割成上中下三个部分，这样处理速度会更快，更精确(divide the image to upper, middle and lower parts, which will result in faster and more accurate processing) #165
    for r in roi: #166
        roi_h = roi_h_list[n] #167
        n += 1        #168
        if n <= 2: #169
            continue #170
        blobs = frame_gb[r[0]:r[1], r[2]:r[3]] #171
        frame_lab = cv2.cvtColor(blobs, cv2.COLOR_BGR2LAB)  # 将图像转换到LAB空间(convert the image to LAB space) #172
        

        for i in color_range_list: #175
            if i in __target_color: #176
                detect_color = i #177
                
                frame_mask = cv2.inRange(frame_lab, #179
                                             (color_range_list[detect_color]['min'][0], #180
                                              color_range_list[detect_color]['min'][1], #181
                                              color_range_list[detect_color]['min'][2]), #182
                                             (color_range_list[detect_color]['max'][0], #183
                                              color_range_list[detect_color]['max'][1], #184
                                              color_range_list[detect_color]['max'][2]))  #对原图像和掩模进行位运算(perform bitwise operation to original image and mask) #185
                opened = cv2.morphologyEx(frame_mask, cv2.MORPH_OPEN, np.ones((6, 6), np.uint8))  # 开运算(opening operation) #186
                closed = cv2.morphologyEx(opened, cv2.MORPH_CLOSE, np.ones((6, 6), np.uint8))  # 闭运算(closing operation) #187
              
        if __target_color == '' or __target_color == 'None' or __target_color == None: #189
            line_centerx = -1 #190
            return img #191
        cnts = cv2.findContours(closed , cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)[-2]#找出所有轮廓(find out all the contours) #192
        cnt_large, area = getAreaMaxContour(cnts)#找到最大面积的轮廓(find the contour with the maximal area) #193
        if cnt_large is not None:#如果轮廓不为空(if the contour is not none) #194
            rect = cv2.minAreaRect(cnt_large)#最小外接矩形(the minimal bounding rectangle) #195
            box = np.intp(cv2.boxPoints(rect))#最小外接矩形的四个顶点(the four vertices of the minimum bounding rectangle) #196
            for i in range(4): #197
                box[i, 1] = box[i, 1] + (n - 1)*roi_h + roi[0][0] #198
                box[i, 1] = int(Misc.map(box[i, 1], 0, size[1], 0, img_h)) #199
            for i in range(4):                 #200
                box[i, 0] = int(Misc.map(box[i, 0], 0, size[0], 0, img_w)) #201
                
            cv2.drawContours(img, [box], -1, (0,0,255,255), 2)#画出四个点组成的矩形(draw a rectangle formed by connecting the four points) #203
            
            #获取矩形的对角点(get the diagonal points of the rectangle) #205
            pt1_x, pt1_y = box[0, 0], box[0, 1] #206
            pt3_x, pt3_y = box[2, 0], box[2, 1]             #207
            center_x, center_y = (pt1_x + pt3_x) / 2, (pt1_y + pt3_y) / 2#中心点(center point) #208
            cv2.circle(img, (int(center_x), int(center_y)), 5, (0,0,255), -1)#画出中心点(draw the center point) #209
            
            center_.append([center_x, center_y])                         #211
            #按权重不同对上中下三个中心点进行求和(summing up the three central points of the top, middle, and bottom with different weights) #212
            centroid_x_sum += center_x * r[4] #213
            weight_sum += r[4] #214

    if weight_sum != 0: #216
        #求最终得到的中心点(calculate the final resulting center point) #217
        cv2.circle(img, (line_centerx, int(center_y)), 10, (0,255,255), -1)#画出中心点(draw the center point) #218
        line_centerx = int(centroid_x_sum / weight_sum)   #219
        print('line_centerx',line_centerx) #220
        start_time = 0 #221
    else: #222
        if start_time == 0: #223
            start_time = time.time() #224
        elapsed_time = time.time() - start_time #225
        if elapsed_time > 3: #226
            rospy.loginfo("Weight sum is zero for more than 3 seconds. Terminating program.") #227
            go_home_service() #228
            enable_running = False   #229
            line_centerx = -1 #230
        else: #231
            line_centerx = -1 #232

    return img #234

def image_callback(ros_image): #236
    global lock, enable_running, __isRunning, line_centerx, img_centerx #237
    
    image = np.ndarray(shape=(ros_image.height, ros_image.width, 3), dtype=np.uint8, #239
                       buffer=ros_image.data)  # 将自定义图像消息转化为图像(convert the customize image information to image) #240
    frame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) #241
    
    frame_result = frame #243
    with lock: #244
        if enable_running: #245
            frame_result = run(frame) #246
            if __isRunning and line_centerx != -1: #247
                error = line_centerx - img_centerx #248
                move_robot(error) #249
            elif __isRunning and line_centerx == -1: #250
                PuppyVelocityPub.publish(x=0, y=0, yaw_rate=0) #251
            
            
            cv2.imshow('Frame', frame_result) #254
            key = cv2.waitKey(1) #255
        else: #256
            cv2.destroyAllWindows() #257

def enter_func(msg): #259
    init() #260
    rospy.loginfo("enter visual_patrol_demo") #261
    rospy.sleep(0.02) #262
    
    response = TriggerResponse() #264
    response.success = True #265
    response.message = "start" #266
    return response #267
   
    
def set_running(msg): #270
    rospy.loginfo("enter visual_patrol_demo") #271
    global enable_running #272
    
    if msg.data: #274
        enable_running = msg.data #275
    else: #276
        enable_running = False #277
    
    response = SetBoolResponse() #279
    response.success = True #280
    return response #281
    
    
def set_target(msg): #284
    global __target_color #285
    
    rospy.loginfo("%s", msg) #287
    __target_color = msg.data #288
    
    response = SetBoolResponse() #290
    response.success = True #291
        
    return [True, 'set_target'] #293


def cleanup(): #296
    global is_shutdown #297
    is_shutdown = True #298
    PuppyVelocityPub.publish(x=0, y=0, yaw_rate=0) #299
    print('is_shutdown') #300
if __name__ == '__main__': #301
    rospy.init_node(ROS_NODE_NAME, log_level=rospy.INFO) #302
    
    
    #PP = rospy.get_param('/puppy_control/PuppyPose') #305
    param_name = '/puppy_control/PuppyPose' #306
    while not rospy.has_param(param_name): #307
        rospy.logwarn("parameter %s not available yet. Waiting....",param_name) #308
        rospy.sleep(1) #309
    PP = rospy.get_param('/puppy_control/PuppyPose') #310
    
    PuppyPose = PP['LookDown_20deg'].copy() #312
    
    GaitConfig = {'overlap_time':0.1, 'swing_time':0.15, 'clearance_time':0.0, 'z_clearance':3} #314
    
    image_pub = rospy.Publisher('/%s/image_result'%ROS_NODE_NAME, Image, queue_size=1)  # register result image publisher #316
    rgb_pub = rospy.Publisher('/ros_robot_controller/set_rgb', RGBsState, queue_size=1) #317
    image_sub = rospy.Subscriber('/usb_cam/image_raw', Image, image_callback) #318

    PuppyGaitConfigPub = rospy.Publisher('/puppy_control/gait', Gait, queue_size=1) #320
    PuppyVelocityPub = rospy.Publisher('/puppy_control/velocity', Velocity, queue_size=1) #321
    PuppyPosePub = rospy.Publisher('/puppy_control/pose', Pose, queue_size=1) #322
    rospy.sleep(0.2) #323
    
    go_home_service = rospy.ServiceProxy('/puppy_control/go_home', Empty) #325
    enter_srv = rospy.Service('/%s/enter'%ROS_NODE_NAME, Trigger, enter_func) #326
    running_srv = rospy.Service('/%s/enable_running'%ROS_NODE_NAME, SetBool, set_running) #327
    set_target_srv = rospy.Service('/%s/set_color_target'%ROS_NODE_NAME, SetString, set_target) #328
    rospy.on_shutdown(cleanup) #329
    
    __target_color = 'red' #331
    
    try: #333
        rospy.spin() #334
    except KeyboardInterrupt: #335
        print("Shutting down") #336
    finally: #337
        cv2.destroyAllWindows() #338
