feat(training): Use video instead of single frames

Video postprocessing can be done on a more powerful computer

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)

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

training/requirements.txt

@@ -0,0 +1 @@
+ffmpeg-python

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