SmartTank/Collect.py

# 这段程序是用来采集图像并分类成三个文件夹的
import cv2
from urllib import request
import time
import pygame
import threading

cap = cv2.VideoCapture('http://192.168.31.100:8080/camera')
collecttime = time.time()
status = 'stop'
index = 0
done = False
speed = 20

# 控制子线程
def changeMotion(tank_state, tank_direct):
    global status
    if tank_state == 1 and tank_direct == 0:
        print("forward")
        print('send request')
        req = request.urlopen(
            'http://192.168.31.100:8080/control?direct=forward')
        res = req.read()
        print(res)
        status = 'forward'
    elif tank_state == -1 and tank_direct == 0:
        print("backward")
        print('send request')
        req = request.urlopen(
            'http://192.168.31.100:8080/control?direct=backward')
        res = req.read()
        print(res)
        status = 'stop'
    elif tank_state == 0:
        print("stop")
        print('send request')
        req = request.urlopen('http://192.168.31.100:8080/control?direct=stop')
        res = req.read()
        print(res)
        status = 'stop'
    elif tank_state == 1 and tank_direct == -1:
        print("left")
        print('send request')
        req = request.urlopen('http://192.168.31.100:8080/control?direct=left')
        res = req.read()
        print(res)
        status = 'left'
    elif tank_state == 1 and tank_direct == 1:
        print("right")
        print('send request')
        req = request.urlopen(
            'http://192.168.31.100:8080/control?direct=right')
        res = req.read()
        print(res)
        status = 'right'
    elif tank_state == -1 and tank_direct == -1:
        print("right")
        print('send request')
        req = request.urlopen(
            'http://192.168.31.100:8080/control?direct=right')
        res = req.read()
        print(res)
        status = 'right'
    elif tank_state == 1 and tank_direct == 1:
        print("left")
        print('send request')
        req = request.urlopen('http://192.168.31.100:8080/control?direct=left')
        res = req.read()
        print(res)
        status = 'left'


def control():
    supportControllers = ['Xbox One S Controller']
    forword_button = [(4, 6)]
    backword_button = [(5, 7)]
    start_button = [(7, 11)]
    # 模块初始化
    pygame.init()
    pygame.joystick.init()

    # 查找支持的手柄
    number = pygame.joystick.get_count()
    for i in range(0, number):
        j = pygame.joystick.Joystick(i)
        name = j.get_name()
        if name in supportControllers:
            print(name)
            search_index=supportControllers.index(name)
            print(search_index)
            f_b = forword_button[search_index]
            b_b = backword_button[search_index]
            s_b = start_button[search_index]
    joystick=pygame.joystick.Joystick(0)
    joystick.init()
    buttonCount = joystick.get_numbuttons()
    print(buttonCount)
    # XBox S手柄在蓝牙连接和USB连接时按钮数量不一样
    if len(f_b) > 1:
        if buttonCount == 11:
            b6_index = f_b[0]
            b7_index = b_b[0]
            b11_index = s_b[0]
        elif buttonCount == 17:
            b6_index = f_b[1]
            b7_index = b_b[1]
            b11_index = s_b[1]
    clock = pygame.time.Clock()
    # 0停止，1前，-1后
    tank_state = 0
    # 0无，-1左，1右
    tank_direct = 0
    while not done:
        for event_ in pygame.event.get():
            if event_.type == pygame.JOYHATMOTION:
                hat = joystick.get_hat(0)
                # 获取方向键状态
                tank_direct = hat[0]
                print("方向键："+str(tank_direct))
                changeMotion(tank_state, tank_direct)
            elif event_.type == pygame.JOYBUTTONDOWN or event_.type == pygame.JOYBUTTONUP:
                b6 = joystick.get_button(b6_index)
                b7 = joystick.get_button(b7_index)
                b11 = joystick.get_button(b11_index)
                print(str(b6)+" "+str(b7))
                if b11 == 1:
                    #设置初始速度
                    print('send request')
                    req = request.urlopen(
                        'http://192.168.31.100:8080/control?speed='+str(speed))
                    res = req.read()
                    print(res)
                if b6 == 1:
                    tank_state = -1
                    changeMotion(tank_state, tank_direct)
                if b7 == 1:
                    tank_state = 1
                    changeMotion(tank_state, tank_direct)
                if b6+b7 == 0:
                    tank_state = 0
                    changeMotion(tank_state, tank_direct)
    pygame.quit()


thread = threading.Thread(target=control)
thread.start()

# 显示主线程
while True:
    ret, frame = cap.read()
    cv2.imshow("capture",  frame)
    currenttime = time.time()
    if currenttime-collecttime > 2:
        if status == 'forward' or status == 'left' or status == 'right':
            cv2.imwrite('/home/xiongweixp/Documents/collect/' +
                        status+'/'+str(index)+'.jpg', frame)
            index = index+1
            collecttime = currenttime
    key = cv2.waitKey(1)
    if key == ord('q'):
        done = True
        break