From b07c33f10a99917f2fa4f4cd44d6de9c9e2eea0f Mon Sep 17 00:00:00 2001 From: Mihir Patil Date: Wed, 25 May 2022 13:47:00 -0400 Subject: [PATCH] feat(training): Use video instead of single frames Video postprocessing can be done on a more powerful computer --- common/driving.py | 47 ++++++++++++++ training/processVideo.py | 26 ++++++++ training/requirements.txt | 1 + training/training_server.py | 126 +++++++++++++----------------------- 4 files changed, 118 insertions(+), 82 deletions(-) create mode 100644 common/driving.py create mode 100644 training/processVideo.py create mode 100644 training/requirements.txt diff --git a/common/driving.py b/common/driving.py new file mode 100644 index 0000000..660f715 --- /dev/null +++ b/common/driving.py @@ -0,0 +1,47 @@ +import pigpio + +# GPIO Pins for motor controller +drivePower = 4 +steerPower = 17 +driveBack = 27 +driveFwd = 22 +steerLeft = 23 +steerRight = 24 + + +# GPIO +pi = pigpio.pi() +pi.set_PWM_frequency(drivePower, 8000) +pi.write(steerPower, 1) + +# Throttle + +def stop(): + pi.clear_bank_1(1 << driveBack | 1 << driveFwd) + pi.set_PWM_dutycycle(drivePower, 0) + +def setThrottle(throttleVal): + if throttleVal == 0: + stop() + else: + pi.set_PWM_dutycycle(drivePower, abs(throttleVal) + 128) + if throttleVal >= 0: + pi.write(driveFwd, 1) + pi.write(driveBack, 0) + else: + pi.write(driveFwd, 0) + pi.write(driveBack, 1) + + + +# Steering +def left(): + pi.write(steerLeft, 1) + pi.write(steerRight, 0) + +def center(): + pi.clear_bank_1(1 << steerLeft | 1 << steerRight) + +def right(): + pi.write(steerLeft, 0) + pi.write(steerRight, 1) \ No newline at end of file diff --git a/training/processVideo.py b/training/processVideo.py new file mode 100644 index 0000000..89c34d0 --- /dev/null +++ b/training/processVideo.py @@ -0,0 +1,26 @@ +import ffmpeg +import shutil +from time import time_ns +import sys, os + +def process_video(start_time): + if not os.path.exists(f"raw/temp_{start_time}"): + os.makedirs(f"raw/temp_{start_time}") + print(f"raw/out_{start_time}.h265") + stream = (ffmpeg + .input(f"raw/out_{start_time}.h265") + .output(f"raw/temp_{start_time}/%05d.jpg")) + stream.run() + + with open(f"raw/commands_{start_time}.txt", "r") as commandFile: + lines = commandFile.readlines() + for i in range(len(lines)): + line = lines[i] + line = line[0:-1] + if line in ["stop", "", "kill"]: + continue + shutil.copyfile(f"raw/temp_{start_time}/{i:05}.jpg", + f"data/{line}/{int(time_ns())}.jpg") + +if __name__ == "__main__": + process_video(sys.argv[1]) \ No newline at end of file diff --git a/training/requirements.txt b/training/requirements.txt new file mode 100644 index 0000000..9abb303 --- /dev/null +++ b/training/requirements.txt @@ -0,0 +1 @@ +ffmpeg-python \ No newline at end of file diff --git a/training/training_server.py b/training/training_server.py index 0905701..e9e42d5 100644 --- a/training/training_server.py +++ b/training/training_server.py @@ -1,18 +1,10 @@ -import asyncio +import threading import bluetooth -import pigpio import depthai as dai -import ffmpeg -import os +import os, sys from time import time - -# GPIO Pins for motor controller -drivePower = 4 -steerPower = 17 -driveBack = 27 -driveFwd = 22 -steerLeft = 23 -steerRight = 24 +sys.path.insert(1, os.path.join(sys.path[0], '..')) +import common.driving as driving # Codes steer = 0 @@ -22,35 +14,26 @@ right = 2 # Directories - directories = ["data/left", "data/center", "data/right", "raw"] for directory in directories: if not os.path.exists(directory): os.makedirs(directory) -# GPIO -pi = pigpio.pi() -pi.set_PWM_frequency(drivePower, 8000) -pi.write(steerPower, 1) - # DepthAI pipeline = dai.Pipeline() -rgb = pipeline.create(dai.node.ColorCamera) -enc = pipeline.create(dai.node.VideoEncoder) -vidOut = pipeline.create(dai.node.XLinkOut) - -vidOut.setStreamName("video") +camRgb = pipeline.create(dai.node.ColorCamera) +videoEnc = pipeline.create(dai.node.VideoEncoder) +xout = pipeline.create(dai.node.XLinkOut) -rgb.setFps(60) -rgb.setBoardSocket(dai.CameraBoardSocket.RGB) -rgb.video.link(enc.input) -enc.bitstream.link(vidOut.input) -rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P) -enc.setDefaultProfilePreset( - rgb.getFps(), dai.VideoEncoderProperties.Profile.H265_MAIN) +xout.setStreamName('video') +camRgb.setBoardSocket(dai.CameraBoardSocket.RGB) +camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P) +videoEnc.setDefaultProfilePreset(60, dai.VideoEncoderProperties.Profile.H265_MAIN) +camRgb.video.link(videoEnc.input) +videoEnc.bitstream.link(xout.input) # Bluetooth server_sock = bluetooth.BluetoothSocket(bluetooth.RFCOMM) @@ -66,88 +49,67 @@ uuid, bluetooth.SERIAL_PORT_CLASS], profiles=[bluetooth.SERIAL_PORT_PROFILE]) -print("Waiting for connection on RFCOMM channel", port) +command = "stop" +def write_frame(queue, videoFile, commandFile): + global command + while True: + videoData = queue.get() + if command == "kill": + break + videoData.getData().tofile(videoFile) + commandFile.write(command+"\n") -async def write_frame(queue, videoFile, commandFile, command): - videoData = queue.tryGet() - if videoData is not None: - videoData.getData - commandFile.write(command) - -async def process_video(start_time): - stream = (ffmpeg - .input(f"raw/{start_time}.h265") - .output("raw/temp_{start_time}/%05d.jpg", "copy")) - stream.run() - - with open(f"raw/commands_{start_time}", "r") as commandFile: - lines = commandFile.readlines() - for i in range(len(lines)): - line = lines[i] - if line == "stop": - continue - os.rename(f"raw/temp_{start_time}/{i:05}.jpg", - f"data/{line}/{int(time()*1000)}.jpg") +driving.stop() with dai.Device(pipeline) as device: - videoQueue = device.getOutputQueue("video", maxSize=1, blocking=False) - + videoQueue = device.getOutputQueue("video", maxSize=30, blocking=False) try: while True: # Keep connecting to Bluetooth clients + print("Waiting for connection on RFCOMM channel", port) client_sock, client_info = server_sock.accept() print("Connected to", client_info) - currentCommand = "stop" start_time = int(time()*1000) with open(f"raw/out_{start_time}.h265", "wb") as videoOut, \ open(f"raw/commands_{start_time}.txt", "w") as commandOut: + thread = threading.Thread(target=write_frame,args=(videoQueue, videoOut, commandOut)) try: + thread.start() while True: data = client_sock.recv(1024) if not data: break - print(data) if data[0] == steer: - direction = "center" + command = "center" if data[1] == 0: - pi.clear_bank_1( - 1 << steerLeft | 1 << steerRight) + driving.center() elif data[1] == 2: - direction = "right" - pi.write(steerLeft, 0) - pi.write(steerRight, 1) + command = "right" + driving.right() else: - direction = "left" - pi.write(steerLeft, 1) - pi.write(steerRight, 0) - currentCommand = direction + command = "left" + driving.left() elif data[0] == throttle: throttleVal = int.from_bytes( data[1:2], "little", signed=True) - pi.set_PWM_dutycycle(drivePower, - abs(throttleVal) + 128) + driving.setThrottle(throttleVal) if throttleVal == 0: - pi.clear_bank_1(1 << driveBack | 1 << driveFwd) - pi.set_PWM_dutycycle(drivePower, 0) - print("stop") - currentCommand = "stop" - elif throttleVal >= 0: - pi.write(driveFwd, 1) - pi.write(driveBack, 0) - print("forward", abs(throttleVal)+128) - else: - pi.write(driveFwd, 0) - pi.write(driveBack, 1) - print("back", abs(throttleVal)+128) - asyncio.run(write_frame( - videoQueue, videoOut, commandOut, currentCommand)) - asyncio.run(process_video(start_time)) + command = "stop" + except OSError as err: print("Error:", err) finally: print("Disconnected") + command = "kill" + thread.join() client_sock.close() + driving.stop() + except KeyboardInterrupt: server_sock.close() + for file in os.scandir("raw"): + if file.is_file(): + os.chmod(file.path,os.stat.S) + driving.stop() \ No newline at end of file