Data Logger w. on/off

Signed-off-by: Ken Haagh <[email protected]>

people_tracking/src/face_detection.py

@@ -83,6 +83,7 @@ def recognize_faces(self, input_images) -> Tuple[bool, Any]:
 
             for encoding in face_encoding:
                 matches = face_recognition.compare_faces(self.known_face_encodings, encoding, tolerance=0.6)  # True
+                # print(face_recognition.face_distance(self.known_face_encodings, encoding))
                 # if match, False if no match.
                 if any(matches):
                     return True, idx

people_tracking/src/people_tracker.py

@@ -24,6 +24,9 @@
 laptop = sys.argv[1]
 name_subscriber_RGB = 'video_frames' if laptop == "True" else '/hero/head_rgbd_sensor/rgb/image_raw'
 
+save_data = False if sys.argv[3] == "False" else True
+
+
 Persons = namedtuple("Persons",
                      ["nr_batch", "time", "nr_persons", "x_positions", "y_positions", "z_positions", "colour_vectors",
                       "face_detected"])
@@ -33,8 +36,9 @@
 
 class PeopleTracker:
     def __init__(self) -> None:
-        csv_file = open(csv_file_path, 'w', newline='')
-        self.csv_writer = csv.writer(csv_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
+        if save_data:
+            csv_file = open(csv_file_path, 'w', newline='')
+            self.csv_writer = csv.writer(csv_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
         # ROS Initialize
         rospy.init_node(NODE_NAME, anonymous=True)
         self.subscriber_hoc = rospy.Subscriber(TOPIC_PREFIX + 'HoC', ColourTarget, self.callback_hoc,
@@ -167,10 +171,10 @@ def callback_hoc(self, data: ColourCheckedTarget) -> None:
             self.detections[idx] = Persons(nr_batch, time, nr_persons, x_positions, y_positions, z_positions,
                                            colour_vectors, face_detected)
             self.new_detections = True
-
-            self.csv_writer.writerow([nr_batch, time, nr_persons,
-                                 x_positions, y_positions, z_positions,
-                                 colour_vectors, face_detected])
+            if save_data:
+                self.csv_writer.writerow([nr_batch, time, nr_persons,
+                                     x_positions, y_positions, z_positions,
+                                     colour_vectors, face_detected])
             rospy.loginfo(f"hoc: {nr_batch}")
 
 
@@ -194,9 +198,10 @@ def callback_face(self, data: ColourCheckedTarget) -> None:
             self.detections[idx] = Persons(nr_batch, time, nr_persons, x_positions, y_positions, z_positions,
                                            colour_vectors, face_detections)
             self.new_detections = True
-            self.csv_writer.writerow([nr_batch, time, nr_persons,
-                                      x_positions, y_positions, z_positions,
-                                      colour_vectors, face_detections])
+            if save_data:
+                self.csv_writer.writerow([nr_batch, time, nr_persons,
+                                          x_positions, y_positions, z_positions,
+                                          colour_vectors, face_detections])
             rospy.loginfo(f"face: {nr_batch}")
 
             current_timestamp = rospy.get_time()
@@ -225,9 +230,10 @@ def callback_persons(self, data: DetectedPerson) -> None:
         self.detections.append(
             Persons(nr_batch, time, nr_persons, x_positions, y_positions, z_positions, colour_vectors, face_detected))
         self.new_detections = True
-        self.csv_writer.writerow([nr_batch, time, nr_persons,
-                                  x_positions, y_positions, z_positions,
-                                  colour_vectors, face_detected])
+        if save_data:
+            self.csv_writer.writerow([nr_batch, time, nr_persons,
+                                      x_positions, y_positions, z_positions,
+                                      colour_vectors, face_detected])
         rospy.loginfo(f"pos: {nr_batch}")
         current_timestamp = rospy.get_time()
         if self.last_timestamp_da is not None:
@@ -441,7 +447,7 @@ def loop(self):
             current_time = rospy.get_time()
             if self.new_detections and current_time - time_old > 0.2:
                 time_old = current_time
-                # self.track_person()
+                self.track_person()
                 self.new_detections = False
                 rospy.loginfo(self.detections[-1])
 
@@ -468,17 +474,18 @@ def remove_outside_batches(lst: List, start_batch: int = 0, end_batch: int = flo
 import time
 
 if __name__ == '__main__':
-    try:
-        rospack = rospkg.RosPack()
-        package_path = rospack.get_path("people_tracking")
-        full_path = os.path.join(package_path, 'data/')
-
-        # Make sure the directory exists
-        os.makedirs(full_path, exist_ok=True)
-        time = time.ctime(time.time())
-        csv_file_path = os.path.join(full_path, f'{time}_test.csv')
-    except:
-        pass
+    if save_data:
+        try:
+            rospack = rospkg.RosPack()
+            package_path = rospack.get_path("people_tracking")
+            full_path = os.path.join(package_path, 'data/')
+
+            # Make sure the directory exists
+            os.makedirs(full_path, exist_ok=True)
+            time = time.ctime(time.time())
+            csv_file_path = os.path.join(full_path, f'{time}_test.csv')
+        except:
+            pass
 
     try:
         node_pt = PeopleTracker()

people_tracking/src/person_detection.py

@@ -19,6 +19,7 @@
 laptop = sys.argv[1]
 name_subscriber_RGB = 'video_frames' if laptop == "True" else '/hero/head_rgbd_sensor/rgb/image_raw'
 depth_camera = False if sys.argv[2] == "False" else True
+save_data = False if sys.argv[3] == "False" else True
 
 class PersonDetector:
     def __init__(self) -> None:
@@ -42,7 +43,6 @@ def __init__(self) -> None:
 
         self.bridge = CvBridge()
 
-
         # depth
         rospy.wait_for_service(TOPIC_PREFIX + 'depth/depth_data')
         self.depth_proxy = rospy.ServiceProxy(TOPIC_PREFIX + 'depth/depth_data', Depth)
@@ -74,6 +74,7 @@ def image_callback(self, data):
         # rospy.loginfo("%s, %s",data.header.seq, data.header.stamp.secs)
         self.latest_image_time = data.header.stamp.secs#float(rospy.get_time())
         self.batch_nr = data.header.seq
+
         # rospy.loginfo("rgb: %s t: %s",data.header.seq, data.header.stamp.secs)
 
     @staticmethod
@@ -110,6 +111,9 @@ def process_latest_image(self):
         cv_image = self.bridge.imgmsg_to_cv2(latest_image, desired_encoding='passthrough')
         cv_image = cv2.GaussianBlur(cv_image, (5, 5), 0)
 
+        # Save Image
+        if save_data:
+            cv2.imwrite(f"{save_path}{batch_nr}.jpg", cv_image)
         # People detection
         classes, segmentations, bounding_box_corners = self.detect(self.model, cv_image)
         if classes is None or segmentations is None:
@@ -199,7 +203,22 @@ def main_loop(self):
             rospy.sleep(0.001)
 
 
+import os
+import rospkg
+import time
+
 if __name__ == '__main__':
+    if save_data:
+        try:
+            rospack = rospkg.RosPack()
+            package_path = rospack.get_path("people_tracking")
+            time = time.ctime(time.time())
+            save_path = os.path.join(package_path, f'data/{time}_test/')
+            os.makedirs(save_path, exist_ok=True)           # Make sure the directory exists
+            print(save_path)
+        except:
+            pass
+
     try:
         print(f"Use Depth: {depth_camera}, Camera Source: {name_subscriber_RGB}")
         node_pd = PersonDetector()

people_tracking/test/people_tracking.launch

@@ -1,57 +1,58 @@
 <!--
     On HERO run the following command: (Sets correct camera source and turns on depth data)
-    $ roslaunch people_tracking people_tracking.launch laptop:=False depth_camera:=True
+    $ roslaunch people_tracking people_tracking.launch laptop:=False depth_camera:=True save_data:= False
 -->
 <launch>
   <arg name="laptop" default="True" />
   <arg name="depth_camera" default="False" />
+  <arg name="save_data" default="False"/>
 
   <node
     pkg="people_tracking"
     type="depth_image.py"
     name="depth"
     output="screen"
-    args="$(arg laptop) $(arg depth_camera)"
+    args="$(arg laptop) $(arg depth_camera) $(arg save_data)"
   />
 
   <node
     pkg="people_tracking"
     type="person_detection.py"
     name="person_detection"
     output="screen"
-    args="$(arg laptop) $(arg depth_camera)"
+    args="$(arg laptop) $(arg depth_camera) $(arg save_data)"
   />
 
   <node
     pkg="people_tracking"
     type="colour_check.py"
     name="HoC"
     output="screen"
-    args="$(arg laptop) $(arg depth_camera)"
+    args="$(arg laptop) $(arg depth_camera) $(arg save_data)"
   />
 
   <node
     pkg="people_tracking"
     type="face_detection.py"
     name="Face"
     output="screen"
-    args="$(arg laptop) $(arg depth_camera)"
+    args="$(arg laptop) $(arg depth_camera) $(arg save_data)"
   />
 
  <node
     pkg="people_tracking"
     type="people_tracker.py"
     name="people_tracker"
     output="screen"
-    args="$(arg laptop) $(arg depth_camera)"
+    args="$(arg laptop) $(arg depth_camera) $(arg save_data)"
   />
 
   <node
     pkg="people_tracking"
     type="webcam.py"
     name="webcam_pub"
     output="screen"
-    args="$(arg laptop) $(arg depth_camera)"
+    args="$(arg laptop) $(arg depth_camera) $(arg save_data)"
   />
 
 </launch>