Spaces:
Running
on
Zero
Running
on
Zero
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| def warp_image_with_flow(source_image, source_mask, target_image, flow) -> np.ndarray: | |
| """ | |
| Warp the target to source image using the given flow vectors. | |
| Flow vectors indicate the displacement from source to target. | |
| Args: | |
| source_image: np.ndarray of shape (H, W, 3), normalized to [0, 1] | |
| target_image: np.ndarray of shape (H, W, 3), normalized to [0, 1] | |
| flow: np.ndarray of shape (H, W, 2) | |
| source_mask: non_occluded mask represented in source image. | |
| Returns: | |
| warped_image: target_image warped according to flow into frame of source image | |
| np.ndarray of shape (H, W, 3), normalized to [0, 1] | |
| """ | |
| # assert source_image.shape[-1] == 3 | |
| # assert target_image.shape[-1] == 3 | |
| assert flow.shape[-1] == 2 | |
| # Get the shape of the source image | |
| height, width = source_image.shape[:2] | |
| target_height, target_width = target_image.shape[:2] | |
| # Create mesh grid | |
| x, y = np.meshgrid(np.arange(width), np.arange(height)) | |
| # Apply flow displacements | |
| flow_x, flow_y = flow[..., 0], flow[..., 1] | |
| x_new = np.clip(x + flow_x, 0, target_width - 1) + 0.5 | |
| y_new = np.clip(y + flow_y, 0, target_height - 1) + 0.5 | |
| x_new = (x_new / target_image.shape[1]) * 2 - 1 | |
| y_new = (y_new / target_image.shape[0]) * 2 - 1 | |
| warped_image = F.grid_sample( | |
| torch.from_numpy(target_image).permute(2, 0, 1)[None, ...].float(), | |
| torch.from_numpy(np.stack([x_new, y_new], axis=-1)).float()[None, ...], | |
| mode="bilinear", | |
| align_corners=False, | |
| ) | |
| warped_image = warped_image[0].permute(1, 2, 0).numpy() | |
| if source_mask is not None: | |
| warped_image = warped_image * (source_mask > 0.5) | |
| return warped_image | |
| def visualize_flow(flow, flow_scale): | |
| """ | |
| Visualize optical flow with direction modulating color and magnitude modulating saturation in HSV color space. | |
| Args: | |
| flow (np.ndarray): Flow array of shape (H, W, 2), where the first dimension | |
| represents (flow_x, flow_y). | |
| flow_scale (float): The scaling factor for the magnitude of the flow. | |
| Returns: | |
| np.ndarray: An RGB image visualizing the flow. | |
| """ | |
| # Convert CHW to HWC | |
| flow_hwc = flow | |
| # Compute the magnitude and angle of the flow | |
| magnitude = np.sqrt(np.square(flow_hwc[..., 0]) + np.square(flow_hwc[..., 1])) | |
| angle = np.arctan2(flow_hwc[..., 1], flow_hwc[..., 0]) # Angle in radians (-pi, pi) | |
| # Normalize the magnitude with the provided flow scale | |
| magnitude = magnitude / flow_scale | |
| magnitude = np.clip(magnitude, 0, 1) # Clip values to [0, 1] for saturation | |
| # Convert angle from radians to degrees (used for color hue in HSV) | |
| angle_deg = np.degrees(angle) % 360 # Convert angle to [0, 360] degrees | |
| # Create an HSV image: hue is based on angle, saturation on magnitude, and value is always 1 | |
| hsv_image = np.zeros((flow_hwc.shape[0], flow_hwc.shape[1], 3), dtype=np.uint8) | |
| hsv_image[..., 0] = angle_deg / 2 # OpenCV expects hue in range [0, 180] | |
| hsv_image[..., 1] = magnitude * 255 # Saturation in range [0, 255] | |
| hsv_image[..., 2] = 255 # Value always max (brightest) | |
| # Convert HSV image to RGB using OpenCV | |
| rgb_image = cv2.cvtColor(hsv_image, cv2.COLOR_HSV2BGR) | |
| return rgb_image | |