modularize and add documentations

Code/Control/Laptop_Code/ball_detection/ball_detection.py

@@ -52,9 +52,21 @@ def returnModel(device, labelSet, modelLoc, modelFile):
 
 #### Live Detection
 def detectLive(model, minDetectionScore = 0.90, showSight = True):
-    gbx  = 0
-    gby  = 0
-    dist = 0
+    '''
+    Descrption:
+        ball detection ML Functions
+
+    Input:
+        the ML modelLoc
+
+    Output:
+        -1 means not found
+        otherwise return the actual ball information
+    '''
+
+    gbx  = -1
+    gby  = -1
+    dist = -1
 
     header = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/92.0.4515.159 Safari/537.36."}
     req      = Request(url, headers = header)
@@ -89,7 +101,8 @@ def detectLive(model, minDetectionScore = 0.90, showSight = True):
     if showSight:
       cv2.imshow('Ball Detection', frame)
       key = cv2.waitKey(1) & 0xFF
-    
+
+    dist = int(dist)
     return gbx, gby, dist
 
 ### Supplementary Functions
@@ -154,4 +167,4 @@ def depth_estimate(figDims, boundingBox):
 colorLow    =  60
 colorHigh   = 180
 cropFactor  =   0.90
-distWtsFile = './ball_detection/distance-detection-torch/distance-detection-weights-3-2.0sd-20210808_212550.json'
\ No newline at end of file
+distWtsFile = './ball_detection/distance-detection-torch/distance-detection-weights-3-2.0sd-20210808_212550.json'

Code/Control/Laptop_Code/main.py

@@ -5,160 +5,135 @@ import ball_detection.ball_detection as ball_detection
 from constants import *
 
 
-# for ESP 32 Serial IO
-PORT = "COM5" # for Alienware
-serial_port = serial.Serial(PORT,115200)
-serial_port.close()
-serial_port.open()
+def serial_port_in(serial_port):
+    '''
+    Description:
+        Take all ESP32_Master serial port's printIn and take all necessary input object
+
+    Input:
+        serial_port     :    serial.Serail object
+
+    Output:
+        tx, ty, tz, rx, ry, rz, AT_dist, DebugM
+    '''
+    line = serial_port.readline()
+    if line == b'SERIAL_IN_START\r\n':
+        # DEBUG Verbose
+        print("initiating one round of serial in ...")
+
+        for i in range(7):
+            line = serial_port.readline()
+            val = int(line.decode())
+
+            if   i == 0:    tx = val
+            elif i == 1:    ty = val
+            elif i == 2:    tz = val
+            elif i == 3:    rx = val
+            elif i == 4:    ry = val
+            elif i == 5:    rz = val
+            elif i == 6:    AT_dist = val
+
+        line = serial_port.readline()
+        debugM = line.decode()
+
+    # DEBUG Verbose
+    print("tx:{}".format(tx))
+    print("ty:{}".format(ty))
+    print("tz:{}".format(tz))
+    print("rx:{}".format(rx))
+    print("ry:{}".format(ry))
+    print("rz:{}".format(rz))
+    print("dist:{}".format(AT_dist))
+    print(debugM)
+
+    return tx, ty, tz, rx, ry, rz, AT_dist, DebugM
+
+
+
+def serial_port_out(serial_port, pwm1, pwm2, pwm3, pwm4, dir1, dir2, dir3, dir4):
+    '''
+    Description:
+        Feed to ESP32_Master to send ESP32_Slave necessary information
 
+    Input:
+        serial_port                                     :   serial.Serail object
+        pwm1, pwm2, pwm3, pwm4, dir1, dir2, dir3, dir4  :   variables to send
 
-gbc  = 0
-gbx  = 0
-gby  = 0
-dist = 0
+    Output:
+        None
+    '''
 
+    output_message = str(pwm1) + str(pwm2) + str(pwm3) + str(pwm4) + dir1 + dir2 +  dir3 +  dir4 + '\n'
+    serial_port.write(output_message.encode())
 
-waitTime = 0.10
 
-model = ball_detection.returnModel(device, labelSet, modelLoc, modelFile)
 
+def main_control(gbx, gby, gb_dist, tx, ty, tz, rx, ry, rz, AT_dist, debugM):
+    '''
+    Description:
+        Given green ball information and AT information, the main control logic
+        to manuver the blimp motors
 
-def goal_detect(R):
-  goalDetect = 1
-  if R[1] == 0 and R[2] == 0:
-    goalDetect = 0
+    Input:
+        gbx, gby, gb_dist                   :   green ball information
+        tx, ty, tz, rx, ry, rz, AT_dist     :   AirTag information
+        debugM                              :   Debug Message
 
+    Output:
+        pwm1, pwm2, pwm3, pwm4, dir1, dir2, dir3, dir4  :   Blimp motor manuver parameters
+    '''
 
-def ball_capture(d_LIDAR):
-  if (d_LIDAR < LIDAR_Thres) and (d_LIDAR > 0):  # Ball captured
-    return 1
-  return 0
+    pass
+    pwm1, pwm2, pwm3, pwm4 = 0 , 0 , 0 , 0   # Not moving
+    dir1, dir2, dir2, dir2 = '+', '+', '+', '+'
 
+    return pwm1, pwm2, pwm3, pwm4, dir1, dir2, dir3, dir4
 
-def main(GB, T, R, d_LIDAR, Debug):
-  # ballDetect = cam.exe_color_detection_biggestblob(threshold[0], threshold[1], threshold[2], threshold[3], threshold[4], threshold[5], x, y)
-  ballDetect  = ball_detect(GB[0], GB[1])
-  ballCapture = ball_capture(d_LIDAR)
-  goalDetect  = goal_detect(R)
 
-  if ballCapture:   # Ball captured
-    if goalDetect:  # Goal detected
-      pass
-    else:  # Goal not detected
-      pass
-  else:  # Ball not captured
-    if ballDetect:  # Ball detected
-      pass
-    else:  # Ball not detected
-      pass
 
+# ===== Main Functions =====
 
 if __name__ == '__main__':
+    # =========== SET UP ============
+    # Defining Variables for ESP 32 Serial I/O
+    PORT = "COM5" # for Alienware
+    serial_port = serial.Serial(PORT, 115200)
+    serial_port.close()
+    serial_port.open()
 
-  while True:
-    #gbx, gby, dist = ball_detection.detectLive(model, minDetectionScore, showSight = True)
-    #dist = int(dist) # cm
+    # Weit Time
+    waitTime = 0.10
 
-    line = serial_port.readline()
-    if line == b'SERIAL_IN_START\r\n':
-        print("initiating one round of serial in ...")
+    # Loading the PyTorch ML model for ball detection
+    model = ball_detection.returnModel(device, labelSet, modelLoc, modelFile)
 
-        for i in range(8):
-            if i < 7:
-                line = serial_port.readline()
-                value = line.decode()
-                val = int(value)
-                if i == 0:
-                    tx = val
-                elif i == 1:
-                    ty = val
-                elif i == 2:
-                    tz = val
-                elif i == 3:
-                    rx = val
-                elif i == 4:
-                    ry = val
-                elif i == 5:
-                    rz = val
-                elif i == 6:
-                    dist = val
-            else:
-                line = serial_port.readline()
-                debugM = line.decode()
-
-        print("tx:{}".format(tx))
-        print("ty:{}".format(ty))
-        print("tz:{}".format(tz))
-        print("rx:{}".format(rx))
-        print("ry:{}".format(ry))
-        print("rz:{}".format(rz))
-        print("dist:{}".format(dist))
-        print(debugM)
-
-    time.sleep(waitTime)
-
-    pwm1 = 225
-    pwm2 = 225
-    pwm3 = 225
-    pwm4 = 225
-    dir1 = '+'
-    dir2 = '-'
-    dir3 = '+'
-    dir4 = '-'
 
-    output_message = str(pwm1) + str(pwm2) + str(pwm3) + str(pwm4) + dir1 + dir2 +  dir3 +  dir4 + '\n'
-    serial_port.write(output_message.encode())
+    # =========== DECLARE VARIABLES ===========
+    # ESP CAM In
+    gbx, gby  = -1, -1   # by default (-1 means no found green ball)
+    gb_dist = -1         # by default (-1 means no found green ball)
 
-    #main(GB, T, R, d_LIDAR, Debug)
+    # Serial Port In
+    tx, ty, tz = -1, -1, -1  # by default (-1 means no found AirTag)
+    rx, ry, rz = -1, -1, -1
+    AT_dist = -1
+    debugM = 'Testing'
 
-    '''
-    if gbx != 0 and gby != 0:
-        """
-        value = 'gbx' + str(gbx) + '\n'
-        print( value , end = '')
-        serial_port.write(value.encode())
-        time.sleep(waitTime)
+    # Serial Port Out
+    pwm1, pwm2, pwm3, pwm4 = 0 , 0 , 0 , 0   # Not moving
+    dir1, dir2, dir2, dir2 = '+', '+', '+', '+'
 
-        value = 'gby' + str(gby) + '\n'
-        print( value)
-        serial_port.write(value.encode())
-        time.sleep(waitTime)
-        """
-        if gbx < 100:
-            strgbx = '0' + str(gbx)
-        else:
-            strgbx = str(gbx)
-        if gby < 100:
-            strgby = '0' + str(gby)
-        else:
-            strgby = str(gby)
-        if dist < 100:
-            strdist = '0' + str(dist)
-        else:
-            strdist = str(dist)
-
-
-        value = 'gbc' + strgbx + strgby + strdist + '\n'
-        # value = 'gbc' + str(dist) + '\n'
-        print(value)
-        serial_port.write(value.encode())
-        time.sleep(waitTime)
 
-    else:
-        """
-        value = 'gbx-1\n'
-        serial_port.write(value.encode())
+    # =========== LOOP FOREVER===========
+    while True:
+        # ===== STEP 1: TAKE ALL INPUT =====
+        gbx, gby, gb_dist = ball_detection.detectLive(model, minDetectionScore, showSight = True)
+        tx, ty, tz, rx, ry, rz, AT_dist, debugM = serial_port_in(serial_port)
         time.sleep(waitTime)
 
-        value = 'gby-1\n'
-        serial_port.write(value.encode())
-        """
-        strgbx = '000'
-        strgby = '000'
-        strdist = '000'
-        value = 'gbc' + strgbx + strgby + strdist + '\n'
-        print(value)
-        serial_port.write(value.encode())
+        # ===== STEP 2: MAIN CONTROL LOOP =====
+        pwm1, pwm2, pwm3, pwm4, dir1, dir2, dir3, dir4 = main_control(gbx, gby, gb_dist, tx, ty, tz, rx, ry, rz, AT_dist, debugM)
+
+        # ===== STEP 3: FEED ALL OUTPUT =====
+        serial_port_out(serial_port, pwm1, pwm2, pwm3, pwm4, dir1, dir2, dir3, dir4)
         time.sleep(waitTime)
-        print("no green ball detected")
-    '''