#By: mengfei - 周日 1月 15 2023
import sensor, image, time,pyb,omv,math,utime,tf,lcd,gc
from pyb import UART,Pin,Timer,Servo
from umotor import UMOTOR
from pid import PID
from button import BUTTON
from centroidtracking import CentroidTracking
from pycommon import PYCOMMON
common=PYCOMMON()
objtraking=CentroidTracking(maxDisappeared=5) #目标追踪算法，质心追踪maxDisappeared表示最大丢失5帧代表目标消失
motor=UMOTOR()   #声明电机驱动，B0,B1,B4,B5
button=BUTTON()  #声明按键，梦飞openmv只有一个按键D8，因此直接内部指定了按键
objects = list()
last_objects=objects
lcd.init()
lcd.set_direction(2)
# 颜色跟踪阈值(L Min, L Max, A Min, A Max, B Min, B Max)
global THRESHOLD
THRESHOLD  =[ (58, 86, 7, 53, 33, 80),(22, 50, 21, 55, 37, 66)]  #黄色乒乓球# 可以根据白天夜晚光 多增加几组阈值，目前取了白天和晚上灯光下的场景阈值
s_last=0
flag_lost=0
#dis_pid  = PID(p=0.5, i=0.01)
#len_pid  = PID(p=0.25 ,i=0.01)
dis_pid  = PID(p=0.5, i=0.01)
len_pid  = PID(p=0.15 ,i=0.01)
#功能： 小球颜色识别
#输入： 图像
#输出： object_s,max_blob
def color_detect(img):
    DIS_RATE=27000          #测距离系数，使用图像像素多少估算距离
    #全局变量
    global last_objects,objects,flag_lost
    object_s=0
    max_blob=None
    if last_objects:  #如果上一张图像中找到了色块，就在色块周边扩展，然后在扩展后的ROI区域进行寻找色块，这样有利于追踪
         for b in objects:
            x1 = b[0]-7
            y1 = b[1]-7
            w1 = b[2]+12
            h1 = b[3]+12
         roi2 = (x1,y1,w1,h1)
         objects = img.find_blobs(THRESHOLD,roi = roi2,area_threshold=100,merge=True)
         last_objects = objects
    else:       #如果没有找到色块，就重新计算全图的色块，并将全图色块信息保存到上一张图像上，这样下一次就可以再进行ROI区域寻找色块
        objects = img.find_blobs(THRESHOLD,area_threshold=100,merge=True)
        last_objects = objects
    xrect=[]
    if objects:
        for d in objects :
            xrect.append(d.rect())
    objects=objtraking.update(xrect)
    #step1 有目标，就进行追踪和抓取操作
    if objects:  #如果找到了色块，就计算最大的色块位置，并做舵机追踪
        flag_lost=0
        object_s=0
        max_blob = common.find_max_object(objects)
        #色块的位置和大小需要满足图像本身大小的要求，以免越界
        object_s=DIS_RATE/(max_blob[2]*2) #计算距离
        img.draw_rectangle((max_blob[0],max_blob[1],max_blob[2],max_blob[3]))
        cx=int(max_blob[0]+max_blob[2]/2)
        cy=int(max_blob[1]+max_blob[3]/2)
        img.draw_cross(cx, cy) # cx, cy画中心点
        img.draw_string(cx,cy, "%.2f mm"%(object_s))  #显示目标的距离参数
    return object_s,max_blob

#功能： 自动找球控制
#输入： 图像，小球距离,最大面积的小球元组
#输出： 无
def find_ball_auto_control(img,mode=0):
    global flag_lost
    object_s=0
    object_s,max_blob=color_detect(img)
    if object_s>0:
        #step1 计算目标距离图像中心的偏差，做机械臂上下调整达到可抓取位置
        pan_error=0
        flag_lost=0
        pan_error = img.width()/2-(max_blob[2]/2+max_blob[0])
        power_s,power_l=common.car_traking(object_s,pan_error,400,30,80,len_pid,dis_pid)
        motor.run(power_s+power_l,power_s-power_l)
    else :
        motor.run(0,0) #停止等待目标出现
        flag_lost=flag_lost+1
        if flag_lost>5:#连续5帧没有
            flag_lost=0

##############################摄像头初始化部分#####################
sensor.reset() # Initialize the camera sensor.
sensor.set_pixformat(sensor.RGB565) # use RGB565.彩色图像
sensor.set_framesize(sensor.QCIF)
sensor.skip_frames(10) # Let new settings take affect.
sensor.set_auto_whitebal(False) # 除7670以外其他摄像头关闭白平衡做识别
clock = time.clock() # Tracks FPS. 设置一个定时器用来计算帧率

click_timer=time.ticks() #计时参数
while True:
    clock.tick()
    img = sensor.snapshot()
    find_ball_auto_control(img)
    lcd.display(img)
    THRESHOLD=common.auto_color_mask(img,lcd,THRESHOLD,15,30)
    print(clock.fps(), "fps", end="\n\n")
    if button.state():
        click_timer=time.ticks()          #开始计时
        while button.state():  pass       #等待按键抬起
        if time.ticks()-click_timer>2000: #按键时长超过2s
            break                         #循环退出，回到主界面
    else :
        click_timer=time.ticks()#计时更新