PR #13967 from AviaAv: improve depth test

Author: Nir-Az

unit-tests/live/frames/test-depth.py

@@ -6,8 +6,9 @@
 import pyrealsense2 as rs
 from rspy import test, log
 from rspy import tests_wrapper as tw
+import numpy as np
 import os
-import math
+import time
 
 # Start depth + color streams and go through the frame to make sure it is showing a depth image
 # Color stream is only used to display the way the camera is facing
@@ -31,31 +32,35 @@
     cfg.enable_stream(rs.stream.color, rs.format.bgr8, 30)
 
 
-def display_image(img):
-    """
-    Display a given image and exits when the x button or esc key are pressed
-    """
-    import cv2
+pipeline = rs.pipeline(ctx)
+pipeline_profile = pipeline.start(cfg)
+pipeline.wait_for_frames()
+time.sleep(2)
 
-    window_title = "Output Stream"
-    cv2.imshow(window_title, img)
-    while cv2.getWindowProperty(window_title, cv2.WND_PROP_VISIBLE) > 0:
-        k = cv2.waitKey(33)
-        if k == 27:  # Esc key to stop
-            cv2.destroyAllWindows()
-            break
-        elif k == -1:  # normally -1 returned
-            pass
-
-
-def frames_to_image(depth, color, save, display, laser_enabled):
+def frames_to_image(depth, color, save, display):
     """
     This function gets depth and color frames, transforms them to an image (numpy array)
     and then save and/or display
 
     If color frame is given, it will also concatenate it with the depth frame before doing the given action
     """
-    import numpy as np
+
+    def display_image(img):
+        """
+        Display a given image and exits when the x button or esc key are pressed
+        """
+        import cv2
+
+        window_title = "Output Stream"
+        cv2.imshow(window_title, img)
+        while cv2.getWindowProperty(window_title, cv2.WND_PROP_VISIBLE) > 0:
+            k = cv2.waitKey(33)
+            if k == 27:  # Esc key to stop
+                cv2.destroyAllWindows()
+                break
+            elif k == -1:  # normally -1 returned
+                pass
+
     import cv2
 
     colorizer = rs.colorizer()
@@ -75,81 +80,43 @@ def frames_to_image(depth, color, save, display, laser_enabled):
         img = np.concatenate((depth_image, color_image), axis=1)
 
     if save:
-        file_name = f"output_stream{'_laser_on' if laser_enabled else '_laser_off'}.png"
+        file_name = f"output_stream.png"
         log.i("Saved image in", os.getcwd() + "\\" + file_name)
         cv2.imwrite(file_name, img)
     if display:
         display_image(img)
 
 
-def get_frames(config, laser_enabled):
-    pipeline = rs.pipeline(ctx)
-    pipeline_profile = pipeline.start(config)
-
-    sensor = pipeline_profile.get_device().first_depth_sensor()
-    if laser_enabled and sensor.supports(rs.option.laser_power):
-        sensor.set_option(rs.option.laser_power, sensor.get_option_range(rs.option.laser_power).max)
-    sensor.set_option(rs.option.emitter_enabled, 1 if laser_enabled else 0)
-
-    # to get a proper image, we sometimes need to wait a few frames, like when the camera is facing a light source
-    frames = pipeline.wait_for_frames()
-    for i in range(30):
-        frames = pipeline.wait_for_frames()
-    depth = frames.get_depth_frame()
-    color = frames.get_color_frame()
-    pipeline.stop()
-    return depth, color
-
+def get_distances(depth_frame):
+    MAX_METERS = 10  # max distance that can be detected, in meters
 
-def round_to_units(num):
-    """
-    Input: Distance of a certain point, in meters
-    Output: Distance according to the given detail level, in cm
-    """
-    in_cm = round(num, 2) * 100  # convert to cm
-    return math.floor(in_cm / DETAIL_LEVEL) * DETAIL_LEVEL  # rounds the distance according to the given unit
+    depth_m = np.asanyarray(depth_frame.get_data()).astype(np.float32) * depth_frame.get_units()
 
+    valid_mask = (depth_m < MAX_METERS)
+    valid_depths = depth_m[valid_mask]  # ignore invalid pixels
 
-def sort_dict(my_dict):
-    my_keys = list(my_dict.keys())
-    my_keys.sort()
-    sorted_dict = {i: my_dict[i] for i in my_keys}
-    return sorted_dict
+    # convert to cm and round according to DETAIL_LEVEL
+    rounded_depths = (np.floor(valid_depths * 100.0 / DETAIL_LEVEL) * DETAIL_LEVEL).astype(np.int32)
 
+    unique_vals, counts = np.unique(rounded_depths, return_counts=True)
+    
+    dists = dict(zip(unique_vals.tolist(), counts.tolist()))
+    total = valid_depths.size
 
-def get_distances(depth):
-    MAX_METERS = 10  # max distance that can be detected, in meters
-    dists = {}
-    total = 0
-    for y in range(depth.get_height()):
-        for x in range(depth.get_width()):
-            dist = depth.get_distance(x, y)
-            if dist >= MAX_METERS:  # out of bounds, assuming it is a junk value
-                continue
-
-            dist = round_to_units(dist)  # round according to DETAIL_LEVEL
-
-            if dists.get(dist) is not None:
-                dists[dist] += 1
-            else:
-                dists[dist] = 1
-            total += 1
-
-    dists = sort_dict(dists)  # for debug convenience
     log.d("Distances detected in frame are:", dists)
     return dists, total
 
 
-def is_depth_meaningful(config, laser_enabled=True, save_image=False, show_image=False):
+def is_depth_meaningful(save_image=False, show_image=False):
     """
     Checks if the camera is showing a frame with a meaningful depth.
     DETAIL_LEVEL is setting how close distances need to be, to be considered the same
 
     returns true if frame shows meaningful depth
     """
-
-    depth, color = get_frames(config, laser_enabled)
-
+    frames = pipeline.wait_for_frames()
+    depth = frames.get_depth_frame()
+    color = frames.get_color_frame()
     if not depth:
         log.f("Error getting depth frame")
         return False
@@ -160,32 +127,34 @@ def is_depth_meaningful(config, laser_enabled=True, save_image=False, show_image
 
     # save or display image (only possible through manual debugging)
     if save_image or show_image:
-        frames_to_image(depth, color, save_image, show_image, laser_enabled)
-
-    # Goes over the distances found, and checks if any distance is the same on more than 90% of the pixels
-    meaningful_depth = True
-    for key in dists:
-        if dists[key] > total*0.9:
-            meaningful_depth = False
-            break
+        frames_to_image(depth, color, save_image, show_image)
+
+    # If any distance is the same on more than 90% of the pixels, there is no meaningful depth
+    meaningful_depth = not any(v > total * 0.9 for v in dists.values())
     num_blank_pixels = dists[0]
     return meaningful_depth, num_blank_pixels
 
 
 ################################################################################################
 
 test.start("Testing depth frame - laser ON -", dev.get_info(rs.camera_info.name))
-is_there_depth = False
-max_black_pixels = float('inf')
+# enable laser
+sensor = pipeline_profile.get_device().first_depth_sensor()
+if sensor.supports(rs.option.laser_power):
+    sensor.set_option(rs.option.laser_power, sensor.get_option_range(rs.option.laser_power).max)
+sensor.set_option(rs.option.emitter_enabled, 1)  # should be set to 0 for laser off
+
+has_depth = False
 
-# we perform the check on a few different frames to make sure we get the best indication if we have depth
+# we check a few different frames to try and detect depth
 for frame_num in range(FRAMES_TO_CHECK):
-    result, laser_black_pixels = is_depth_meaningful(cfg, laser_enabled=True, save_image=DEBUG_MODE, show_image=DEBUG_MODE)
-    is_there_depth = is_there_depth or result  # we check if we found depth at any frame checked
-    max_black_pixels = min(max_black_pixels, laser_black_pixels)
+    has_depth, laser_black_pixels = is_depth_meaningful(save_image=DEBUG_MODE, show_image=DEBUG_MODE)
+    if has_depth:
+        break
 
-test.check(is_there_depth is True)
+test.check(has_depth is True)
 test.finish()
 
+pipeline.stop()
 tw.stop_wrapper( dev )
 test.print_results_and_exit()