diff --git a/utils/helpers.py b/utils/helpers.py index e4bf89a..17a221e 100644 --- a/utils/helpers.py +++ b/utils/helpers.py @@ -91,22 +91,20 @@ def get_dataloader(params, mode="train"): ) return data_loader - -def draw_gaze(frame, bbox, pitch, yaw, thickness=2, color=(0, 0, 255)): - """Draws gaze direction on a frame given bounding box and gaze angles.""" - # Unpack bounding box coordinates - x_min, y_min, x_max, y_max = map(int, bbox[:4]) - - # Calculate center of the bounding box - x_center = (x_min + x_max) // 2 - y_center = (y_min + y_max) // 2 - +def draw_gaze(frame, pitch, yaw, thickness=2, color=(0, 0, 255), length_scale=0.6): + """Draws a longer gaze direction on a frame given gaze angles (pitch and yaw).""" + # Handle grayscale frames by converting them to BGR if len(frame.shape) == 2 or frame.shape[2] == 1: frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR) + # Calculate the center of the image (assuming it's a face image) + h, w = frame.shape[:2] + x_center = w // 2 + y_center = h // 2 + # Calculate the direction of the gaze - length = x_max - x_min + length = int(w * length_scale) # Scale the arrow length based on image width dx = int(-length * np.sin(pitch) * np.cos(yaw)) dy = int(-length * np.sin(yaw)) @@ -122,10 +120,11 @@ def draw_gaze(frame, bbox, pitch, yaw, thickness=2, color=(0, 0, 255)): color=color, thickness=thickness, line_type=cv2.LINE_AA, - tipLength=0.25 + tipLength=0.3 # Adjust tip length for better visibility ) + def draw_bbox(image, bbox, color=(0, 255, 0), thickness=2, proportion=0.2): x_min, y_min, x_max, y_max = map(int, bbox[:4])