Update yolo_tracker_save_hqsync.py

Add option to set auto focus range.

yolo_tracker_save_hqsync.py

@@ -30,6 +30,8 @@
 - optional arguments:
   "-4k"      crop detections from (+ save HQ frames in) 4K resolution (default: 1080p)
              -> decreases pipeline speed to ~3.4 fps (1080p: ~13.2 fps)
+  "-af"      set auto focus range in cm (min distance, max distance)
+             -> e.g. "-af 14 20" to restrict auto focus range to 14-20 cm
   "-ae"      use bounding box coordinates from detections to set auto exposure region
              -> can improve image quality of crops and thereby classification accuracy
   "-crop"    default:  save cropped detections with aspect ratio 1:1 ("-crop square") OR
@@ -74,6 +76,8 @@
 parser = argparse.ArgumentParser()
 parser.add_argument("-4k", "--four_k_resolution", action="store_true",
     help="crop detections from (+ save HQ frames in) 4K resolution (default: 1080p)")
+parser.add_argument("-af", "--af_range", nargs=2, type=int,
+    help="set auto focus range in cm (min distance, max distance)", metavar=("cm_min", "cm_max"))
 parser.add_argument("-ae", "--bbox_ae_region", action="store_true",
     help="use bounding box coordinates from detections to set auto exposure region")
 parser.add_argument("-crop", "--crop_bbox", choices=["square", "tight"], default="square", type=str,
@@ -261,8 +265,8 @@ def get_synced_frame(frames_list, tracks_seq):
 xout_tracker.setStreamName("track")
 script.outputs["track_out"].link(xout_tracker.input)  # synced tracker output
 
-if args.bbox_ae_region:
-    # Create XLinkIn node to send control commands to OAK device (e.g. set auto exposure region)
+if args.af_range or args.bbox_ae_region:
+    # Create XLinkIn node to send control commands to color camera node
     xin_ctrl = pipeline.create(dai.node.XLinkIn)
     xin_ctrl.setStreamName("control")
     xin_ctrl.out.link(cam_rgb.inputControl)
@@ -275,6 +279,31 @@ def frame_norm(frame, bbox):
     return (np.clip(np.array(bbox), 0, 1) * norm_vals).astype(int)
 
 
+def set_focus_range():
+    """Convert closest cm values to lens position values and set auto focus range."""
+    cm_lenspos_dict = {
+        6: 250,
+        8: 220,
+        10: 190,
+        12: 170,
+        14: 160,
+        16: 150,
+        20: 140,
+        25: 135,
+        30: 130,
+        40: 125,
+        60: 120
+    }
+
+    closest_cm_min = min(cm_lenspos_dict.keys(), key=lambda k: abs(k - args.af_range[0]))
+    closest_cm_max = min(cm_lenspos_dict.keys(), key=lambda k: abs(k - args.af_range[1]))
+    lenspos_min = cm_lenspos_dict[closest_cm_max]
+    lenspos_max = cm_lenspos_dict[closest_cm_min]
+
+    af_ctrl = dai.CameraControl().setAutoFocusLensRange(lenspos_min, lenspos_max)
+    q_ctrl.send(af_ctrl)
+
+
 def bbox_set_exposure_region(xmin_roi, ymin_roi, xmax_roi, ymax_roi):
     """Use bounding box coordinates from detections to set auto exposure region."""
     xmin_roi = max(0.001, xmin_roi)
@@ -343,6 +372,7 @@ def save_crop_metadata(frame, bbox):
         "y_max": round(ymax, 4),
         "file_path": path_crop
     }
+
     with open(save_path / f"metadata_{rec_start}.csv", "a", encoding="utf-8") as metadata_file:
         metadata = csv.DictWriter(metadata_file, fieldnames=[
             "rec_ID", "timestamp", "label", "confidence", "track_ID",
@@ -389,6 +419,7 @@ def save_logs():
         temp_oak = round(device.getChipTemperature().average)
     except RuntimeError:
         temp_oak = "NA"
+
     try:
         logs = {
             "rec_ID": rec_id,
@@ -405,6 +436,7 @@ def save_logs():
         }
     except IndexError:
         logs = {}
+
     with open(save_path.parent / f"info_log_{rec_start[:8]}.csv", "a", encoding="utf-8") as log_file:
         log_info = csv.DictWriter(log_file, fieldnames=[
             "rec_ID", "timestamp", "temp_pi", "temp_oak", "pi_mem_available", "pi_cpu_used",
@@ -423,6 +455,7 @@ def record_log():
         unique_ids = df_meta["track_ID"].nunique()
     except pd.errors.EmptyDataError:
         unique_ids = 0
+
     logs_rec = {
         "rec_ID": rec_id,
         "rec_start_date": rec_start[:8],
@@ -435,6 +468,7 @@ def record_log():
         "chargelevel_start": chargelevel_start,
         "chargelevel_end": chargelevel
     }
+
     with open(save_path.parents[1] / "record_log.csv", "a", encoding="utf-8") as log_rec_file:
         log_rec = csv.DictWriter(log_rec_file, fieldnames=[
             "rec_ID", "rec_start_date", "rec_start_time", "rec_end_time", "rec_time_min",
@@ -451,6 +485,7 @@ def save_zip():
     with ZipFile(f"{save_path.parent}.zip", "a") as zip_file:
         for file in save_path.rglob("*"):
             zip_file.write(file, file.relative_to(save_path.parent))
+
     shutil.rmtree(save_path.parent, ignore_errors=True)
 
 
@@ -471,10 +506,14 @@ def save_zip():
     q_frame = device.getOutputQueue(name="frame", maxSize=4, blocking=False)
     q_track = device.getOutputQueue(name="track", maxSize=4, blocking=False)
 
-    if args.bbox_ae_region:
-        # Create input queue to send commands to the OAK device (e.g. set auto exposure region)
+    if args.af_range or args.bbox_ae_region:
+        # Create input queue to send control commands to OAK camera
         q_ctrl = device.getInputQueue(name="control", maxSize=16, blocking=False)
 
+    if args.af_range:
+        # Set auto focus range to specified cm values
+        set_focus_range()
+
     # Set start time of recording and create empty list to save charge level (if < 10)
     start_time = time.monotonic()
     chargelevel_list = []