| 48 | |
| 49 | |
| 50 | def prepare_traj_visual(full_pred_tracks, full_visibility_tracks, original_height, original_width, selected_frames, |
| 51 | dot_radius, target_width, target_height, name = ""): |
| 52 | |
| 53 | # Prepare the color |
| 54 | target_color_codes = all_color_codes[:len(full_pred_tracks[0])] # This means how many objects in total we have |
| 55 | invisible_color_code = (0, 0, 0) |
| 56 | |
| 57 | # Prepare the traj image |
| 58 | traj_img_lists = [] |
| 59 | |
| 60 | # Set a new dot radius based on the resolution fluctuating |
| 61 | dot_radius_resize = int( dot_radius * original_height / 384 ) # This is set with respect to default 384 height, will be adjust based on the height change |
| 62 | |
| 63 | |
| 64 | # Iterate all object instance |
| 65 | for temporal_idx, obj_points in enumerate(full_pred_tracks): # Iterate all downsampled frames, should be 13 |
| 66 | |
| 67 | # Init the base img for the traj figures |
| 68 | base_img = np.zeros((original_height, original_width, 3)).astype(np.float32) # Use the original image size |
| 69 | base_img.fill(255) # Whole white frames |
| 70 | |
| 71 | # Prepare visibility points |
| 72 | visible_masks = full_visibility_tracks[temporal_idx] |
| 73 | |
| 74 | |
| 75 | # Iterate for the per object |
| 76 | for obj_idx, points in enumerate(obj_points): |
| 77 | |
| 78 | # Basic setting |
| 79 | color_code = target_color_codes[obj_idx] # Color across frames should be consistent |
| 80 | visible_mask = visible_masks[obj_idx] |
| 81 | |
| 82 | # Process all points in this current object |
| 83 | for point_idx, (horizontal, vertical) in enumerate(points): |
| 84 | if horizontal < 0 or horizontal >= original_width or vertical < 0 or vertical >= original_height: |
| 85 | continue # If the point is already out of the range, Don't draw |
| 86 | |
| 87 | is_visible = visible_mask[point_idx] |
| 88 | |
| 89 | # Draw square around the target position |
| 90 | vertical_start = min(original_height, max(0, vertical - dot_radius_resize)) |
| 91 | vertical_end = min(original_height, max(0, vertical + dot_radius_resize)) # Diameter, used to be 10, but want smaller if there are too many points now |
| 92 | horizontal_start = min(original_width, max(0, horizontal - dot_radius_resize)) |
| 93 | horizontal_end = min(original_width, max(0, horizontal + dot_radius_resize)) |
| 94 | |
| 95 | # Paint |
| 96 | if is_visible: |
| 97 | base_img[vertical_start:vertical_end, horizontal_start:horizontal_end, :] = color_code |
| 98 | else: |
| 99 | base_img[vertical_start:vertical_end, horizontal_start:horizontal_end, :] = invisible_color_code |
| 100 | |
| 101 | |
| 102 | # Resize frames Don't use negative and don't resize in [0,1] |
| 103 | base_img = cv2.resize(base_img, (target_width, target_height), interpolation = cv2.INTER_CUBIC) |
| 104 | |
| 105 | # Dilate (Default to be True) |
| 106 | base_img = cv2.filter2D(base_img, -1, blur_kernel).astype(np.uint8) |
| 107 | |