upload files

README.md

@@ -1,10 +1,10 @@
 ---
-title: AllTracker PointVersion
-emoji: 🔥
-colorFrom: purple
-colorTo: purple
+title: DenseTrack
+emoji: 🏃
+colorFrom: pink
+colorTo: pink
 sdk: gradio
-sdk_version: 5.38.2
+sdk_version: 5.21.0
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,328 @@
+import os
+import random
+import time
+import datetime
+import cv2
+import numpy as np
+import torch
+import imageio
+from pathlib import Path
+from tqdm import tqdm
+import gradio as gr
+
+# Import your custom modules
+import utils.loss
+import utils.samp
+import utils.data
+import utils.improc
+import utils.misc
+import utils.saveload
+from nets.blocks import InputPadder
+from nets.net34 import Net
+from tensorboardX import SummaryWriter
+import imageio
+from demo_dense_visualize import Tracker
+import spaces
+
+# -------------------- Utility Functions -------------------- #
+def count_parameters(model):
+    total_params = 0
+    for name, parameter in model.named_parameters():
+        if not parameter.requires_grad:
+            continue
+        total_params += parameter.numel()
+    print('Total params: %.2f M' % (total_params/1e6))
+    return total_params
+
+def seed_everything(seed: int):
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+# -------------------- Step 1: Extract the First Frame -------------------- #
+def extract_first_frame(video_path, _, tracker):
+    """
+    Opens the video, extracts the first frame, resizes it (largest dimension 1024),
+    and returns:
+      - the frame for display (to be annotated),
+      - the video file path (to store in state),
+      - a copy of the original first frame (to be used when adding points)
+    """
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        return None, None, None
+    ret, frame = cap.read()
+    cap.release()
+    if not ret:
+        return None, video_path, None
+    # Convert BGR to RGB
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    scale = min(tracker.target_res / frame_rgb.shape[0], tracker.target_res / frame_rgb.shape[1])
+    frame_resized = cv2.resize(frame_rgb, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
+    # Return the displayed frame, the video file path, and a copy of the original frame for point drawing.
+    return frame_resized, video_path, frame_resized.copy(), []
+
+# -------------------- Callback to Add a Clicked Point -------------------- #
+def add_point(orig_frame, points, evt: gr.SelectData):
+    """
+    Called when the user clicks on the displayed first frame.
+    - orig_frame: The original first frame image (numpy array).
+    - points: The current list of point coordinates.
+    - evt: Event data from the image click (expects evt.index as (x, y)).
+    
+    Returns the updated image (with circles drawn at all points)
+    and the updated list of points.
+    """
+    if points is None:
+        points = []
+    # evt.index contains the (x, y) coordinates of the click.
+    x, y = evt.index
+    new_points = points.copy()
+    new_points.append([x, y])
+    # Draw circles on a copy of the original image.
+    updated_frame = orig_frame.copy()
+    for (px, py) in new_points:
+        cv2.circle(updated_frame, (int(round(px)), int(round(py))), radius=5, color=(0,255,0), thickness=-1)
+        
+    # print(updated_frame.shape)
+    return updated_frame, new_points
+
+# -------------------- Step 2: Process Video & Track Points -------------------- #
+@torch.no_grad()
+@spaces.GPU
+def process_video_with_points(video_path, click_points, tracker):
+    """
+    Runs the dense flow prediction over the entire video, tracking the user-selected points.
+    Args:
+        video_path: Path to the uploaded video.
+        click_points: List of [x, y] coordinates selected on the first frame.
+                      (Coordinates are in the same (resized) space as the displayed first frame.)
+        tracker: The tracker instance to use for processing.
+    Returns:
+        A path to the output video with tracked points overlaid.
+    """
+    if len(click_points) == 0:
+        print("No points selected for tracking.")
+        return "Error: No points selected for tracking."
+
+    # Open the video.
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        return "Error: Could not open video."
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    # List to store frames with overlaid points.
+    output_frames = []
+    # Initialize the points with those selected on the first frame.
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    pbar = tqdm(total=total_frames, desc="Processing video")
+    
+    tracker.reset()
+    
+    frame_disps = []
+    try:
+        while True:
+            torch.cuda.empty_cache()
+            ret, frame = cap.read()
+            if not ret:
+                break
+
+            # Convert frame from BGR to RGB and resize as in your original code.
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            scale = min(tracker.target_res / frame_rgb.shape[0], tracker.target_res / frame_rgb.shape[1])
+            frame_disp = cv2.resize(frame_rgb, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
+            frame_disps.append(frame_disp)
+            
+            flows = tracker.track(frame_rgb)
+            
+            if flows is not None:
+                flows_np = flows[0].cpu().numpy()
+
+                for i, flow_np in enumerate(flows_np):
+                    # --- Update tracked points using the flow ---
+                    current_points = []
+                    for (x, y) in click_points:
+                        xi = int(round(x))
+                        yi = int(round(y))
+                        # print('xi, yi', xi, yi)
+                        if 0 <= yi < flow_np.shape[1] and 0 <= xi < flow_np.shape[2]:
+                            dx = flow_np[0, yi, xi]
+                            dy = flow_np[1, yi, xi]
+                            # print('dx, dy', dx, dy)
+                        else:
+                            dx, dy = 0.0, 0.0
+                        current_points.append([x + dx, y + dy])
+
+                    # Draw the updated points on the frame.
+                    for (x, y) in current_points:
+                        cv2.circle(frame_disps[i], (int(round(x)), int(round(y))), radius=5, color=(0,255,0), thickness=-1)
+                    output_frames.append(frame_disps[i])
+                frame_disps = []
+            pbar.update(1)
+        
+    except RuntimeError as e:
+        # Check if the error message indicates an OOM error.
+        if "out of memory" in str(e).lower():
+            torch.cuda.empty_cache()
+            pbar.close()
+            cap.release()
+            print("Error: Out of Memory during video processing.")
+            return "Error: Out of Memory during video processing. Please try a smaller video or lower resolution."
+        else:
+            # Re-raise if it's another type of error.
+            raise e
+    pbar.close()
+    cap.release()
+
+    # -------------------- Save the Output Video -------------------- #
+    output_path = "tracked_output.mp4"
+    print(len(output_frames), output_frames[0].shape)
+    imageio.mimwrite(output_path, output_frames, fps=fps)
+
+    return output_path
+
+# -------------------- Model Initialization -------------------- #
+def load_model():
+    """Initialize and load the tracking model."""
+    # Adjust these paths as needed.
+    init_dir = '648Ai4i4i3n4s_1e-5m_c5c_stage3_from_kub_ns_wa_kk_lsh_dyk_46470'
+    ckpt_dir = 'checkpoints'
+    load_dir = os.path.join(ckpt_dir, init_dir)
+
+    # Create the model and load weights.
+    model = Net(16)
+    count_parameters(model)
+    _ = utils.saveload.load(
+        None,
+        load_dir,
+        model,
+        optimizer=None,
+        scheduler=None,
+        ignore_load=None,
+        strict=True,
+        verbose=False,
+        weights_only=False,
+    )
+    model.cuda()
+    for n, p in model.named_parameters():
+        p.requires_grad = False
+    model.eval()
+    
+    return model
+
+def create_tracker(model):
+    """Create tracker instance with the loaded model."""
+    return Tracker(
+        model=model,
+        mean=torch.tensor([0.485, 0.456, 0.406]).cuda().reshape(1, 3, 1, 1),
+        std=torch.tensor([0.229, 0.224, 0.225]).cuda().reshape(1, 3, 1, 1),
+        S=16,
+        stride=8,
+        inference_iters=4,
+        target_res=1024,
+    )
+
+# -------------------- Wrappers to Update Tracker Based on UI Settings -------------------- #
+def extract_with_config(video_path, points, resolution, window_index, tracker):
+    """
+    Update the tracker configuration using the slider values, then extract the first frame.
+    - resolution: Target resolution from slider (e.g., 512, 768, 1024).
+    - window_index: An index (0–3) to be mapped to sliding window lengths {0:2, 1:4, 2:8, 3:16}.
+    - tracker: The tracker instance to configure.
+    """
+    tracker.target_res = resolution
+    mapping = {0: 2, 1: 4, 2: 8, 3: 16}
+    tracker.S = mapping.get(int(window_index), 16)
+    return extract_first_frame(video_path, points, tracker)
+
+@torch.no_grad()
+@spaces.GPU
+def process_with_config(video_path, click_points, resolution, window_index, tracker):
+    """
+    Update the tracker configuration using the slider values, then process the video.
+    """
+    tracker.target_res = resolution
+    mapping = {0: 2, 1: 4, 2: 8, 3: 16}
+    tracker.S = mapping.get(int(window_index), 16)
+    return process_video_with_points(video_path, click_points, tracker)
+
+def main():
+    """Main function that initializes the model and runs the Gradio interface."""
+    # Set torch matmul precision
+    torch.set_float32_matmul_precision('medium')
+    
+    # Initialize random seeds
+    seed_everything(42)
+    torch.set_grad_enabled(False)
+    
+    # Load model and create tracker
+    print("Loading model...")
+    model = load_model()
+    tracker = create_tracker(model)
+    print("Model loaded successfully!")
+    
+    # -------------------- Gradio Interface -------------------- #
+    # The interface is built in two steps:
+    # 1. Upload a video and extract the first frame.
+    # 2. Annotate the first frame with multiple points (using gr.Points),
+    #    then run tracking on the video.
+    with gr.Blocks() as demo:
+        gr.Markdown("## Dense Flow Tracking with Clickable Points")
+        
+        with gr.Row():
+            with gr.Column():
+                video_input = gr.Video(label="Upload Video", value="data/244754_medium.mp4")
+                extract_btn = gr.Button("Extract First Frame")
+                # Add sliders for resolution and sliding window length.
+                resolution_slider = gr.Slider(minimum=512, maximum=1024, step=256, value=1024, label="Target Resolution")
+                # The slider below outputs an index 0-3; we'll map it to {0:2, 1:4, 2:8, 3:16}
+                window_slider = gr.Slider(minimum=0, maximum=3, step=1, value=3, label="Sliding Window Length (Index: 0->2, 1->4, 2->8, 3->16)")
+            with gr.Column():
+                # This image will display the first frame and be interactive.
+                first_frame_display = gr.Image(label="First Frame (Click to add points)", interactive=True)
+                clear_pts_btn = gr.Button("Clear Points")
+        
+        # Hidden states: video file path, original first frame, and accumulated click points.
+        video_state = gr.State(None)
+        orig_frame_state = gr.State(None)
+        points_state = gr.State([])
+
+        track_btn = gr.Button("Track Points")
+        output_video = gr.Video(label="Tracked Video")
+        
+        # When "Extract First Frame" is clicked, extract and display the first frame.
+        extract_btn.click(
+            fn=lambda *args: extract_with_config(*args, tracker),
+            inputs=[video_input, points_state, resolution_slider, window_slider],
+            outputs=[first_frame_display, video_state, orig_frame_state, points_state]
+        )
+        
+        clear_pts_btn.click(
+            fn=lambda orig_frame, points: (orig_frame, []),
+            inputs=[orig_frame_state, points_state],
+            outputs=[first_frame_display, points_state]
+        )
+        
+        # When the user clicks on the image, add a point.
+        first_frame_display.select(
+            fn=add_point,
+            inputs=[orig_frame_state, points_state],
+            outputs=[first_frame_display, points_state]
+        )
+        
+        # When "Track Points" is clicked, process the video using the accumulated points.
+        track_btn.click(
+            fn=lambda *args: process_with_config(*args, tracker),
+            inputs=[video_state, points_state, resolution_slider, window_slider],
+            outputs=output_video
+        )
+
+    demo.launch()
+
+if __name__ == '__main__':
+    main()

checkpoints/648Ai4i4i3n4s_1e-5m_c5c_stage3_from_kub_ns_wa_kk_lsh_dyk_46470/model-000600000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15b029e1396feb453d361988afafab02d94993f4b43191b4dfcc44aac10fffa5
+size 198027808

data/244754_medium.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a26ffc0ea63adfb3e2b44f5f6ea384a6071f74cd8453e7b6d386a829d8c11fcb
+size 42948491

demo_dense_visualize.py

@@ -0,0 +1,229 @@
+import os
+import random
+import torch
+import signal
+import socket
+import sys
+import json
+import torch.nn.functional as F
+import numpy as np
+import argparse
+from pathlib import Path
+import torch.optim as optim
+from torch.cuda.amp import GradScaler
+from lightning_fabric import Fabric 
+
+import utils.loss
+import utils.samp
+import utils.data
+import utils.improc
+import utils.misc
+import utils.saveload
+from tensorboardX import SummaryWriter
+import datetime
+import time
+import cv2
+import imageio
+from nets.blocks import InputPadder
+from tqdm import tqdm
+# from pytorch_lightning.callbacks import BaseFinetuning
+from utils.visualizer import Visualizer
+from torchvision.transforms.functional import resize
+
+import torch
+import requests
+from PIL import Image, ImageDraw
+from transformers import AutoProcessor, AutoModelForCausalLM
+import numpy as np
+
+
+torch.set_float32_matmul_precision('medium')
+
+def run_example(processor, model, task_prompt, image, text_input=None):
+    if text_input is None:
+        prompt = task_prompt
+    else:
+        prompt = task_prompt + text_input
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to('cuda', torch.float32)
+    generated_ids = model.generate(
+      input_ids=inputs["input_ids"],
+      pixel_values=inputs["pixel_values"],
+      max_new_tokens=1024,
+      num_beams=3
+    )
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+
+    parsed_answer = processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
+
+    return parsed_answer
+
+
+def polygons_to_mask(image, prediction, fill_value=255):
+    """
+    Converts polygons into a mask.
+    
+    Parameters:
+    - image: A PIL Image instance whose size will be used for the mask.
+    - prediction: Dictionary containing 'polygons' and 'labels'. 
+                  'polygons' is a list where each element is a list of sub-polygons.
+    - fill_value: The pixel value used to fill the polygon areas (default 255 for a binary mask).
+    
+    Returns:
+    - A NumPy array representing the mask (same width and height as the input image).
+    """
+    # Create a blank grayscale mask image with the same size as the original image.
+    mask = Image.new('L', image.size, 0)
+    draw = ImageDraw.Draw(mask)
+    
+    # Iterate over each set of polygons
+    for polygons in prediction['polygons']:
+        # Each element in "polygons" can be a sub-polygon
+        for poly in polygons:
+            # Ensure the polygon is in the right shape and has at least 3 points.
+            poly_arr = np.array(poly).reshape(-1, 2)
+            if poly_arr.shape[0] < 3:
+                print('Skipping invalid polygon:', poly_arr)
+                continue
+            # Convert the polygon vertices into a list for drawing.
+            poly_list = poly_arr.reshape(-1).tolist()
+            # Draw the polygon on the mask with the fill_value.
+            draw.polygon(poly_list, fill=fill_value)
+    
+    # Convert the PIL mask image to a NumPy array and return it.
+    return np.array(mask)
+
+
+class Tracker:
+    def __init__(self, model, mean, std, S, stride, inference_iters, target_res, device='cuda'):
+        """
+        Initializes the Tracker.
+
+        Args:
+            model: The model used to compute feature maps and forward window flow.
+            mean: Tensor or value used for normalizing the input.
+            std: Tensor or value used for normalizing the input.
+            S: Window size for the tracker.
+            stride: The stride used when updating the window.
+            inference_iters: Number of inference iterations.
+            device: Torch device, defaults to 'cuda'.
+        """
+        self.model = model
+        self.mean = mean
+        self.std = std
+        self.S = S
+        self.stride = stride
+        self.inference_iters = inference_iters
+        self.device = device
+        self.target_res = target_res
+
+        self.padder = None
+        self.cnt = 0
+        self.fmap_anchor = None
+        self.fmaps2 = None
+        self.flows8 = None
+        self.visconfs8 = None
+        self.flows = []  # List to store computed flows
+        self.visibs = []  # List to store visibility confidences
+        self.rgbs = []  # List to store RGB frames
+
+    def reset(self):
+        """Reset the tracker state."""
+        self.padder = None
+        self.cnt = 0
+        self.fmap_anchor = None
+        self.fmaps2 = None
+        self.flows8 = None
+        self.visconfs8 = None
+        self.flows = []
+        self.visibs = []
+        self.rgbs = []
+        
+    def preprocess(self, rgb_frame):
+        # Resize frame (scale to keep maximum dimension ~1024)
+        scale = min(self.target_res / rgb_frame.shape[0], self.target_res / rgb_frame.shape[1])
+        rgb_resized = cv2.resize(rgb_frame, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
+        
+        # Convert to tensor, normalize and move to device.
+        rgb_tensor = torch.from_numpy(rgb_resized).permute(2, 0, 1).float().unsqueeze(0).to(self.device)
+        rgb_tensor = rgb_tensor / 255.0
+        
+        self.rgbs.append(rgb_tensor.cpu())
+            
+        # import pdb; pdb.set_trace()
+        rgb_tensor = (rgb_tensor - self.mean) / self.std
+        return rgb_tensor
+
+    @torch.no_grad()
+    def track(self, rgb_frame):
+        """
+        Process a single RGB frame and return the computed flow when available.
+
+        Args:
+            rgb_frame: A NumPy array containing the RGB frame.
+                       (Assumed to be in RGB; if coming from OpenCV, convert it before passing.)
+        
+        Returns:
+            flow_predictions: The predicted flow for the current frame (or None if not enough frames have been processed).
+        """
+        torch.cuda.empty_cache()
+
+        rgb_tensor = self.preprocess(rgb_frame)
+        
+        # Initialize padder on the first frame.
+        if self.cnt == 0:
+            self.padder = InputPadder(rgb_tensor.shape)
+        rgb_padded = self.padder.pad(rgb_tensor)[0]
+        _, _, H_pad, W_pad = rgb_padded.shape
+        C = 256  # Feature map channel dimension (could be parameterized if needed)
+        H8, W8 = H_pad // 8, W_pad // 8
+
+        # Accumulate feature maps until the window is full.
+        if self.cnt == 0:
+            self.fmap_anchor = self.model.get_fmaps(rgb_padded, 1, 1, None, False, False).reshape(1, C, H8, W8)
+            self.fmaps2 = self.fmap_anchor[:, None]
+            self.cnt += 1
+            return None
+            
+        new_fmap = self.model.get_fmaps(rgb_padded, 1, 1, None, False, False).reshape(1, 1, C, H8, W8)
+        self.fmaps2 = torch.cat([self.fmaps2[:, (1 if self.fmaps2.shape[1] >= self.S else 0):].detach().clone(), new_fmap], dim=1)
+        
+        # need to track
+        if self.cnt - self.S + 1 >= 0 and (self.cnt - self.S + 1) % self.stride == 0:
+            # Initialize or update temporary flow buffers.
+            iter_num = self.inference_iters
+            if self.flows8 is None:
+                self.flows8 = torch.zeros((self.S, 2, H_pad // 8, W_pad // 8), device=self.device)
+                self.visconfs8 = torch.zeros((self.S, 2, H_pad // 8, W_pad // 8), device=self.device)
+                # iter_num = self.inference_iters
+            else:
+                self.flows8 = torch.cat([
+                    self.flows8[self.stride:self.stride + self.S // 2].detach().clone(),
+                    self.flows8[self.stride + self.S // 2 - 1:self.stride + self.S // 2].detach().clone().repeat(self.S // 2, 1, 1, 1)
+                ])
+                self.visconfs8 = torch.cat([
+                    self.visconfs8[self.stride:self.stride + self.S // 2].detach().clone(),
+                    self.visconfs8[self.stride + self.S // 2 - 1:self.stride + self.S // 2].detach().clone().repeat(self.S // 2, 1, 1, 1)
+                ])
+
+            # import pdb; pdb.set_trace()
+            # Compute flow predictions using the model's forward window.
+            flow_predictions, visconf_predictions, self.flows8, self.visconfs8, _ = self.model.forward_window(
+                self.fmap_anchor,
+                self.fmaps2,
+                self.visconfs8,
+                iters=iter_num,
+                flowfeat=None,
+                flows8=self.flows8,
+                is_training=False
+            )
+            flow_predictions = self.padder.unpad(flow_predictions[-1][0 if self.cnt == self.S - 1 else -self.stride:])
+            visconf_predictions = self.padder.unpad(torch.sigmoid(visconf_predictions[-1][0 if self.cnt == self.S - 1 else -self.stride:]))
+
+            self.cnt += 1
+            self.flows.append(flow_predictions.cpu())
+            self.visibs.append(visconf_predictions.cpu())
+
+            return flow_predictions, visconf_predictions
+        
+        self.cnt += 1
+        return None

nets/blocks.py

@@ -0,0 +1,1337 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import nn, Tensor
+from itertools import repeat
+import collections
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence
+from functools import partial
+import einops
+import math
+from torchvision.ops.misc import Conv2dNormActivation, Permute
+from torchvision.ops.stochastic_depth import StochasticDepth
+
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+    return parse
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+to_2tuple = _ntuple(2)
+
+class InputPadder:
+    """ Pads images such that dimensions are divisible by a certain stride """
+    def __init__(self, dims, mode='sintel'):
+        self.ht, self.wd = dims[-2:]
+        pad_ht = (((self.ht // 64) + 1) * 64 - self.ht) % 64
+        pad_wd = (((self.wd // 64) + 1) * 64 - self.wd) % 64
+        if mode == 'sintel':
+            self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
+        else:
+            self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
+
+    def pad(self, *inputs):
+        return [F.pad(x, self._pad, mode='replicate') for x in inputs]
+
+    def unpad(self, x):
+        ht, wd = x.shape[-2:]
+        c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
+        return x[..., c[0]:c[1], c[2]:c[3]]
+
+def bilinear_sampler(
+        input, coords,
+        align_corners=True,
+        padding_mode="border",
+        normalize_coords=True):
+    # func from mattie (oct9)
+    if input.ndim not in [4, 5]:
+        raise ValueError("input must be 4D or 5D.")
+    
+    if input.ndim == 4 and not coords.ndim == 4:
+        raise ValueError("input is 4D, but coords is not 4D.")
+
+    if input.ndim == 5 and not coords.ndim == 5:
+        raise ValueError("input is 5D, but coords is not 5D.")
+
+    if coords.ndim == 5:
+        coords = coords[..., [1, 2, 0]]  # t x y -> x y t to match what grid_sample() expects.
+
+    if normalize_coords:
+        if align_corners:
+            # Normalize coordinates from [0, W/H - 1] to [-1, 1].
+            coords = (
+                coords
+                * torch.tensor([2 / max(size - 1, 1) for size in reversed(input.shape[2:])], device=coords.device)
+                - 1
+            )
+        else:
+            # Normalize coordinates from [0, W/H] to [-1, 1].
+            coords = coords * torch.tensor([2 / size for size in reversed(input.shape[2:])], device=coords.device) - 1
+            
+    return F.grid_sample(input, coords, align_corners=align_corners, padding_mode=padding_mode)
+
+
+class CorrBlock:
+    def __init__(self, fmap1, fmap2, corr_levels, corr_radius):
+        self.num_levels = corr_levels
+        self.radius = corr_radius
+        self.corr_pyramid = []
+        # all pairs correlation
+        for i in range(self.num_levels):
+            corr = CorrBlock.corr(fmap1, fmap2, 1)
+            batch, h1, w1, dim, h2, w2 = corr.shape
+            corr = corr.reshape(batch*h1*w1, dim, h2, w2)
+            fmap2 = F.interpolate(fmap2, scale_factor=0.5, mode='area')
+            # print('corr', corr.shape)
+            self.corr_pyramid.append(corr)
+
+    def __call__(self, coords, dilation=None):
+        r = self.radius
+        coords = coords.permute(0, 2, 3, 1)
+        batch, h1, w1, _ = coords.shape
+
+        if dilation is None:
+            dilation = torch.ones(batch, 1, h1, w1, device=coords.device)
+
+        # print(dilation.max(), dilation.mean(), dilation.min())
+        out_pyramid = []
+        for i in range(self.num_levels):
+            corr = self.corr_pyramid[i]
+            device = coords.device
+            dx = torch.linspace(-r, r, 2*r+1, device=device)
+            dy = torch.linspace(-r, r, 2*r+1, device=device)
+            delta = torch.stack(torch.meshgrid(dy, dx), axis=-1)
+            delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
+            delta_lvl = delta_lvl * dilation.view(batch * h1 * w1, 1, 1, 1)
+            centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i
+            coords_lvl = centroid_lvl + delta_lvl
+            corr = bilinear_sampler(corr, coords_lvl)
+            corr = corr.view(batch, h1, w1, -1)
+            out_pyramid.append(corr)
+
+        out = torch.cat(out_pyramid, dim=-1)
+        out = out.permute(0, 3, 1, 2).contiguous().float()  
+        return out
+
+    @staticmethod
+    def corr(fmap1, fmap2, num_head):
+        batch, dim, h1, w1 = fmap1.shape
+        h2, w2 = fmap2.shape[2:]
+        fmap1 = fmap1.view(batch, num_head, dim // num_head, h1*w1)
+        fmap2 = fmap2.view(batch, num_head, dim // num_head, h2*w2) 
+        corr = fmap1.transpose(2, 3) @ fmap2
+        corr = corr.reshape(batch, num_head, h1, w1, h2, w2).permute(0, 2, 3, 1, 4, 5)
+        return corr  / torch.sqrt(torch.tensor(dim).float())
+    
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution without padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0)
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1)
+
+class LayerNorm2d(nn.LayerNorm):
+    def forward(self, x: Tensor) -> Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.permute(0, 3, 1, 2)
+        return x
+    
+class CNBlock1d(nn.Module):
+    def __init__(
+        self,
+        dim,
+        output_dim,
+        layer_scale: float = 1e-6,
+        stochastic_depth_prob: float = 0,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        dense=True,
+        use_attn=True,
+        use_mixer=False,
+        use_conv=False,
+        use_convb=False,
+        use_layer_scale=True,
+    ) -> None:
+        super().__init__()
+        self.dense = dense
+        self.use_attn = use_attn
+        self.use_mixer = use_mixer
+        self.use_conv = use_conv
+        self.use_layer_scale = use_layer_scale
+
+        if use_attn:
+            assert not use_mixer
+            assert not use_conv
+            assert not use_convb
+        
+        if norm_layer is None:
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        if use_attn:
+            num_heads = 8
+            self.block = AttnBlock(
+                hidden_size=dim,
+                num_heads=num_heads,
+                mlp_ratio=4,
+                attn_class=Attention,
+            )
+        elif use_mixer:
+            self.block = MLPMixerBlock(
+                S=16,
+                dim=dim,
+                depth=1,
+                expansion_factor=2,
+            )
+        elif use_conv:
+            self.block = nn.Sequential(
+                nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim, bias=True, padding_mode='zeros'),
+                Permute([0, 2, 1]),
+                norm_layer(dim),
+                nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
+                nn.GELU(),
+                nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
+                Permute([0, 2, 1]),
+            )
+        elif use_convb:
+            self.block = nn.Sequential(
+                nn.Conv1d(dim, dim, kernel_size=3, padding=1, bias=True, padding_mode='zeros'),
+                Permute([0, 2, 1]),
+                norm_layer(dim),
+                nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
+                nn.GELU(),
+                nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
+                Permute([0, 2, 1]),
+            )
+        else:
+            assert(False) # choose attn, mixer, or conv please
+
+        if self.use_layer_scale:
+            self.layer_scale = nn.Parameter(torch.ones(dim, 1) * layer_scale)
+        else:
+            self.layer_scale = 1.0
+            
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
+
+        if output_dim != dim:
+            self.final = nn.Conv1d(dim, output_dim, kernel_size=1, padding=0)
+        else:
+            self.final = nn.Identity()
+
+    def forward(self, input, S=None):
+        if self.dense:
+            assert S is not None
+            BS,C,H,W = input.shape
+            B = BS//S
+
+            # if S<7:
+            #     return input
+
+            input = einops.rearrange(input, '(b s) c h w -> (b h w) c s', b=B, s=S, c=C, h=H, w=W)
+
+            if self.use_mixer or self.use_attn:
+                # mixer/transformer blocks want B,S,C
+                result = self.layer_scale * self.block(input.permute(0,2,1)).permute(0,2,1)
+            else:
+                result = self.layer_scale * self.block(input)
+            result = self.stochastic_depth(result)
+            result += input
+            result = self.final(result)
+
+            result = einops.rearrange(result, '(b h w) c s -> (b s) c h w', b=B, s=S, c=C, h=H, w=W)
+        else:
+            B,S,C = input.shape
+
+            if S<7:
+                return input
+
+            input = einops.rearrange(input, 'b s c -> b c s', b=B, s=S, c=C)
+
+            result = self.layer_scale * self.block(input)
+            result = self.stochastic_depth(result)
+            result += input
+
+            result = self.final(result)
+
+            result = einops.rearrange(result, 'b c s -> b s c', b=B, s=S, c=C)
+        
+        return result
+    
+class CNBlock2d(nn.Module):
+    def __init__(
+        self,
+        dim,
+        output_dim,
+        layer_scale: float = 1e-6,
+        stochastic_depth_prob: float = 0,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        use_layer_scale=True,
+    ) -> None:
+        super().__init__()
+        self.use_layer_scale = use_layer_scale
+        if norm_layer is None:
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.block = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim, bias=True, padding_mode='zeros'),
+            Permute([0, 2, 3, 1]),
+            norm_layer(dim),
+            nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
+            nn.GELU(),
+            nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
+            Permute([0, 3, 1, 2]),
+        )
+        if self.use_layer_scale:
+            self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
+        else:
+            self.layer_scale = 1.0
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
+        
+        if output_dim != dim:
+            self.final = nn.Conv2d(dim, output_dim, kernel_size=1, padding=0)
+        else:
+            self.final = nn.Identity()
+
+    def forward(self, input, S=None):
+        result = self.layer_scale * self.block(input)
+        result = self.stochastic_depth(result)
+        result += input
+        result = self.final(result)
+        return result
+
+class CNBlockConfig:
+    # Stores information listed at Section 3 of the ConvNeXt paper
+    def __init__(
+        self,
+        input_channels: int,
+        out_channels: Optional[int],
+        num_layers: int,
+        downsample: bool,
+    ) -> None:
+        self.input_channels = input_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+        self.downsample = downsample
+
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "input_channels={input_channels}"
+        s += ", out_channels={out_channels}"
+        s += ", num_layers={num_layers}"
+        s += ", downsample={downsample}"
+        s += ")"
+        return s.format(**self.__dict__)
+    
+class ConvNeXt(nn.Module):
+    def __init__(
+            self,
+            block_setting: List[CNBlockConfig],
+            stochastic_depth_prob: float = 0.0,
+            layer_scale: float = 1e-6,
+            num_classes: int = 1000,
+            block: Optional[Callable[..., nn.Module]] = None,
+            norm_layer: Optional[Callable[..., nn.Module]] = None,
+            init_weights=True):
+        super().__init__()
+
+        self.init_weights = init_weights
+        
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock2d
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        layers: List[nn.Module] = []
+
+        # Stem
+        firstconv_output_channels = block_setting[0].input_channels
+        layers.append(
+            Conv2dNormActivation(
+                3,
+                firstconv_output_channels,
+                kernel_size=4,
+                stride=4,
+                padding=0,
+                norm_layer=norm_layer,
+                activation_layer=None,
+                bias=True,
+            )
+        )
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        for cnf in block_setting:
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                # adjust stochastic depth probability based on the depth of the stage block
+                sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                stage.append(block(cnf.input_channels, cnf.input_channels, layer_scale, sd_prob))
+                stage_block_id += 1
+            layers.append(nn.Sequential(*stage))
+            if cnf.out_channels is not None:
+                if cnf.downsample:
+                    layers.append(
+                        nn.Sequential(
+                            norm_layer(cnf.input_channels),
+                            nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                        )
+                    )
+                else:
+                    # we convert the 2x2 downsampling layer into a 3x3 with dilation2 and replicate padding.
+                    # replicate padding compensates for the fact that this kernel never saw zero-padding.
+                    layers.append(
+                        nn.Sequential(
+                            norm_layer(cnf.input_channels),
+                            nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=3, stride=1, padding=2, dilation=2, padding_mode='zeros'),
+                        )
+                    )
+
+        self.features = nn.Sequential(*layers)
+        
+        # self.final_conv = conv1x1(block_setting[-1].input_channels, output_dim)
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        if self.init_weights:
+            from torchvision.models import convnext_tiny, ConvNeXt_Tiny_Weights
+            pretrained_dict = convnext_tiny(weights=ConvNeXt_Tiny_Weights.DEFAULT).state_dict()
+            # from torchvision.models import convnext_base, ConvNeXt_Base_Weights
+            # pretrained_dict = convnext_base(weights=ConvNeXt_Base_Weights.DEFAULT).state_dict()
+            model_dict = self.state_dict()
+            pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+            # if self.input_dim == 6:
+            #     for k, v in pretrained_dict.items():
+            #         if k == 'conv1.weight':
+            #             pretrained_dict[k] = torch.cat((v, v), dim=1)
+
+            # del pretrained_dict['features.4.1.weight']
+
+            for k, v in pretrained_dict.items():
+                if k == 'features.4.1.weight': # this is the layer normally in charge of 2x2 downsampling
+                    # convert to 3x3 filter
+                    pretrained_dict[k] = F.interpolate(v, (3, 3), mode='bicubic', align_corners=True) * (4/9.0)
+                    
+                    # print('v0', v[0,0])
+                    # print('d0', pretrained_dict[k][0,0])
+            
+            model_dict.update(pretrained_dict)
+            self.load_state_dict(model_dict, strict=False)
+        
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        # with torch.no_grad():
+        x = self.features(x)
+        # x = self.final_conv(x)
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+class Mlp(nn.Module):
+    """MLP as used in Vision Transformer, MLP-Mixer and related networks"""
+
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        norm_layer=None,
+        bias=True,
+        drop=0.0,
+        use_conv=False,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        bias = to_2tuple(bias)
+        drop_probs = to_2tuple(drop)
+        linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
+
+        self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0])
+        self.act = act_layer()
+        self.drop1 = nn.Dropout(drop_probs[0])
+        self.norm = (
+            norm_layer(hidden_features) if norm_layer is not None else nn.Identity()
+        )
+        self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1])
+        self.drop2 = nn.Dropout(drop_probs[1])
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop1(x)
+        x = self.fc2(x)
+        x = self.drop2(x)
+        return x
+
+class Attention(nn.Module):
+    def __init__(
+        self, query_dim, context_dim=None, num_heads=8, dim_head=48, qkv_bias=False
+    ):
+        super().__init__()
+        inner_dim = dim_head * num_heads
+        context_dim = default(context_dim, query_dim)
+        self.scale = dim_head**-0.5
+        self.heads = num_heads
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=qkv_bias)
+        self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias=qkv_bias)
+        self.to_out = nn.Linear(inner_dim, query_dim)
+
+    def forward(self, x, context=None, attn_bias=None):
+        B, N1, C = x.shape
+        H = self.heads
+        q = self.to_q(x)
+        context = default(context, x)
+        k, v = self.to_kv(context).chunk(2, dim=-1)
+        q, k, v = map(lambda t: einops.rearrange(t, 'b n (h d) -> b h n d', h=self.heads), (q, k, v))
+        x = F.scaled_dot_product_attention(q, k, v) # scale default is already dim^-0.5
+        x = einops.rearrange(x, 'b h n d -> b n (h d)')
+        return self.to_out(x)
+    
+class CrossAttnBlock(nn.Module):
+    def __init__(
+        self, hidden_size, context_dim, num_heads=1, mlp_ratio=4.0, **block_kwargs
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_context = nn.LayerNorm(hidden_size)
+        self.cross_attn = Attention(
+            hidden_size,
+            context_dim=context_dim,
+            num_heads=num_heads,
+            qkv_bias=True,
+            **block_kwargs
+        )
+
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+            drop=0,
+        )
+
+    def forward(self, x, context, mask=None):
+        attn_bias = None
+        if mask is not None:
+            if mask.shape[1] == x.shape[1]:
+                mask = mask[:, None, :, None].expand(
+                    -1, self.cross_attn.heads, -1, context.shape[1]
+                )
+            else:
+                mask = mask[:, None, None].expand(
+                    -1, self.cross_attn.heads, x.shape[1], -1
+                )
+
+            max_neg_value = -torch.finfo(x.dtype).max
+            attn_bias = (~mask) * max_neg_value
+        x = x + self.cross_attn(
+            self.norm1(x), context=self.norm_context(context), attn_bias=attn_bias
+        )
+        x = x + self.mlp(self.norm2(x))
+        return x
+    
+class AttnBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        attn_class: Callable[..., nn.Module] = Attention,
+        mlp_ratio=4.0,
+        **block_kwargs
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = attn_class(hidden_size, num_heads=num_heads, qkv_bias=True, dim_head=hidden_size//num_heads)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+            drop=0,
+        )
+
+    def forward(self, x, mask=None):
+        attn_bias = mask
+        if mask is not None:
+            mask = (
+                (mask[:, None] * mask[:, :, None])
+                .unsqueeze(1)
+                .expand(-1, self.attn.num_heads, -1, -1)
+            )
+            max_neg_value = -torch.finfo(x.dtype).max
+            attn_bias = (~mask) * max_neg_value
+
+        x = x + self.attn(self.norm1(x), attn_bias=attn_bias)
+        x = x + self.mlp(self.norm2(x))
+        return x
+    
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn="group", stride=1):
+        super(ResidualBlock, self).__init__()
+
+        self.conv1 = nn.Conv2d(
+            in_planes,
+            planes,
+            kernel_size=3,
+            padding=1,
+            stride=stride,
+            padding_mode="zeros",
+        )
+        self.conv2 = nn.Conv2d(
+            planes, planes, kernel_size=3, padding=1, padding_mode="zeros"
+        )
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == "group":
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+
+        elif norm_fn == "batch":
+            self.norm1 = nn.BatchNorm2d(planes)
+            self.norm2 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.BatchNorm2d(planes)
+
+        elif norm_fn == "instance":
+            self.norm1 = nn.InstanceNorm2d(planes)
+            self.norm2 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == "none":
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            if not stride == 1:
+                self.norm3 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+
+        else:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3
+            )
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x + y)
+
+
+class BasicEncoder(nn.Module):
+    def __init__(self, input_dim=3, output_dim=128, stride=4):
+        super(BasicEncoder, self).__init__()
+        self.stride = stride
+        self.norm_fn = "instance"
+        self.in_planes = output_dim // 2
+        self.norm1 = nn.InstanceNorm2d(self.in_planes)
+        self.norm2 = nn.InstanceNorm2d(output_dim * 2)
+
+        self.conv1 = nn.Conv2d(
+            input_dim,
+            self.in_planes,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            padding_mode="zeros",
+        )
+        self.relu1 = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(output_dim // 2, stride=1)
+        self.layer2 = self._make_layer(output_dim // 4 * 3, stride=2)
+        self.layer3 = self._make_layer(output_dim, stride=2)
+        self.layer4 = self._make_layer(output_dim, stride=2)
+
+        self.conv2 = nn.Conv2d(
+            output_dim * 3 + output_dim // 4,
+            output_dim * 2,
+            kernel_size=3,
+            padding=1,
+            padding_mode="zeros",
+        )
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(output_dim * 2, output_dim, kernel_size=1)
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.InstanceNorm2d)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        _, _, H, W = x.shape
+
+        # with torch.no_grad():
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        a = self.layer1(x)
+        b = self.layer2(a)
+        c = self.layer3(b)
+        d = self.layer4(c)
+
+        def _bilinear_intepolate(x):
+            return F.interpolate(
+                x,
+                (H // self.stride, W // self.stride),
+                mode="bilinear",
+                align_corners=True,
+            )
+
+        a = _bilinear_intepolate(a)
+        b = _bilinear_intepolate(b)
+        c = _bilinear_intepolate(c)
+        d = _bilinear_intepolate(d)
+
+        x = self.conv2(torch.cat([a, b, c, d], dim=1))
+        x = self.norm2(x)
+        x = self.relu2(x)
+        x = self.conv3(x)
+        return x
+
+class EfficientUpdateFormer(nn.Module):
+    """
+    Transformer model that updates track estimates.
+    """
+
+    def __init__(
+        self,
+        space_depth=6,
+        time_depth=6,
+        input_dim=320,
+        hidden_size=384,
+        num_heads=8,
+        output_dim=130,
+        mlp_ratio=4.0,
+        num_virtual_tracks=64,
+        add_space_attn=True,
+        linear_layer_for_vis_conf=False,
+        use_time_conv=False,
+        use_time_mixer=False,
+    ):
+        super().__init__()
+        self.out_channels = 2
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.input_transform = torch.nn.Linear(input_dim, hidden_size, bias=True)
+        if linear_layer_for_vis_conf:
+            self.flow_head = torch.nn.Linear(hidden_size, output_dim - 2, bias=True)
+            self.vis_conf_head = torch.nn.Linear(hidden_size, 2, bias=True)
+        else:
+            self.flow_head = torch.nn.Linear(hidden_size, output_dim, bias=True)
+        self.num_virtual_tracks = num_virtual_tracks
+        self.virual_tracks = nn.Parameter(
+            torch.randn(1, num_virtual_tracks, 1, hidden_size)
+        )
+        self.add_space_attn = add_space_attn
+        self.linear_layer_for_vis_conf = linear_layer_for_vis_conf
+
+        if use_time_conv:
+            self.time_blocks = nn.ModuleList(
+                [
+                    CNBlock1d(hidden_size, hidden_size, dense=False)
+                    for _ in range(time_depth)
+                ]
+            )
+        elif use_time_mixer:
+            self.time_blocks = nn.ModuleList(
+                [
+                    MLPMixerBlock(
+                        S=16,
+                        dim=hidden_size,
+                        depth=1,
+                    )
+                    for _ in range(time_depth)
+                ]
+            )
+        else:
+            self.time_blocks = nn.ModuleList(
+                [
+                    AttnBlock(
+                        hidden_size,
+                        num_heads,
+                        mlp_ratio=mlp_ratio,
+                        attn_class=Attention,
+                    )
+                    for _ in range(time_depth)
+                ]
+            )
+
+        if add_space_attn:
+            self.space_virtual_blocks = nn.ModuleList(
+                [
+                    AttnBlock(
+                        hidden_size,
+                        num_heads,
+                        mlp_ratio=mlp_ratio,
+                        attn_class=Attention,
+                    )
+                    for _ in range(space_depth)
+                ]
+            )
+            self.space_point2virtual_blocks = nn.ModuleList(
+                [
+                    CrossAttnBlock(
+                        hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio
+                    )
+                    for _ in range(space_depth)
+                ]
+            )
+            self.space_virtual2point_blocks = nn.ModuleList(
+                [
+                    CrossAttnBlock(
+                        hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio
+                    )
+                    for _ in range(space_depth)
+                ]
+            )
+            assert len(self.time_blocks) >= len(self.space_virtual2point_blocks)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+            torch.nn.init.trunc_normal_(self.flow_head.weight, std=0.001)
+            if self.linear_layer_for_vis_conf:
+                torch.nn.init.trunc_normal_(self.vis_conf_head.weight, std=0.001)
+
+        def _trunc_init(module):
+            """ViT weight initialization, original timm impl (for reproducibility)"""
+            if isinstance(module, nn.Linear):
+                torch.nn.init.trunc_normal_(module.weight, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+        self.apply(_basic_init)
+
+    def forward(self, input_tensor, mask=None, add_space_attn=True):
+        tokens = self.input_transform(input_tensor)
+
+        B, _, T, _ = tokens.shape
+        virtual_tokens = self.virual_tracks.repeat(B, 1, T, 1)
+        tokens = torch.cat([tokens, virtual_tokens], dim=1)
+
+        _, N, _, _ = tokens.shape
+        j = 0
+        layers = []
+        for i in range(len(self.time_blocks)):
+            time_tokens = tokens.contiguous().view(B * N, T, -1)  # B N T C -> (B N) T C
+            time_tokens = self.time_blocks[i](time_tokens)
+
+            tokens = time_tokens.view(B, N, T, -1)  # (B N) T C -> B N T C
+            if (
+                add_space_attn
+                and hasattr(self, "space_virtual_blocks")
+                and (i % (len(self.time_blocks) // len(self.space_virtual_blocks)) == 0)
+            ):
+                space_tokens = (
+                    tokens.permute(0, 2, 1, 3).contiguous().view(B * T, N, -1)
+                )  # B N T C -> (B T) N C
+
+                point_tokens = space_tokens[:, : N - self.num_virtual_tracks]
+                virtual_tokens = space_tokens[:, N - self.num_virtual_tracks :]
+
+                virtual_tokens = self.space_virtual2point_blocks[j](
+                    virtual_tokens, point_tokens, mask=mask
+                )
+
+                virtual_tokens = self.space_virtual_blocks[j](virtual_tokens)
+                point_tokens = self.space_point2virtual_blocks[j](
+                    point_tokens, virtual_tokens, mask=mask
+                )
+
+                space_tokens = torch.cat([point_tokens, virtual_tokens], dim=1)
+                tokens = space_tokens.view(B, T, N, -1).permute(
+                    0, 2, 1, 3
+                )  # (B T) N C -> B N T C
+                j += 1
+        tokens = tokens[:, : N - self.num_virtual_tracks]
+
+        flow = self.flow_head(tokens)
+        if self.linear_layer_for_vis_conf:
+            vis_conf = self.vis_conf_head(tokens)
+            flow = torch.cat([flow, vis_conf], dim=-1)
+
+        return flow
+
+
+class MMPreNormResidual(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        return self.fn(self.norm(x)) + x
+    
+def MMFeedForward(dim, expansion_factor=4, dropout=0., dense=nn.Linear):
+    return nn.Sequential(
+        dense(dim, dim * expansion_factor),
+        nn.GELU(),
+        nn.Dropout(dropout),
+        dense(dim * expansion_factor, dim),
+        nn.Dropout(dropout)
+    )
+    
+def MLPMixer(S, input_dim, dim, output_dim, depth=6, expansion_factor=4, dropout=0., do_reduce=False):
+    # input is coming in as B,S,C, as standard for mlp and transformer
+    # chan_first treats S as the channel dim, and transforms it to a new S
+    # chan_last treats C as the channel dim, and transforms it to a new C 
+    chan_first, chan_last = partial(nn.Conv1d, kernel_size=1), nn.Linear
+    if do_reduce:
+        return nn.Sequential(
+            nn.Linear(input_dim, dim),
+            *[nn.Sequential(
+                MMPreNormResidual(dim, MMFeedForward(S, expansion_factor, dropout, chan_first)),
+                MMPreNormResidual(dim, MMFeedForward(dim, expansion_factor, dropout, chan_last))
+            ) for _ in range(depth)],
+            nn.LayerNorm(dim),
+            Reduce('b n c -> b c', 'mean'),
+            nn.Linear(dim, output_dim)
+        )
+    else:
+        return nn.Sequential(
+            nn.Linear(input_dim, dim),
+            *[nn.Sequential(
+                MMPreNormResidual(dim, MMFeedForward(S, expansion_factor, dropout, chan_first)),
+                MMPreNormResidual(dim, MMFeedForward(dim, expansion_factor, dropout, chan_last))
+            ) for _ in range(depth)],
+        )
+
+def MLPMixerBlock(S, dim, depth=1, expansion_factor=4, dropout=0., do_reduce=False):
+    # input is coming in as B,S,C, as standard for mlp and transformer
+    # chan_first treats S as the channel dim, and transforms it to a new S
+    # chan_last treats C as the channel dim, and transforms it to a new C 
+    chan_first, chan_last = partial(nn.Conv1d, kernel_size=1), nn.Linear
+    return nn.Sequential(
+        *[nn.Sequential(
+            MMPreNormResidual(dim, MMFeedForward(S, expansion_factor, dropout, chan_first)),
+            MMPreNormResidual(dim, MMFeedForward(dim, expansion_factor, dropout, chan_last))
+        ) for _ in range(depth)],
+    )
+    
+    
+class MlpUpdateFormer(nn.Module):
+    """
+    Transformer model that updates track estimates.
+    """
+
+    def __init__(
+        self,
+        space_depth=6,
+        time_depth=6,
+        input_dim=320,
+        hidden_size=384,
+        num_heads=8,
+        output_dim=130,
+        mlp_ratio=4.0,
+        num_virtual_tracks=64,
+        add_space_attn=True,
+        linear_layer_for_vis_conf=False,
+    ):
+        super().__init__()
+        self.out_channels = 2
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.input_transform = torch.nn.Linear(input_dim, hidden_size, bias=True)
+        if linear_layer_for_vis_conf:
+            self.flow_head = torch.nn.Linear(hidden_size, output_dim - 2, bias=True)
+            self.vis_conf_head = torch.nn.Linear(hidden_size, 2, bias=True)
+        else:
+            self.flow_head = torch.nn.Linear(hidden_size, output_dim, bias=True)
+        self.num_virtual_tracks = num_virtual_tracks
+        self.virual_tracks = nn.Parameter(
+            torch.randn(1, num_virtual_tracks, 1, hidden_size)
+        )
+        self.add_space_attn = add_space_attn
+        self.linear_layer_for_vis_conf = linear_layer_for_vis_conf
+        self.time_blocks = nn.ModuleList(
+            [
+                MLPMixer(
+                    S=16,
+                    input_dim=hidden_size,
+                    dim=hidden_size,
+                    output_dim=hidden_size,
+                    depth=1,
+                )
+                for _ in range(time_depth)
+            ]
+        )
+
+        if add_space_attn:
+            self.space_virtual_blocks = nn.ModuleList(
+                [
+                    AttnBlock(
+                        hidden_size,
+                        num_heads,
+                        mlp_ratio=mlp_ratio,
+                        attn_class=Attention,
+                    )
+                    for _ in range(space_depth)
+                ]
+            )
+            self.space_point2virtual_blocks = nn.ModuleList(
+                [
+                    CrossAttnBlock(
+                        hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio
+                    )
+                    for _ in range(space_depth)
+                ]
+            )
+            self.space_virtual2point_blocks = nn.ModuleList(
+                [
+                    CrossAttnBlock(
+                        hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio
+                    )
+                    for _ in range(space_depth)
+                ]
+            )
+            assert len(self.time_blocks) >= len(self.space_virtual2point_blocks)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+            torch.nn.init.trunc_normal_(self.flow_head.weight, std=0.001)
+            if self.linear_layer_for_vis_conf:
+                torch.nn.init.trunc_normal_(self.vis_conf_head.weight, std=0.001)
+
+        def _trunc_init(module):
+            """ViT weight initialization, original timm impl (for reproducibility)"""
+            if isinstance(module, nn.Linear):
+                torch.nn.init.trunc_normal_(module.weight, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+        self.apply(_basic_init)
+
+    def forward(self, input_tensor, mask=None, add_space_attn=True):
+        tokens = self.input_transform(input_tensor)
+
+        B, _, T, _ = tokens.shape
+        virtual_tokens = self.virual_tracks.repeat(B, 1, T, 1)
+        tokens = torch.cat([tokens, virtual_tokens], dim=1)
+
+        _, N, _, _ = tokens.shape
+        j = 0
+        layers = []
+        for i in range(len(self.time_blocks)):
+            time_tokens = tokens.contiguous().view(B * N, T, -1)  # B N T C -> (B N) T C
+            time_tokens = self.time_blocks[i](time_tokens)
+
+            tokens = time_tokens.view(B, N, T, -1)  # (B N) T C -> B N T C
+            if (
+                add_space_attn
+                and hasattr(self, "space_virtual_blocks")
+                and (i % (len(self.time_blocks) // len(self.space_virtual_blocks)) == 0)
+            ):
+                space_tokens = (
+                    tokens.permute(0, 2, 1, 3).contiguous().view(B * T, N, -1)
+                )  # B N T C -> (B T) N C
+
+                point_tokens = space_tokens[:, : N - self.num_virtual_tracks]
+                virtual_tokens = space_tokens[:, N - self.num_virtual_tracks :]
+
+                virtual_tokens = self.space_virtual2point_blocks[j](
+                    virtual_tokens, point_tokens, mask=mask
+                )
+
+                virtual_tokens = self.space_virtual_blocks[j](virtual_tokens)
+                point_tokens = self.space_point2virtual_blocks[j](
+                    point_tokens, virtual_tokens, mask=mask
+                )
+
+                space_tokens = torch.cat([point_tokens, virtual_tokens], dim=1)
+                tokens = space_tokens.view(B, T, N, -1).permute(
+                    0, 2, 1, 3
+                )  # (B T) N C -> B N T C
+                j += 1
+        tokens = tokens[:, : N - self.num_virtual_tracks]
+
+        flow = self.flow_head(tokens)
+        if self.linear_layer_for_vis_conf:
+            vis_conf = self.vis_conf_head(tokens)
+            flow = torch.cat([flow, vis_conf], dim=-1)
+
+        return flow
+    
+class BasicMotionEncoder(nn.Module):
+    def __init__(self, corr_channel, dim=128, pdim=2):
+        super(BasicMotionEncoder, self).__init__()
+        self.pdim = pdim
+        self.convc1 = nn.Conv2d(corr_channel, dim*4, 1, padding=0)
+        self.convc2 = nn.Conv2d(dim*4, dim+dim//2, 3, padding=1)
+        if pdim==2 or pdim==4:
+            self.convf1 = nn.Conv2d(pdim, dim*2, 5, padding=2)
+            self.convf2 = nn.Conv2d(dim*2, dim//2, 3, padding=1)
+            self.conv = nn.Conv2d(dim*2, dim-pdim, 3, padding=1)
+        else:
+            self.conv = nn.Conv2d(dim+dim//2+pdim, dim, 3, padding=1)
+
+    def forward(self, flow, corr):
+        cor = F.relu(self.convc1(corr))
+        cor = F.relu(self.convc2(cor))
+        if self.pdim==2 or self.pdim==4:
+            flo = F.relu(self.convf1(flow))
+            flo = F.relu(self.convf2(flo))
+            cor_flo = torch.cat([cor, flo], dim=1)
+            out = F.relu(self.conv(cor_flo))
+            return torch.cat([out, flow], dim=1)
+        else:
+            # the flow is already encoded to something nice
+            cor_flo = torch.cat([cor, flow], dim=1)
+            return F.relu(self.conv(cor_flo))
+            # return torch.cat([out, flow], dim=1)
+    
+def conv133_encoder(input_dim, dim, expansion_factor=4):
+    return nn.Sequential(
+        nn.Conv2d(input_dim, dim*expansion_factor, kernel_size=1),
+        nn.GELU(),
+        nn.Conv2d(dim*expansion_factor, dim*expansion_factor, kernel_size=3, padding=1),
+        nn.GELU(),
+        nn.Conv2d(dim*expansion_factor, dim, kernel_size=3, padding=1),
+    )        
+    
+class BasicUpdateBlock(nn.Module):
+    def __init__(self, corr_channel, num_blocks, hdim=128, cdim=128):
+        # flowfeat is hdim; ctxfeat is dim. typically hdim==cdim.
+        super(BasicUpdateBlock, self).__init__()
+        self.encoder = BasicMotionEncoder(corr_channel, dim=cdim)
+        self.compressor = conv1x1(2*cdim+hdim, hdim)
+        
+        self.refine = []
+        for i in range(num_blocks):
+            self.refine.append(CNBlock1d(hdim, hdim))
+            self.refine.append(CNBlock2d(hdim, hdim))
+        self.refine = nn.ModuleList(self.refine)
+
+    def forward(self, flowfeat, ctxfeat, corr, flow, S, upsample=True):
+        BS,C,H,W = flowfeat.shape
+        B = BS//S
+
+        # with torch.no_grad():
+        motion_features = self.encoder(flow, corr)
+        flowfeat = self.compressor(torch.cat([flowfeat, ctxfeat, motion_features], dim=1))
+            
+        for blk in self.refine:
+            flowfeat = blk(flowfeat, S)
+        return flowfeat
+    
+class FullUpdateBlock(nn.Module):
+    def __init__(self, corr_channel, num_blocks, hdim=128, cdim=128, pdim=2, use_attn=False):
+        # flowfeat is hdim; ctxfeat is dim. typically hdim==cdim.
+        super(FullUpdateBlock, self).__init__()
+        self.encoder = BasicMotionEncoder(corr_channel, dim=cdim, pdim=pdim)
+
+        # note we have hdim==cdim
+        # compressor chans:
+        #   dim for flowfeat
+        #   dim for ctxfeat
+        #   dim for motion_features
+        #   pdim for flow (if p 2, like if we give sincos(relflow))
+        #   2 for visconf
+        
+        if pdim==2:
+            # hdim==cdim
+            # dim for flowfeat
+            # dim for ctxfeat
+            # dim for motion_features
+            # 2 for visconf
+            self.compressor = conv1x1(2*cdim+hdim+2, hdim) 
+        else:
+            # we concatenate the flow info again, to not lose it (e.g., from the relu)
+            self.compressor = conv1x1(2*cdim+hdim+2+pdim, hdim)
+        
+        self.refine = []
+        for i in range(num_blocks):
+            self.refine.append(CNBlock1d(hdim, hdim, use_attn=use_attn))
+            self.refine.append(CNBlock2d(hdim, hdim))
+        self.refine = nn.ModuleList(self.refine)
+
+    def forward(self, flowfeat, ctxfeat, visconf, corr, flow, S, upsample=True):
+        BS,C,H,W = flowfeat.shape
+        B = BS//S
+
+        # print('flowfeat', flowfeat.shape)
+        # print('ctxfeat', ctxfeat.shape)
+        # print('visconf', visconf.shape)
+        # print('corr', corr.shape)
+        # print('flow', flow.shape)
+
+        motion_features = self.encoder(flow, corr)
+        
+        # print('cat', torch.cat([flowfeat, ctxfeat, motion_features, visconf, flow], dim=1).shape)
+        flowfeat = self.compressor(torch.cat([flowfeat, ctxfeat, motion_features, visconf], dim=1))
+        for blk in self.refine:
+            flowfeat = blk(flowfeat, S)
+        return flowfeat
+
+class MixerUpdateBlock(nn.Module):
+    def __init__(self, corr_channel, num_blocks, hdim=128, cdim=128):
+        # flowfeat is hdim; ctxfeat is dim. typically hdim==cdim.
+        super(MixerUpdateBlock, self).__init__()
+        self.encoder = BasicMotionEncoder(corr_channel, dim=cdim)
+        self.compressor = conv1x1(2*cdim+hdim, hdim)
+        
+        self.refine = []
+        for i in range(num_blocks):
+            self.refine.append(CNBlock1d(hdim, hdim, use_mixer=True))
+            self.refine.append(CNBlock2d(hdim, hdim))
+        self.refine = nn.ModuleList(self.refine)
+
+    def forward(self, flowfeat, ctxfeat, corr, flow, S, upsample=True):
+        BS,C,H,W = flowfeat.shape
+        B = BS//S
+
+        # with torch.no_grad():
+        motion_features = self.encoder(flow, corr)
+        flowfeat = self.compressor(torch.cat([flowfeat, ctxfeat, motion_features], dim=1))
+            
+        for ii, blk in enumerate(self.refine):
+            flowfeat = blk(flowfeat, S)
+        return flowfeat
+    
+class FacUpdateBlock(nn.Module):
+    def __init__(self, corr_channel, num_blocks, hdim=128, cdim=128, pdim=84, use_attn=False):
+        super(FacUpdateBlock, self).__init__()
+        self.corr_encoder = conv133_encoder(corr_channel, cdim)
+        # note we have hdim==cdim
+        # compressor chans:
+        #   dim for flowfeat
+        #   dim for ctxfeat
+        #   dim for corr
+        #   pdim for flow
+        #   2 for visconf
+        self.compressor = conv1x1(2*cdim+hdim+2+pdim, hdim)
+        self.refine = []
+        for i in range(num_blocks):
+            self.refine.append(CNBlock1d(hdim, hdim, use_attn=use_attn))
+            self.refine.append(CNBlock2d(hdim, hdim))
+        self.refine = nn.ModuleList(self.refine)
+
+    def forward(self, flowfeat, ctxfeat, visconf, corr, flow, S, upsample=True):
+        BS,C,H,W = flowfeat.shape
+        B = BS//S
+
+        # print('flowfeat', flowfeat.shape)
+        # print('ctxfeat', ctxfeat.shape)
+        # print('visconf', visconf.shape)
+        # print('corr', corr.shape)
+        # print('flow', flow.shape)
+        corr = self.corr_encoder(corr)
+        # print('cat', torch.cat([flowfeat, ctxfeat, motion_features, visconf, flow], dim=1).shape)
+        flowfeat = self.compressor(torch.cat([flowfeat, ctxfeat, corr, visconf, flow], dim=1))
+        for blk in self.refine:
+            flowfeat = blk(flowfeat, S)
+        return flowfeat
+    
+class CleanUpdateBlock(nn.Module):
+    def __init__(self, corr_channel, num_blocks, cdim=128, hdim=256, pdim=84, use_attn=False, use_layer_scale=True):
+        super(CleanUpdateBlock, self).__init__()
+        self.corr_encoder = conv133_encoder(corr_channel, cdim)
+        # compressor chans:
+        #   cdim for flowfeat
+        #   cdim for ctxfeat
+        #   cdim for corrfeat
+        #   pdim for flow
+        #   2 for visconf
+        self.compressor = conv1x1(3*cdim+pdim+2, hdim)
+        self.refine = []
+        for i in range(num_blocks):
+            self.refine.append(CNBlock1d(hdim, hdim, use_attn=use_attn, use_layer_scale=use_layer_scale))
+            self.refine.append(CNBlock2d(hdim, hdim, use_layer_scale=use_layer_scale))
+        self.refine = nn.ModuleList(self.refine)
+        self.final_conv = conv1x1(hdim, cdim)
+
+    def forward(self, flowfeat, ctxfeat, visconf, corr, flow, S, upsample=True):
+        BS,C,H,W = flowfeat.shape
+        B = BS//S
+        corrfeat = self.corr_encoder(corr)
+        flowfeat = self.compressor(torch.cat([flowfeat, ctxfeat, corrfeat, flow, visconf], dim=1))
+        for blk in self.refine:
+            flowfeat = blk(flowfeat, S)
+        flowfeat = self.final_conv(flowfeat)
+        return flowfeat
+
+class RelUpdateBlock(nn.Module):
+    def __init__(self, corr_channel, num_blocks, cdim=128, hdim=128, pdim=4, use_attn=True, use_mixer=False, use_conv=False, use_convb=False, use_layer_scale=True, no_time=False, no_space=False, final_space=False, no_ctx=False):
+        super(RelUpdateBlock, self).__init__()
+        self.motion_encoder = BasicMotionEncoder(corr_channel, dim=hdim, pdim=pdim) # B,hdim,H,W
+        self.no_ctx = no_ctx
+        if no_ctx:
+            self.compressor = conv1x1(cdim+hdim+2, hdim)
+        else:
+            self.compressor = conv1x1(2*cdim+hdim+2, hdim)
+        self.refine = []
+        for i in range(num_blocks):
+            if not no_time:
+                self.refine.append(CNBlock1d(hdim, hdim, use_attn=use_attn, use_mixer=use_mixer, use_conv=use_conv, use_convb=use_convb, use_layer_scale=use_layer_scale))
+            if not no_space:
+                self.refine.append(CNBlock2d(hdim, hdim, use_layer_scale=use_layer_scale))
+        if final_space:
+            self.refine.append(CNBlock2d(hdim, hdim, use_layer_scale=use_layer_scale))
+        self.refine = nn.ModuleList(self.refine)
+        self.final_conv = conv1x1(hdim, cdim)
+
+    def forward(self, flowfeat, ctxfeat, visconf, corr, flow, S, upsample=True):
+        BS,C,H,W = flowfeat.shape
+        B = BS//S
+        motion_features = self.motion_encoder(flow, corr)
+        if self.no_ctx:
+            flowfeat = self.compressor(torch.cat([flowfeat, motion_features, visconf], dim=1))
+        else:
+            flowfeat = self.compressor(torch.cat([flowfeat, ctxfeat, motion_features, visconf], dim=1))
+        for blk in self.refine:
+            flowfeat = blk(flowfeat, S)
+        flowfeat = self.final_conv(flowfeat)
+        return flowfeat

nets/net34.py

@@ -0,0 +1,647 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import utils.samp
+import utils.misc
+import numpy as np
+
+from nets.blocks import CNBlockConfig, ConvNeXt, conv1x1, RelUpdateBlock, InputPadder, CorrBlock, BasicEncoder
+
+# init mostly from raft82 in alltrack
+
+class Net(nn.Module):
+    def __init__(
+            self,
+            seqlen,
+            noise_level=0,
+            use_attn=True,
+            use_mixer=False,
+            use_conv=False,
+            use_convb=False,
+            use_basicencoder=False,
+            use_sinmotion=False,
+            use_relmotion=False,
+            use_sinrelmotion=False,
+            use_feats8=False,
+            no_time=False,
+            no_space=False,
+            final_space=False,
+            no_split=False,
+            no_ctx=False,
+            full_split=False,
+            half_corr=False,
+            corr_levels=5,
+            corr_radius=4,
+            num_blocks=3,
+            dim=128,
+            hdim=128,
+            init_weights=True,
+    ):
+        super(Net, self).__init__()
+
+        self.dim = dim
+        self.hdim = hdim
+
+        self.noise_level = noise_level
+        self.no_time = no_time
+        self.no_space = no_space
+        self.final_space = final_space
+        self.seqlen = seqlen
+        self.corr_levels = corr_levels
+        self.corr_radius = corr_radius
+        self.corr_channel = self.corr_levels * (self.corr_radius * 2 + 1) ** 2
+        self.num_blocks = num_blocks
+
+        self.use_feats8 = use_feats8
+        self.use_basicencoder = use_basicencoder
+        self.use_sinmotion = use_sinmotion
+        self.use_relmotion = use_relmotion
+        self.use_sinrelmotion = use_sinrelmotion
+        self.no_split = no_split
+        self.no_ctx = no_ctx
+        self.full_split = full_split
+        self.half_corr = half_corr
+
+        if use_basicencoder:
+            if self.full_split:
+                self.fnet = BasicEncoder(input_dim=3, output_dim=self.dim, stride=8)
+                self.cnet = BasicEncoder(input_dim=3, output_dim=self.dim, stride=8)
+            else:
+                if self.no_split:
+                    self.fnet = BasicEncoder(input_dim=3, output_dim=self.dim, stride=8)
+                else:
+                    self.fnet = BasicEncoder(input_dim=3, output_dim=self.dim*2, stride=8)
+        else:
+            block_setting = [
+                CNBlockConfig(96, 192, 3, True), # 4x
+                CNBlockConfig(192, 384, 3, False), # 8x
+                CNBlockConfig(384, None, 9, False), # 8x
+            ]
+            self.cnn = ConvNeXt(block_setting, stochastic_depth_prob=0.0, init_weights=init_weights)
+            if self.no_split:
+                self.dot_conv = conv1x1(384, dim)
+            else:
+                self.dot_conv = conv1x1(384, dim*2)
+            
+        # # conv for iter 0 results
+        # self.init_conv = conv3x3(2 * dim, 2 * dim)
+        self.upsample_weight = nn.Sequential(
+            # convex combination of 3x3 patches
+            nn.Conv2d(dim, dim * 2, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(dim * 2, 64 * 9, 1, padding=0)
+        )
+        self.flow_head = nn.Sequential(
+            nn.Conv2d(dim, 2*dim, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(2*dim, 2, kernel_size=3, padding=1)
+        )
+        self.visconf_head = nn.Sequential(
+            nn.Conv2d(dim, 2*dim, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(2*dim, 2, kernel_size=3, padding=1)
+        )
+
+        if self.use_sinrelmotion:
+            self.pdim = 84 # 32*2
+        elif self.use_relmotion:
+            self.pdim = 4
+        elif self.use_sinmotion:
+            self.pdim = 42
+        else:
+            self.pdim = 2
+            
+        self.update_block = RelUpdateBlock(self.corr_channel, self.num_blocks, cdim=dim, hdim=hdim, pdim=self.pdim,
+                                           use_attn=use_attn, use_mixer=use_mixer, use_conv=use_conv, use_convb=use_convb,
+                                           use_layer_scale=True, no_time=no_time, no_space=no_space, final_space=final_space,
+                                           no_ctx=no_ctx)
+
+        time_line = torch.linspace(0, seqlen-1, seqlen).reshape(1, seqlen, 1)
+        self.register_buffer("time_emb", utils.misc.get_1d_sincos_pos_embed_from_grid(self.dim, time_line[0])) # 1,S,C
+        
+    def fetch_time_embed(self, t, dtype, is_training=False):
+        S = self.time_emb.shape[1]
+        # print('fetching time_embed for t', t, '(we have %d)' % (S))
+        # print('self.time_emb', self.time_emb.shape)
+        if t == S:
+            return self.time_emb.to(dtype)
+        elif t==1:
+            if is_training:
+                ind = np.random.choice(S)
+                return self.time_emb[:,ind:ind+1].to(dtype)
+            else:
+                return self.time_emb[:,1:2].to(dtype)
+        else:
+            time_emb = self.time_emb.float()
+            time_emb = F.interpolate(time_emb.permute(0, 2, 1), size=t, mode="linear").permute(0, 2, 1) 
+            return time_emb.to(dtype)
+    
+    def coords_grid(self, batch, ht, wd, device, dtype):
+        coords = torch.meshgrid(torch.arange(ht, device=device, dtype=dtype), torch.arange(wd, device=device, dtype=dtype))
+        coords = torch.stack(coords[::-1], dim=0)
+        return coords[None].repeat(batch, 1, 1, 1)
+
+    def initialize_flow(self, img):
+        """ Flow is represented as difference between two coordinate grids flow = coords2 - coords1"""
+        N, C, H, W = img.shape
+        coords1 = self.coords_grid(N, H//8, W//8, device=img.device)
+        coords2 = self.coords_grid(N, H//8, W//8, device=img.device)
+        return coords1, coords2
+
+    def upsample_data(self, flow, mask):
+        """ Upsample [H/8, W/8, C] -> [H, W, C] using convex combination """
+        N, C, H, W = flow.shape
+        mask = mask.view(N, 1, 9, 8, 8, H, W)
+        mask = torch.softmax(mask, dim=2)
+
+        up_flow = F.unfold(8 * flow, [3,3], padding=1)
+        up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)
+
+        up_flow = torch.sum(mask * up_flow, dim=2)
+        up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
+        
+        return up_flow.reshape(N, 2, 8*H, 8*W).to(flow.dtype)
+
+    def get_T_padded_images(self, images, T, S, is_training, stride=None, pad=True):
+        B,T,C,H,W = images.shape
+        indices = None
+        if T > 2:
+            step = S // 2 if stride is None else stride
+            # starts = list(range(step,max(T,S),step))
+            # indices = [mid-step for mid in mids]
+            indices = []
+            start = 0
+            while start + S < T:
+                indices.append(start)
+                start += step
+            indices.append(start)
+            Tpad = indices[-1]+S-T
+            # print(indices, Tpad)
+            # import pdb; pdb.set_trace()
+            if pad:
+                if is_training:
+                    assert Tpad == 0
+                else:
+                    images = images.reshape(B,1,T,C*H*W)
+                    if Tpad > 0:
+                        padding_tensor = images[:,:,-1:,:].expand(B,1,Tpad,C*H*W)
+                        images = torch.cat([images, padding_tensor], dim=2)
+                    images = images.reshape(B,T+Tpad,C,H,W)
+                    T = T+Tpad
+        else:
+            assert T == 2
+        return images, T, indices
+
+    def get_fmaps(self, images_, B, T, sw, is_training, nograd_backbone):
+        _, _, H_pad, W_pad = images_.shape # revised HW
+
+        C, H8, W8 = self.dim*2, H_pad//8, W_pad//8
+        if self.no_split:
+            C = self.dim
+
+        fmaps_chunk_size = 64
+        if (not is_training) and (T > fmaps_chunk_size):
+            images = images_.reshape(B,T,3,H_pad,W_pad)
+            fmaps = []
+            for t in range(0, T, fmaps_chunk_size):
+                images_chunk = images[:, t : t + fmaps_chunk_size]
+                images_chunk = images_chunk.cuda()
+                if self.use_basicencoder:
+                    if self.full_split:
+                        fmaps_chunk1 = self.fnet(images_chunk.reshape(-1, 3, H_pad, W_pad))
+                        fmaps_chunk2 = self.cnet(images_chunk.reshape(-1, 3, H_pad, W_pad))
+                        fmaps_chunk = torch.cat([fmaps_chunk1, fmaps_chunk2], axis=1)
+                    else:
+                        fmaps_chunk = self.fnet(images_chunk.reshape(-1, 3, H_pad, W_pad))
+                else:
+                    fmaps_chunk = self.cnn(images_chunk.reshape(-1, 3, H_pad, W_pad))
+                    if t==0 and sw is not None and sw.save_this:
+                        sw.summ_feat('1_model/fmap_raw', fmaps_chunk[0:1])
+                    fmaps_chunk = self.dot_conv(fmaps_chunk) # B*T,C,H8,W8
+                T_chunk = images_chunk.shape[1]
+                fmaps.append(fmaps_chunk.reshape(B, -1, C, H8, W8))
+            fmaps_ = torch.cat(fmaps, dim=1).reshape(-1, C, H8, W8)
+        else:
+            if not is_training:
+                # sometimes we need to move things to cuda here
+                images_ = images_.cuda()
+            if self.use_basicencoder:
+                if self.full_split:
+                    # if self.half_corr:
+                    fmaps1_ = self.fnet(images_)
+                    fmaps2_ = self.cnet(images_)
+                    fmaps_ = torch.cat([fmaps1_, fmaps2_], axis=1)
+                else:
+                    fmaps_ = self.fnet(images_)
+            else:
+                if nograd_backbone:
+                    with torch.no_grad():
+                        fmaps_ = self.cnn(images_)
+                else:
+                    fmaps_ = self.cnn(images_)
+                if sw is not None and sw.save_this:
+                    sw.summ_feat('1_model/fmap_raw', fmaps_[0:1])
+                fmaps_ = self.dot_conv(fmaps_) # B*T,C,H8,W8
+        return fmaps_
+    
+    def forward(self, images, iters=6, sw=None, nograd_backbone=False, is_training=False, stride=None):
+        B,T,C,H,W = images.shape
+        S = self.seqlen
+        device = images.device
+        dtype = images.dtype
+        # print('images', images.shape, 'device', device)
+
+        T_bak = T
+        if stride is not None:
+            pad = False
+        else:
+            pad = True
+        images, T, indices = self.get_T_padded_images(images, T, S, is_training, stride=stride, pad=pad)
+        # print('indices', indices)
+
+        images = images.contiguous()
+        images_ = images.reshape(B*T,3,H,W)
+        padder = InputPadder(images_.shape)
+        images_ = padder.pad(images_)[0]
+
+        _, _, H_pad, W_pad = images_.shape # revised HW
+        C, H8, W8 = self.dim*2, H_pad//8, W_pad//8
+        C2 = C//2
+        if self.no_split:
+            C = self.dim
+            C2 = C
+
+        fmaps = self.get_fmaps(images_, B, T, sw, is_training, nograd_backbone).reshape(B,T,C,H8,W8)
+        # print('fmaps_', fmaps_.shape)
+        device = fmaps.device
+
+        # if sw is not None and sw.save_this:
+        #     sw.summ_feat('1_model/fmap_dc', fmaps_[0:1])
+
+        # fmaps = fmaps_.reshape(B,T,C,H8,W8)
+        # del fmaps_
+        fmap_anchor = fmaps[:,0]
+
+        # if not is_training:
+        #     del images
+        #     del images_
+
+        if T<=2 or is_training:
+            # note: collecting preds can get expensive on a long video
+            all_flow_preds = []
+            all_visconf_preds = []
+        else:
+            all_flow_preds = None
+            all_visconf_preds = None
+
+        if T > 2: # multiframe tracking
+            
+            # we will store our final outputs in these tensors
+            full_flows = torch.zeros((B,T,2,H,W), dtype=dtype, device=device)
+            full_visconfs = torch.zeros((B,T,2,H,W), dtype=dtype, device=device)
+            # 1/8 resolution 
+            full_flows8 = torch.zeros((B,T,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+            full_visconfs8 = torch.zeros((B,T,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+
+            if is_training and self.noise_level and np.random.rand() < 0.1:
+                # full_flows8 = 2.0*torch.randn((B,T,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+                # print('flows8 += randn4, const on time')
+                # full_flows8 = float(self.noise_level)*torch.randn((B,T,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+                # print('noise const on time')
+                full_flows8 += np.random.rand()*float(self.noise_level)*torch.randn((B,1,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+
+            if self.use_feats8:
+                full_feats8 = torch.zeros((B,T,C2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+            visits = np.zeros((T))
+
+            for ii, ind in enumerate(indices):
+                ara = np.arange(ind,ind+S)
+                if ii < len(indices)-1:
+                    next_ind = indices[ii+1]
+                    next_ara = np.arange(next_ind,next_ind+S)
+                
+                # print("torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024), 'ara', ara)
+                # print('ara', ara)
+                fmaps2 = fmaps[:,ara]
+                flows8 = full_flows8[:,ara].reshape(B*(S),2,H_pad//8,W_pad//8).detach()
+                visconfs8 = full_visconfs8[:,ara].reshape(B*(S),2,H_pad//8,W_pad//8).detach()
+
+                if self.use_feats8:
+                    if ind==0:
+                        feats8 = None
+                    else:
+                        feats8 = full_feats8[:,ara].reshape(B*(S),C2,H_pad//8,W_pad//8).detach()
+                else:
+                    feats8 = None
+
+                flow_predictions, visconf_predictions, flows8, visconfs8, feats8 = self.forward_window(
+                    fmap_anchor, fmaps2, visconfs8, iters=iters, flowfeat=feats8, flows8=flows8,
+                    is_training=is_training)
+
+                unpad_flow_predictions = []
+                unpad_visconf_predictions = []
+                for i in range(len(flow_predictions)):
+                    flow_predictions[i] = padder.unpad(flow_predictions[i])
+                    unpad_flow_predictions.append(flow_predictions[i].reshape(B,S,2,H,W))
+                    visconf_predictions[i] = padder.unpad(torch.sigmoid(visconf_predictions[i]))
+                    # print('visconf_predictions[%d]' % i, visconf_predictions[i].shape)
+                    unpad_visconf_predictions.append(visconf_predictions[i].reshape(B,S,2,H,W))
+
+                full_flows[:,ara] = unpad_flow_predictions[-1].reshape(B,S,2,H,W)
+                full_flows8[:,ara] = flows8.reshape(B,S,2,H_pad//8,W_pad//8)
+                full_visconfs[:,ara] = unpad_visconf_predictions[-1].reshape(B,S,2,H,W)
+                full_visconfs8[:,ara] = visconfs8.reshape(B,S,2,H_pad//8,W_pad//8)
+                if self.use_feats8:
+                    full_feats8[:,ara] = feats8.reshape(B,S,C2,H_pad//8,W_pad//8)
+                visits[ara] += 1
+
+                if is_training:
+                    all_flow_preds.append(unpad_flow_predictions)
+                    all_visconf_preds.append(unpad_visconf_predictions)
+                else:
+                    del unpad_flow_predictions
+                    del unpad_visconf_predictions
+
+                # for the next iter, replace empty data with nearest available preds
+                invalid_idx = np.where(visits==0)[0]
+                valid_idx = np.where(visits>0)[0]
+                for idx in invalid_idx:
+                    nearest = valid_idx[np.argmin(np.abs(valid_idx - idx))]
+                    # print('replacing %d with %d' % (idx, nearest))
+                    full_flows8[:,idx] = full_flows8[:,nearest]
+                    full_visconfs8[:,idx] = full_visconfs8[:,nearest]
+                    if self.use_feats8:
+                        full_feats8[:,idx] = full_feats8[:,nearest]
+        else: # flow
+
+            flows8 = torch.zeros((B,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+            visconfs8 = torch.zeros((B,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+
+            if is_training and self.noise_level and np.random.rand() < 0.1:
+                # full_flows8 = 2.0*torch.randn((B,T,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+                # print('flows8 += randn4, const on time')
+                # full_flows8 = float(self.noise_level)*torch.randn((B,T,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+                # print('noise on flow too')
+                flows8 += np.random.rand()*float(self.noise_level)*torch.randn((B,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+            
+            flow_predictions, visconf_predictions, flows8, visconfs8, feats8 = self.forward_window(
+                fmap_anchor, fmaps[:,1:2], visconfs8, iters=iters, flowfeat=None, flows8=flows8,
+                is_training=is_training)
+            unpad_flow_predictions = []
+            unpad_visconf_predictions = []
+            for i in range(len(flow_predictions)):
+                flow_predictions[i] = padder.unpad(flow_predictions[i])
+                all_flow_preds.append(flow_predictions[i].reshape(B,2,H,W))
+                visconf_predictions[i] = padder.unpad(torch.sigmoid(visconf_predictions[i]))
+                all_visconf_preds.append(visconf_predictions[i].reshape(B,2,H,W))
+            full_flows = all_flow_preds[-1].reshape(B,2,H,W)
+            full_visconfs = all_visconf_preds[-1].reshape(B,2,H,W)
+                
+        if (not is_training) and (T > 2):
+            # print('full_flows', full_flows.shape)
+            # print('T_bak', T_bak)
+            full_flows = full_flows[:,:T_bak]
+            full_visconfs = full_visconfs[:,:T_bak]
+            # print('full_flows trim', full_flows.shape)
+            
+        return full_flows, full_visconfs, all_flow_preds, all_visconf_preds#, bak_flows8
+    
+    def forward_sliding(self, images, iters=6, sw=None, nograd_backbone=False, is_training=False, window_len=None, stride=None):
+        B,T,C,H,W = images.shape
+        S = self.seqlen if window_len is None else window_len
+        device = images.device
+        dtype = images.dtype
+        # print('images', images.shape, 'device', device)
+        stride = S // 2 if stride is None else stride
+
+        T_bak = T
+        images, T, indices = self.get_T_padded_images(images, T, S, is_training, stride)
+        # print('indices', indices)
+        assert stride <= S // 2
+
+        images = images.contiguous()
+        images_ = images.reshape(B*T,3,H,W)
+        padder = InputPadder(images_.shape)
+        images_ = padder.pad(images_)[0]
+
+        _, _, H_pad, W_pad = images_.shape # revised HW
+        C, H8, W8 = self.dim*2, H_pad//8, W_pad//8
+        C2 = C//2
+        if self.no_split:
+            C = self.dim
+            C2 = C
+            
+        all_flow_preds = None
+        all_visconf_preds = None
+        
+        if T<=2:
+            # note: collecting preds can get expensive on a long video
+            all_flow_preds = []
+            all_visconf_preds = []
+            
+            fmaps = self.get_fmaps(images_, B, T, sw, is_training, nograd_backbone).reshape(B,T,C,H8,W8)
+            device = fmaps.device
+            
+            flows8 = torch.zeros((B,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+            visconfs8 = torch.zeros((B,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+
+            if is_training and self.noise_level and np.random.rand() < 0.1:
+                # full_flows8 = 2.0*torch.randn((B,T,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+                # print('flows8 += randn4, const on time')
+                # full_flows8 = float(self.noise_level)*torch.randn((B,T,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+                # print('noise on flow too')
+                flows8 += np.random.rand()*float(self.noise_level)*torch.randn((B,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+                
+            fmap_anchor = fmaps[:,0]
+            
+            flow_predictions, visconf_predictions, flows8, visconfs8, feats8 = self.forward_window(
+                fmap_anchor, fmaps[:,1:2], visconfs8, iters=iters, flowfeat=None, flows8=flows8,
+                is_training=is_training)
+            unpad_flow_predictions = []
+            unpad_visconf_predictions = []
+            for i in range(len(flow_predictions)):
+                flow_predictions[i] = padder.unpad(flow_predictions[i])
+                all_flow_preds.append(flow_predictions[i].reshape(B,2,H,W))
+                visconf_predictions[i] = padder.unpad(torch.sigmoid(visconf_predictions[i]))
+                all_visconf_preds.append(visconf_predictions[i].reshape(B,2,H,W))
+            full_flows = all_flow_preds[-1].reshape(B,2,H,W).detach().cpu()
+            full_visconfs = all_visconf_preds[-1].reshape(B,2,H,W).detach().cpu()
+            
+            return full_flows, full_visconfs, all_flow_preds, all_visconf_preds#, bak_flows8
+
+        assert T > 2 # multiframe tracking
+        
+        if is_training:
+            all_flow_preds = []
+            all_visconf_preds = []
+            
+        # # we will store our final outputs in these tensors cpu
+        full_flows = torch.zeros((B,T,2,H,W), dtype=dtype, device='cpu')
+        full_visconfs = torch.zeros((B,T,2,H,W), dtype=dtype, device='cpu')
+        
+        images_ = images_.reshape(B,T,3,H_pad,W_pad)
+        fmap_anchor = self.get_fmaps(images_[:,:1].reshape(-1,3,H_pad,W_pad), B, 1, sw, is_training, nograd_backbone).reshape(B,C,H8,W8)
+        device = fmap_anchor.device
+        full_visited = torch.zeros((T,), dtype=torch.bool, device=device)
+
+        for ii, ind in enumerate(indices):
+            ara = np.arange(ind,ind+S)
+            if ii == 0:
+                flows8 = torch.zeros((B,S,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+                visconfs8 = torch.zeros((B,S,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+                fmaps2 = self.get_fmaps(images_[:,ara].reshape(-1,3,H_pad,W_pad), B, S, sw, is_training, nograd_backbone).reshape(B,S,C,H8,W8)
+                if is_training and self.noise_level and np.random.rand() < 0.1:
+                    flows8 += np.random.rand()*float(self.noise_level)*torch.randn((B,1,2,H_pad//8,W_pad//8), dtype=dtype, device=device)
+            else:
+                # import pdb; pdb.set_trace()
+                flows8 = torch.cat([flows8[:,stride:stride+S//2], flows8[:,stride+S//2-1:stride+S//2].repeat(1,S//2,1,1,1)], dim=1)
+                visconfs8 = torch.cat([visconfs8[:,stride:stride+S//2], visconfs8[:,stride+S//2-1:stride+S//2].repeat(1,S//2,1,1,1)], dim=1)
+                fmaps2 = torch.cat([fmaps2[:,stride:stride+S//2], 
+                                    self.get_fmaps(images_[:,np.arange(ind+S//2,ind+S)].reshape(-1,3,H_pad,W_pad), B, S//2, sw, is_training, nograd_backbone).reshape(B,S//2,C,H8,W8)], dim=1)
+            
+            flows8 = flows8.reshape(B*S,2,H_pad//8,W_pad//8).detach()
+            visconfs8 = visconfs8.reshape(B*S,2,H_pad//8,W_pad//8).detach()
+            
+            # print("torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024), 'ara', ara)
+            # print('ara', ara)
+            flow_predictions, visconf_predictions, flows8, visconfs8, _ = self.forward_window(
+                fmap_anchor, fmaps2, visconfs8, iters=iters, flowfeat=None, flows8=flows8,
+                is_training=is_training)
+
+            unpad_flow_predictions = []
+            unpad_visconf_predictions = []
+            for i in range(len(flow_predictions)):
+                flow_predictions[i] = padder.unpad(flow_predictions[i])
+                unpad_flow_predictions.append(flow_predictions[i].reshape(B,S,2,H,W))
+                visconf_predictions[i] = padder.unpad(torch.sigmoid(visconf_predictions[i]))
+                # print('visconf_predictions[%d]' % i, visconf_predictions[i].shape)
+                unpad_visconf_predictions.append(visconf_predictions[i].reshape(B,S,2,H,W))
+
+            current_visiting = torch.zeros((T,), dtype=torch.bool, device=device)
+            current_visiting[ara] = True
+            
+            to_fill = current_visiting & (~full_visited)
+            to_fill_sum = to_fill.sum().item()
+            full_flows[:,to_fill] = unpad_flow_predictions[-1].reshape(B,S,2,H,W)[:,-to_fill_sum:].detach().cpu()
+            full_visconfs[:,to_fill] = unpad_visconf_predictions[-1].reshape(B,S,2,H,W)[:,-to_fill_sum:].detach().cpu()
+            full_visited |= current_visiting
+
+            if is_training:
+                all_flow_preds.append(unpad_flow_predictions)
+                all_visconf_preds.append(unpad_visconf_predictions)
+            else:
+                del unpad_flow_predictions
+                del unpad_visconf_predictions
+                
+            flows8 = flows8.reshape(B,S,2,H_pad//8,W_pad//8)
+            visconfs8 = visconfs8.reshape(B,S,2,H_pad//8,W_pad//8)
+                
+        if not is_training:
+            # print('full_flows', full_flows.shape)
+            # print('T_bak', T_bak)
+            full_flows = full_flows[:,:T_bak]
+            full_visconfs = full_visconfs[:,:T_bak]
+            # print('full_flows trim', full_flows.shape)
+            
+        return full_flows, full_visconfs, all_flow_preds, all_visconf_preds#, bak_flows8
+        
+    def forward_window(self, fmap1_single, fmaps2, visconfs8, iters=None, flowfeat=None, flows8=None, sw=None, is_training=False):
+        B,S,C,H8,W8 = fmaps2.shape
+        device = fmaps2.device
+        dtype = fmaps2.dtype
+
+        flow_predictions = []
+        visconf_predictions = []
+
+        # print('fmap1_single', fmap1_single.shape)
+        # print('fmaps2', fmaps2.shape)
+
+        fmap1 = fmap1_single.unsqueeze(1).repeat(1,S,1,1,1) # B,S,C,H,W
+        # print('fmap1', fmap1.shape)
+        fmap1 = fmap1.reshape(B*(S),C,H8,W8).contiguous()
+        # print('fmap1', fmap1.shape)
+
+        fmap2 = fmaps2.reshape(B*(S),C,H8,W8).contiguous()
+
+        visconfs8 = visconfs8.reshape(B*(S),2,H8,W8).contiguous()
+
+        if not self.half_corr:
+            corr_fn = CorrBlock(fmap1, fmap2, self.corr_levels, self.corr_radius)
+
+        coords1 = self.coords_grid(B*(S), H8, W8, device=fmap1.device, dtype=dtype)
+
+        if self.no_split:
+            flowfeat, ctxfeat = fmap1.clone(), fmap1.clone()
+        else:
+            if flowfeat is not None:
+                _, ctxfeat = torch.split(fmap1, [self.dim, self.dim], dim=1)
+            else:
+                flowfeat, ctxfeat = torch.split(fmap1, [self.dim, self.dim], dim=1)
+
+            if self.half_corr:
+                # we discard the cnet output for fmap2
+                # in general maybe we shouldn't compute it, if half_corr pays off
+                flowfeat2, _ = torch.split(fmap2, [self.dim, self.dim], dim=1)
+                corr_fn = CorrBlock(flowfeat, flowfeat2, self.corr_levels, self.corr_radius)
+                
+        # add pos emb to ctxfeat (and not flowfeat), since ctxfeat is untouched across iters
+        time_emb = self.fetch_time_embed(S, ctxfeat.dtype, is_training).reshape(1,S,self.dim,1,1).repeat(B,1,1,1,1)
+        ctxfeat = ctxfeat + time_emb.reshape(B*S,self.dim,1,1)
+
+        if self.no_ctx:
+            flowfeat = flowfeat + time_emb.reshape(B*S,self.dim,1,1)
+            
+        for itr in range(iters):
+            _, _, H8, W8 = flows8.shape
+            flows8 = flows8.detach()
+            coords2 = (coords1 + flows8).detach() # B*S,2,H,W
+            corr = corr_fn(coords2).to(dtype)
+
+            if self.use_relmotion or self.use_sinrelmotion:
+                coords_ = coords2.reshape(B,S,2,H8*W8).permute(0,1,3,2) # B,S,H8*W8,2
+                rel_coords_forward = coords_[:, :-1] - coords_[:, 1:]
+                rel_coords_backward = coords_[:, 1:] - coords_[:, :-1]
+                rel_coords_forward = torch.nn.functional.pad(
+                    rel_coords_forward, (0, 0, 0, 0, 0, 1) # pad the 3rd-last dim (S) by (0,1)
+                )
+                rel_coords_backward = torch.nn.functional.pad(
+                    rel_coords_backward, (0, 0, 0, 0, 1, 0) # pad the 3rd-last dim (S) by (1,0)
+                )
+                rel_coords = torch.cat([rel_coords_forward, rel_coords_backward], dim=-1) # B,S,H8*W8,4
+
+                if self.use_sinrelmotion:
+                    rel_pos_emb_input = utils.misc.posenc(
+                        rel_coords,
+                        min_deg=0,
+                        max_deg=10,
+                    )  # B,S,H*W,pdim
+                    motion = rel_pos_emb_input.reshape(B*S,H8,W8,self.pdim).permute(0,3,1,2).to(dtype) # B*S,pdim,H8,W8
+                else:
+                    motion = rel_coords.reshape(B*S,H8,W8,4).permute(0,3,1,2).to(dtype) # B*S,4,H8,W8
+                
+            else:
+                if self.use_sinmotion:
+                    pos_emb_input = utils.misc.posenc(
+                        flows8.reshape(B,S,H8*W8,2),
+                        min_deg=0,
+                        max_deg=10,
+                    )  # B,S,H*W,pdim
+                    motion = pos_emb_input.reshape(B*S,H8,W8,self.pdim).permute(0,3,1,2).to(dtype) # B*S,pdim,H8,W8
+                else:
+                    motion = flows8
+                    
+            flowfeat = self.update_block(flowfeat, ctxfeat, visconfs8, corr, motion, S)
+            flow_update = self.flow_head(flowfeat)
+            visconf_update = self.visconf_head(flowfeat)
+            weight_update = .25 * self.upsample_weight(flowfeat)
+            flows8 = flows8 + flow_update
+            visconfs8 = visconfs8 + visconf_update
+            flow_up = self.upsample_data(flows8, weight_update)
+            flow_predictions.append(flow_up)
+            visconf_up = self.upsample_data(visconfs8, weight_update)
+            visconf_predictions.append(visconf_up)
+            
+        return flow_predictions, visconf_predictions, flows8, visconfs8, flowfeat#, bak_flows8
+
+    
+

requirements.txt

@@ -0,0 +1,17 @@
+numpy==1.26.4
+imageio==2.19.3
+imageio-ffmpeg==0.4.7
+gradio
+spaces
+matplotlib
+pillow
+torch==2.2.0
+torchvision==0.17.0
+albumentations
+pytorch-lightning==2.2.5
+opencv-python
+scikit-learn
+scikit-image
+einops
+tensorboardX
+transformers

utils/basic.py

@@ -0,0 +1,429 @@
+import os
+import numpy as np
+from os.path import isfile
+import torch
+import torch.nn.functional as F
+EPS = 1e-6
+import copy
+
+def sub2ind(height, width, y, x):
+    return y*width + x
+
+def ind2sub(height, width, ind):
+    y = ind // width
+    x = ind % width
+    return y, x
+    
+def get_lr_str(lr):
+    lrn = "%.1e" % lr # e.g., 5.0e-04
+    lrn = lrn[0] + lrn[3:5] + lrn[-1] # e.g., 5e-4
+    return lrn
+    
+def strnum(x):
+    s = '%g' % x
+    if '.' in s:
+        if x < 1.0:
+            s = s[s.index('.'):]
+        s = s[:min(len(s),4)]
+    return s
+
+def assert_same_shape(t1, t2):
+    for (x, y) in zip(list(t1.shape), list(t2.shape)):
+        assert(x==y)
+
+def print_stats(name, tensor):
+    shape = tensor.shape
+    tensor = tensor.detach().cpu().numpy()
+    print('%s (%s) min = %.2f, mean = %.2f, max = %.2f' % (name, tensor.dtype, np.min(tensor), np.mean(tensor), np.max(tensor)), shape)
+
+def print_stats_py(name, tensor):
+    shape = tensor.shape
+    print('%s (%s) min = %.2f, mean = %.2f, max = %.2f' % (name, tensor.dtype, np.min(tensor), np.mean(tensor), np.max(tensor)), shape)
+
+def print_(name, tensor):
+    tensor = tensor.detach().cpu().numpy()
+    print(name, tensor, tensor.shape)
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+def normalize_single(d):
+    # d is a whatever shape torch tensor
+    dmin = torch.min(d)
+    dmax = torch.max(d)
+    d = (d-dmin)/(EPS+(dmax-dmin))
+    return d
+
+def normalize(d):
+    # d is B x whatever. normalize within each element of the batch
+    out = torch.zeros(d.size(), dtype=d.dtype, device=d.device)
+    B = list(d.size())[0]
+    for b in list(range(B)):
+        out[b] = normalize_single(d[b])
+    return out
+
+def hard_argmax2d(tensor):
+    B, C, Y, X = list(tensor.shape)
+    assert(C==1)
+
+    # flatten the Tensor along the height and width axes
+    flat_tensor = tensor.reshape(B, -1)
+    # argmax of the flat tensor
+    argmax = torch.argmax(flat_tensor, dim=1)
+
+    # convert the indices into 2d coordinates
+    argmax_y = torch.floor(argmax / X) # row
+    argmax_x = argmax % X # col
+
+    argmax_y = argmax_y.reshape(B)
+    argmax_x = argmax_x.reshape(B)
+    return argmax_y, argmax_x
+
+def argmax2d(heat, hard=True):
+    B, C, Y, X = list(heat.shape)
+    assert(C==1)
+
+    if hard:
+        # hard argmax
+        loc_y, loc_x = hard_argmax2d(heat)
+        loc_y = loc_y.float()
+        loc_x = loc_x.float()
+    else:
+        heat = heat.reshape(B, Y*X)
+        prob = torch.nn.functional.softmax(heat, dim=1)
+
+        grid_y, grid_x = meshgrid2d(B, Y, X)
+
+        grid_y = grid_y.reshape(B, -1)
+        grid_x = grid_x.reshape(B, -1)
+        
+        loc_y = torch.sum(grid_y*prob, dim=1)
+        loc_x = torch.sum(grid_x*prob, dim=1)
+        # these are B
+        
+    return loc_y, loc_x
+
+def reduce_masked_mean(x, mask, dim=None, keepdim=False, broadcast=False):
+    # x and mask are the same shape, or at least broadcastably so < actually it's safer if you disallow broadcasting
+    # returns shape-1
+    # axis can be a list of axes
+    if not broadcast:
+        for (a,b) in zip(x.size(), mask.size()):
+            if not a==b:
+                print('some shape mismatch:', x.shape, mask.shape)
+            assert(a==b) # some shape mismatch!
+    # assert(x.size() == mask.size())
+    prod = x*mask
+    if dim is None:
+        numer = torch.sum(prod)
+        denom = EPS+torch.sum(mask)
+    else:
+        numer = torch.sum(prod, dim=dim, keepdim=keepdim)
+        denom = EPS+torch.sum(mask, dim=dim, keepdim=keepdim)
+    mean = numer/denom
+    return mean
+        
+def reduce_masked_median(x, mask, keep_batch=False):
+    # x and mask are the same shape
+    assert(x.size() == mask.size())
+    device = x.device
+
+    B = list(x.shape)[0]
+    x = x.detach().cpu().numpy()
+    mask = mask.detach().cpu().numpy()
+
+    if keep_batch:
+        x = np.reshape(x, [B, -1])
+        mask = np.reshape(mask, [B, -1])
+        meds = np.zeros([B], np.float32)
+        for b in list(range(B)):
+            xb = x[b]
+            mb = mask[b]
+            if np.sum(mb) > 0:
+                xb = xb[mb > 0]
+                meds[b] = np.median(xb)
+            else:
+                meds[b] = np.nan
+        meds = torch.from_numpy(meds).to(device)
+        return meds.float()
+    else:
+        x = np.reshape(x, [-1])
+        mask = np.reshape(mask, [-1])
+        if np.sum(mask) > 0:
+            x = x[mask > 0]
+            med = np.median(x)
+        else:
+            med = np.nan
+        med = np.array([med], np.float32)
+        med = torch.from_numpy(med).to(device)
+        return med.float()
+    
+def pack_seqdim(tensor, B):
+    shapelist = list(tensor.shape)
+    B_, S = shapelist[:2]
+    assert(B==B_)
+    otherdims = shapelist[2:]
+    tensor = torch.reshape(tensor, [B*S]+otherdims)
+    return tensor
+
+def unpack_seqdim(tensor, B):
+    shapelist = list(tensor.shape)
+    BS = shapelist[0]
+    assert(BS%B==0)
+    otherdims = shapelist[1:]
+    S = int(BS/B)
+    tensor = torch.reshape(tensor, [B,S]+otherdims)
+    return tensor
+
+def meshgrid2d(B, Y, X, stack=False, norm=False, device='cuda', on_chans=False):
+    # returns a meshgrid sized B x Y x X
+
+    grid_y = torch.linspace(0.0, Y-1, Y, device=torch.device(device))
+    grid_y = torch.reshape(grid_y, [1, Y, 1])
+    grid_y = grid_y.repeat(B, 1, X)
+
+    grid_x = torch.linspace(0.0, X-1, X, device=torch.device(device))
+    grid_x = torch.reshape(grid_x, [1, 1, X])
+    grid_x = grid_x.repeat(B, Y, 1)
+
+    if norm:
+        grid_y, grid_x = normalize_grid2d(
+            grid_y, grid_x, Y, X)
+
+    if stack:
+        # note we stack in xy order
+        # (see https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample)
+        if on_chans:
+            grid = torch.stack([grid_x, grid_y], dim=1)
+        else:
+            grid = torch.stack([grid_x, grid_y], dim=-1)
+        return grid
+    else:
+        return grid_y, grid_x
+
+def meshgrid2d_py(B, Y, X, stack=False, on_chans=False):
+    grid_y = np.linspace(0.0, Y-1, Y)
+    grid_y = np.reshape(grid_y, [1, Y, 1])
+    grid_y = np.tile(grid_y, [B, 1, X])
+
+    grid_x = np.linspace(0.0, X-1, X)
+    grid_x = np.reshape(grid_x, [1, 1, X])
+    grid_x = np.tile(grid_x, [B, Y, 1])
+
+    if stack:
+        # note we stack in xy order
+        # (see https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample)
+        if on_chans:
+            grid = np.stack([grid_x, grid_y], axis=1)
+        else:
+            grid = np.stack([grid_x, grid_y], axis=-1)
+        return grid
+    else:
+        # outputs are Y x X
+        return grid_y, grid_x
+    
+
+def meshgrid3d(B, Z, Y, X, stack=False, norm=False, device='cuda'):
+    # returns a meshgrid sized B x Z x Y x X
+    
+    grid_z = torch.linspace(0.0, Z-1, Z, device=device)
+    grid_z = torch.reshape(grid_z, [1, Z, 1, 1])
+    grid_z = grid_z.repeat(B, 1, Y, X)
+
+    grid_y = torch.linspace(0.0, Y-1, Y, device=device)
+    grid_y = torch.reshape(grid_y, [1, 1, Y, 1])
+    grid_y = grid_y.repeat(B, Z, 1, X)
+
+    grid_x = torch.linspace(0.0, X-1, X, device=device)
+    grid_x = torch.reshape(grid_x, [1, 1, 1, X])
+    grid_x = grid_x.repeat(B, Z, Y, 1)
+
+    if norm:
+        grid_z, grid_y, grid_x = normalize_grid3d(
+            grid_z, grid_y, grid_x, Z, Y, X)
+
+    if stack:
+        # note we stack in xyz order
+        # (see https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample)
+        grid = torch.stack([grid_x, grid_y, grid_z], dim=-1)
+        return grid
+    else:
+        return grid_z, grid_y, grid_x
+    
+def normalize_grid2d(grid_y, grid_x, Y, X, clamp_extreme=True):
+    # make things in [-1,1]
+    grid_y = 2.0*(grid_y / float(Y-1)) - 1.0
+    grid_x = 2.0*(grid_x / float(X-1)) - 1.0
+    
+    if clamp_extreme:
+        grid_y = torch.clamp(grid_y, min=-2.0, max=2.0)
+        grid_x = torch.clamp(grid_x, min=-2.0, max=2.0)
+        
+    return grid_y, grid_x
+
+def normalize_grid3d(grid_z, grid_y, grid_x, Z, Y, X, clamp_extreme=True):
+    # make things in [-1,1]
+    grid_z = 2.0*(grid_z / float(Z-1)) - 1.0
+    grid_y = 2.0*(grid_y / float(Y-1)) - 1.0
+    grid_x = 2.0*(grid_x / float(X-1)) - 1.0
+    
+    if clamp_extreme:
+        grid_z = torch.clamp(grid_z, min=-2.0, max=2.0)
+        grid_y = torch.clamp(grid_y, min=-2.0, max=2.0)
+        grid_x = torch.clamp(grid_x, min=-2.0, max=2.0)
+    
+    return grid_z, grid_y, grid_x
+
+def gridcloud2d(B, Y, X, norm=False, device='cuda'):
+    # we want to sample for each location in the grid
+    grid_y, grid_x = meshgrid2d(B, Y, X, norm=norm, device=device)
+    x = torch.reshape(grid_x, [B, -1])
+    y = torch.reshape(grid_y, [B, -1])
+    # these are B x N
+    xy = torch.stack([x, y], dim=2)
+    # this is B x N x 2
+    return xy
+
+def gridcloud3d(B, Z, Y, X, norm=False, device='cuda'):
+    # we want to sample for each location in the grid
+    grid_z, grid_y, grid_x = meshgrid3d(B, Z, Y, X, norm=norm, device=device)
+    x = torch.reshape(grid_x, [B, -1])
+    y = torch.reshape(grid_y, [B, -1])
+    z = torch.reshape(grid_z, [B, -1])
+    # these are B x N
+    xyz = torch.stack([x, y, z], dim=2)
+    # this is B x N x 3
+    return xyz
+
+# import re
+# def readPFM(file):
+#     file = open(file, 'rb')
+
+#     color = None
+#     width = None
+#     height = None
+#     scale = None
+#     endian = None
+
+#     header = file.readline().rstrip()
+#     if header == b'PF':
+#         color = True
+#     elif header == b'Pf':
+#         color = False
+#     else:
+#         raise Exception('Not a PFM file.')
+
+#     dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline())
+#     if dim_match:
+#         width, height = map(int, dim_match.groups())
+#     else:
+#         raise Exception('Malformed PFM header.')
+
+#     scale = float(file.readline().rstrip())
+#     if scale < 0: # little-endian
+#         endian = '<'
+#         scale = -scale
+#     else:
+#         endian = '>' # big-endian
+
+#     data = np.fromfile(file, endian + 'f')
+#     shape = (height, width, 3) if color else (height, width)
+
+#     data = np.reshape(data, shape)
+#     data = np.flipud(data)
+#     return data    
+
+def normalize_boxlist2d(boxlist2d, H, W):
+    boxlist2d = boxlist2d.clone()
+    ymin, xmin, ymax, xmax = torch.unbind(boxlist2d, dim=2)
+    ymin = ymin / float(H)
+    ymax = ymax / float(H)
+    xmin = xmin / float(W)
+    xmax = xmax / float(W)
+    boxlist2d = torch.stack([ymin, xmin, ymax, xmax], dim=2)
+    return boxlist2d
+
+def unnormalize_boxlist2d(boxlist2d, H, W):
+    boxlist2d = boxlist2d.clone()
+    ymin, xmin, ymax, xmax = torch.unbind(boxlist2d, dim=2)
+    ymin = ymin * float(H)
+    ymax = ymax * float(H)
+    xmin = xmin * float(W)
+    xmax = xmax * float(W)
+    boxlist2d = torch.stack([ymin, xmin, ymax, xmax], dim=2)
+    return boxlist2d
+
+def unnormalize_box2d(box2d, H, W):
+    return unnormalize_boxlist2d(box2d.unsqueeze(1), H, W).squeeze(1)
+
+def normalize_box2d(box2d, H, W):
+    return normalize_boxlist2d(box2d.unsqueeze(1), H, W).squeeze(1)
+
+def get_gaussian_kernel_2d(channels, kernel_size=3, sigma=2.0, mid_one=False, device='cuda'):
+    C = channels
+    xy_grid = gridcloud2d(C, kernel_size, kernel_size, device=device) # C x N x 2
+
+    mean = (kernel_size - 1)/2.0
+    variance = sigma**2.0
+
+    gaussian_kernel = (1.0/(2.0*np.pi*variance)**1.5) * torch.exp(-torch.sum((xy_grid - mean)**2.0, dim=-1) / (2.0*variance)) # C X N
+    gaussian_kernel = gaussian_kernel.view(C, 1, kernel_size, kernel_size) # C x 1 x 3 x 3
+    kernel_sum = torch.sum(gaussian_kernel, dim=(2,3), keepdim=True)
+
+    gaussian_kernel = gaussian_kernel / kernel_sum # normalize
+
+    if mid_one:
+        # normalize so that the middle element is 1
+        maxval = gaussian_kernel[:,:,(kernel_size//2),(kernel_size//2)].reshape(C, 1, 1, 1)
+        gaussian_kernel = gaussian_kernel / maxval
+
+    return gaussian_kernel
+
+def gaussian_blur_2d(input, kernel_size=3, sigma=2.0, reflect_pad=False, mid_one=False):
+    B, C, Z, X = input.shape
+    kernel = get_gaussian_kernel_2d(C, kernel_size, sigma, mid_one=mid_one, device=input.device)
+    if reflect_pad:
+        pad = (kernel_size - 1)//2
+        out = F.pad(input, (pad, pad, pad, pad), mode='reflect')
+        out = F.conv2d(out, kernel, padding=0, groups=C)
+    else:
+        out = F.conv2d(input, kernel, padding=(kernel_size - 1)//2, groups=C)
+    return out
+
+def gradient2d(x, absolute=False, square=False, return_sum=False):
+    # x should be B x C x H x W
+    dh = x[:, :, 1:, :] - x[:, :, :-1, :]
+    dw = x[:, :, :, 1:] - x[:, :, :, :-1]
+
+    zeros = torch.zeros_like(x)
+    zero_h = zeros[:, :, 0:1, :]
+    zero_w = zeros[:, :, :, 0:1]
+    dh = torch.cat([dh, zero_h], axis=2)
+    dw = torch.cat([dw, zero_w], axis=3)
+    if absolute:
+        dh = torch.abs(dh)
+        dw = torch.abs(dw)
+    if square:
+        dh = dh ** 2
+        dw = dw ** 2
+    if return_sum:
+        return dh+dw
+    else: 
+        return dh, dw
+
+def gradient1d(x, absolute=False, square=False):
+    # x should be B,S,C
+    dx = x[:, 1:] - x[:, :-1]
+    zero = torch.zeros_like(x[:,0:1])
+    dx = torch.cat([dx, zero], axis=1)
+    if absolute:
+        dx = torch.abs(dx)
+    if square:
+        dx = dx ** 2
+    return dx
+    
+def smart_cat(tensor1, tensor2, dim):
+    if tensor1 is None:
+        return tensor2
+    return torch.cat([tensor1, tensor2], dim=dim)

utils/data.py

@@ -0,0 +1,122 @@
+import torch
+import dataclasses
+import torch.nn.functional as F
+from dataclasses import dataclass
+from typing import Any, Optional, Dict
+
+
+@dataclass(eq=False)
+class VideoData:
+    """
+    Dataclass for storing video tracks data.
+    """
+
+    video: torch.Tensor  # B, S, C, H, W
+    trajs: torch.Tensor  # B, S, N, 2
+    visibs: torch.Tensor  # B, S, N
+    # optional data
+    valids: Optional[torch.Tensor] = None  # B, S, N
+    hards: Optional[torch.Tensor] = None  # B, S, N
+    segmentation: Optional[torch.Tensor] = None  # B, S, 1, H, W
+    seq_name: Optional[str] = None
+    dname: Optional[str] = None
+    query_points: Optional[torch.Tensor] = None  # TapVID evaluation format
+    transforms: Optional[Dict[str, Any]] = None
+    aug_video: Optional[torch.Tensor] = None
+
+
+def collate_fn(batch):
+    """
+    Collate function for video tracks data.
+    """
+    video = torch.stack([b.video for b in batch], dim=0)
+    trajs = torch.stack([b.trajs for b in batch], dim=0)
+    visibs = torch.stack([b.visibs for b in batch], dim=0)
+    query_points = segmentation = None
+    if batch[0].query_points is not None:
+        query_points = torch.stack([b.query_points for b in batch], dim=0)
+    if batch[0].segmentation is not None:
+        segmentation = torch.stack([b.segmentation for b in batch], dim=0)
+    seq_name = [b.seq_name for b in batch]
+    dname = [b.dname for b in batch]
+
+    return VideoData(
+        video=video,
+        trajs=trajs,
+        visibs=visibs,
+        segmentation=segmentation,
+        seq_name=seq_name,
+        dname=dname,
+        query_points=query_points,
+    )
+
+
+def collate_fn_train(batch):
+    """
+    Collate function for video tracks data during training.
+    """
+    gotit = [gotit for _, gotit in batch]
+    video = torch.stack([b.video for b, _ in batch], dim=0)
+    trajs = torch.stack([b.trajs for b, _ in batch], dim=0)
+    visibs = torch.stack([b.visibs for b, _ in batch], dim=0)
+    valids = torch.stack([b.valids for b, _ in batch], dim=0)
+    seq_name = [b.seq_name for b, _ in batch]
+    dname = [b.dname for b, _ in batch]
+    query_points = transforms = aug_video = hards = None
+    if batch[0][0].query_points is not None:
+        query_points = torch.stack([b.query_points for b, _ in batch], dim=0)
+    if batch[0][0].hards is not None:
+        hards = torch.stack([b.hards for b, _ in batch], dim=0)
+
+    if batch[0][0].transforms is not None:
+        transforms = [b.transforms for b, _ in batch]
+
+    if batch[0][0].aug_video is not None:
+        aug_video = torch.stack([b.aug_video for b, _ in batch], dim=0)
+    return (
+        VideoData(
+            video=video,
+            trajs=trajs,
+            visibs=visibs,
+            valids=valids,
+            hards=hards,
+            seq_name=seq_name,
+            dname=dname,
+            query_points=query_points,
+            aug_video=aug_video,
+            transforms=transforms,
+        ),
+        gotit,
+    )
+
+
+def try_to_cuda(t: Any) -> Any:
+    """
+    Try to move the input variable `t` to a cuda device.
+
+    Args:
+        t: Input.
+
+    Returns:
+        t_cuda: `t` moved to a cuda device, if supported.
+    """
+    try:
+        t = t.float().cuda()
+    except AttributeError:
+        pass
+    return t
+
+
+def dataclass_to_cuda_(obj):
+    """
+    Move all contents of a dataclass to cuda inplace if supported.
+
+    Args:
+        batch: Input dataclass.
+
+    Returns:
+        batch_cuda: `batch` moved to a cuda device, if supported.
+    """
+    for f in dataclasses.fields(obj):
+        setattr(obj, f.name, try_to_cuda(getattr(obj, f.name)))
+    return obj

utils/geom.py

@@ -0,0 +1,771 @@
+import torch
+import utils.basic
+# import utils.box
+# import utils.vox
+import numpy as np
+import torchvision.ops as ops
+from utils.basic import print_
+
+def split_intrinsics(K):
+    # K is B x 3 x 3 or B x 4 x 4
+    fx = K[:,0,0]
+    fy = K[:,1,1]
+    x0 = K[:,0,2]
+    y0 = K[:,1,2]
+    return fx, fy, x0, y0
+
+# def apply_pix_T_cam_py(pix_T_cam, xyz):
+
+#     fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
+    
+#     # xyz is shaped B x H*W x 3
+#     # returns xy, shaped B x H*W x 2
+    
+#     B, N, C = list(xyz.shape)
+#     assert(C==3)
+    
+#     x, y, z = xyz[:,:,0], xyz[:,:,1], xyz[:,:,2]
+
+#     fx = np.reshape(fx, [B, 1])
+#     fy = np.reshape(fy, [B, 1])
+#     x0 = np.reshape(x0, [B, 1])
+#     y0 = np.reshape(y0, [B, 1])
+
+#     EPS = 1e-4
+#     z = np.clip(z, EPS, None)
+#     x = (x*fx)/(z)+x0
+#     y = (y*fy)/(z)+y0
+#     xy = np.stack([x, y], axis=-1)
+#     return xy
+
+def apply_pix_T_cam(pix_T_cam, xyz):
+
+    fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
+    
+    # xyz is shaped B x H*W x 3
+    # returns xy, shaped B x H*W x 2
+    
+    B, N, C = list(xyz.shape)
+    assert(C==3)
+    
+    x, y, z = torch.unbind(xyz, axis=-1)
+
+    fx = torch.reshape(fx, [B, 1])
+    fy = torch.reshape(fy, [B, 1])
+    x0 = torch.reshape(x0, [B, 1])
+    y0 = torch.reshape(y0, [B, 1])
+
+    EPS = 1e-4
+    z = torch.clamp(z, min=EPS)
+    x = (x*fx)/(z)+x0
+    y = (y*fy)/(z)+y0
+    xy = torch.stack([x, y], axis=-1)
+    return xy
+
+# def apply_4x4_py(RT, xyz):
+#     # print('RT', RT.shape)
+#     B, N, _ = list(xyz.shape)
+#     ones = np.ones_like(xyz[:,:,0:1])
+#     xyz1 = np.concatenate([xyz, ones], 2)
+#     # print('xyz1', xyz1.shape)
+#     xyz1_t = xyz1.transpose(0,2,1)
+#     # print('xyz1_t', xyz1_t.shape)
+#     # this is B x 4 x N
+#     xyz2_t = np.matmul(RT, xyz1_t)
+#     # print('xyz2_t', xyz2_t.shape)
+#     xyz2 = xyz2_t.transpose(0,2,1)
+#     # print('xyz2', xyz2.shape)
+#     xyz2 = xyz2[:,:,:3]
+#     return xyz2
+
+def apply_4x4(RT, xyz):
+    B, N, _ = list(xyz.shape)
+    ones = torch.ones_like(xyz[:,:,0:1])
+    xyz1 = torch.cat([xyz, ones], 2)
+    xyz1_t = torch.transpose(xyz1, 1, 2)
+    # this is B x 4 x N
+    xyz2_t = torch.matmul(RT, xyz1_t)
+    xyz2 = torch.transpose(xyz2_t, 1, 2)
+    xyz2 = xyz2[:,:,:3]
+    return xyz2
+
+def apply_3x3(RT, xy):
+    B, N, _ = list(xy.shape)
+    ones = torch.ones_like(xy[:,:,0:1])
+    xy1 = torch.cat([xy, ones], 2)
+    xy1_t = torch.transpose(xy1, 1, 2)
+    # this is B x 4 x N
+    xy2_t = torch.matmul(RT, xy1_t)
+    xy2 = torch.transpose(xy2_t, 1, 2)
+    xy2 = xy2[:,:,:2]
+    return xy2
+
+def generate_polygon(ctr_x, ctr_y, avg_r, irregularity, spikiness, num_verts):
+    '''
+    Start with the center of the polygon at ctr_x, ctr_y, 
+    Then creates the polygon by sampling points on a circle around the center.
+    Random noise is added by varying the angular spacing between sequential points,
+    and by varying the radial distance of each point from the centre.
+
+    Params:
+        ctr_x, ctr_y - coordinates of the "centre" of the polygon
+        avg_r - in px, the average radius of this polygon, this roughly controls how large the polygon is, really only useful for order of magnitude.
+        irregularity - [0,1] indicating how much variance there is in the angular spacing of vertices. [0,1] will map to [0, 2pi/numberOfVerts]
+        spikiness - [0,1] indicating how much variance there is in each vertex from the circle of radius avg_r. [0,1] will map to [0, avg_r]
+        num_verts
+
+    Returns:
+        np.array [num_verts, 2] - CCW order.
+    '''
+    # spikiness
+    spikiness = np.clip(spikiness, 0, 1) * avg_r
+
+    # generate n angle steps
+    irregularity = np.clip(irregularity, 0, 1) * 2 * np.pi / num_verts
+    lower = (2*np.pi / num_verts) - irregularity
+    upper = (2*np.pi / num_verts) + irregularity
+
+    # angle steps
+    angle_steps = np.random.uniform(lower, upper, num_verts)
+    sc = (2 * np.pi) / angle_steps.sum()
+    angle_steps *= sc
+
+    # get all radii
+    angle = np.random.uniform(0, 2*np.pi)
+    radii = np.clip(np.random.normal(avg_r, spikiness, num_verts), 0, 2 * avg_r)
+
+    # compute all points
+    points = []
+    for i in range(num_verts):
+        x = ctr_x + radii[i] * np.cos(angle)
+        y = ctr_y + radii[i] * np.sin(angle)
+        points.append([x, y])
+        angle += angle_steps[i]
+
+    return np.array(points).astype(int)
+
+
+def get_random_affine_2d(B, rot_min=-5.0, rot_max=5.0, tx_min=-0.1, tx_max=0.1, ty_min=-0.1, ty_max=0.1, sx_min=-0.05, sx_max=0.05, sy_min=-0.05, sy_max=0.05, shx_min=-0.05, shx_max=0.05, shy_min=-0.05, shy_max=0.05):
+    '''
+    Params:
+        rot_min: rotation amount min
+        rot_max: rotation amount max
+
+        tx_min: translation x min
+        tx_max: translation x max
+
+        ty_min: translation y min
+        ty_max: translation y max
+
+        sx_min: scaling x min
+        sx_max: scaling x max
+
+        sy_min: scaling y min
+        sy_max: scaling y max
+
+        shx_min: shear x min
+        shx_max: shear x max
+
+        shy_min: shear y min
+        shy_max: shear y max
+
+    Returns:
+        transformation matrix: (B, 3, 3)
+    '''
+    # rotation
+    if rot_max - rot_min != 0:
+        rot_amount = np.random.uniform(low=rot_min, high=rot_max, size=B)
+        rot_amount = np.pi/180.0*rot_amount
+    else:
+        rot_amount = rot_min
+    rotation = np.zeros((B, 3, 3)) # B, 3, 3
+    rotation[:, 2, 2] = 1
+    rotation[:, 0, 0] = np.cos(rot_amount)
+    rotation[:, 0, 1] = -np.sin(rot_amount)
+    rotation[:, 1, 0] = np.sin(rot_amount)
+    rotation[:, 1, 1] = np.cos(rot_amount)
+
+    # translation
+    translation = np.zeros((B, 3, 3)) # B, 3, 3
+    translation[:, [0,1,2], [0,1,2]] = 1 
+    if (tx_max - tx_min) > 0:
+        trans_x = np.random.uniform(low=tx_min, high=tx_max, size=B)
+        translation[:, 0, 2] = trans_x
+    # else:
+    #     translation[:, 0, 2] = tx_max
+    if ty_max - ty_min != 0:
+        trans_y = np.random.uniform(low=ty_min, high=ty_max, size=B)
+        translation[:, 1, 2] = trans_y
+    # else:
+    #     translation[:, 1, 2] = ty_max
+
+    # scaling
+    scaling = np.zeros((B, 3, 3)) # B, 3, 3
+    scaling[:, [0,1,2], [0,1,2]] = 1 
+    if (sx_max - sx_min) > 0:
+        scale_x = 1 + np.random.uniform(low=sx_min, high=sx_max, size=B)
+        scaling[:, 0, 0] = scale_x
+    # else:
+    #     scaling[:, 0, 0] = sx_max
+    if (sy_max - sy_min) > 0:
+        scale_y = 1 + np.random.uniform(low=sy_min, high=sy_max, size=B)
+        scaling[:, 1, 1] = scale_y
+    # else:
+    #     scaling[:, 1, 1] = sy_max
+
+    # shear
+    shear = np.zeros((B, 3, 3)) # B, 3, 3
+    shear[:, [0,1,2], [0,1,2]] = 1 
+    if (shx_max - shx_min) > 0:
+        shear_x = np.random.uniform(low=shx_min, high=shx_max, size=B)
+        shear[:, 0, 1] = shear_x
+    # else:
+    #     shear[:, 0, 1] = shx_max
+    if (shy_max - shy_min) > 0:
+        shear_y = np.random.uniform(low=shy_min, high=shy_max, size=B)
+        shear[:, 1, 0] = shear_y
+    # else:
+    #     shear[:, 1, 0] = shy_max
+
+    # compose all those
+    rt = np.einsum("ijk,ikl->ijl", rotation, translation)
+    ss = np.einsum("ijk,ikl->ijl", scaling, shear)
+    trans = np.einsum("ijk,ikl->ijl", rt, ss)
+
+    return trans
+
+def get_centroid_from_box2d(box2d):
+    ymin = box2d[:,0]
+    xmin = box2d[:,1]
+    ymax = box2d[:,2]
+    xmax = box2d[:,3]
+    x = (xmin+xmax)/2.0
+    y = (ymin+ymax)/2.0
+    return y, x
+
+def normalize_boxlist2d(boxlist2d, H, W):
+    boxlist2d = boxlist2d.clone()
+    ymin, xmin, ymax, xmax = torch.unbind(boxlist2d, dim=2)
+    ymin = ymin / float(H)
+    ymax = ymax / float(H)
+    xmin = xmin / float(W)
+    xmax = xmax / float(W)
+    boxlist2d = torch.stack([ymin, xmin, ymax, xmax], dim=2)
+    return boxlist2d
+
+def unnormalize_boxlist2d(boxlist2d, H, W):
+    boxlist2d = boxlist2d.clone()
+    ymin, xmin, ymax, xmax = torch.unbind(boxlist2d, dim=2)
+    ymin = ymin * float(H)
+    ymax = ymax * float(H)
+    xmin = xmin * float(W)
+    xmax = xmax * float(W)
+    boxlist2d = torch.stack([ymin, xmin, ymax, xmax], dim=2)
+    return boxlist2d
+
+def unnormalize_box2d(box2d, H, W):
+    return unnormalize_boxlist2d(box2d.unsqueeze(1), H, W).squeeze(1)
+
+def normalize_box2d(box2d, H, W):
+    return normalize_boxlist2d(box2d.unsqueeze(1), H, W).squeeze(1)
+
+def get_size_from_box2d(box2d):
+    ymin = box2d[:,0]
+    xmin = box2d[:,1]
+    ymax = box2d[:,2]
+    xmax = box2d[:,3]
+    height = ymax-ymin
+    width = xmax-xmin
+    return height, width
+
+def crop_and_resize(im, boxlist, PH, PW, boxlist_is_normalized=False):
+    B, C, H, W = im.shape
+    B2, N, D = boxlist.shape
+    assert(B==B2)
+    assert(D==4)
+    # PH, PW is the size to resize to
+
+    # print('im', im.shape)
+
+    # output is B,N,C,PH,PW
+
+    # pt wants xy xy, unnormalized
+    if boxlist_is_normalized:
+        boxlist_unnorm = unnormalize_boxlist2d(boxlist, H, W)
+    else:
+        boxlist_unnorm = boxlist
+        
+    ymin, xmin, ymax, xmax = boxlist_unnorm.unbind(2)
+    # boxlist_pt = torch.stack([boxlist_unnorm[:,1], boxlist_unnorm[:,0], boxlist_unnorm[:,3], boxlist_unnorm[:,2]], dim=1)
+    boxlist_pt = torch.stack([xmin, ymin, xmax, ymax], dim=2)
+    # we want a B-len list of K x 4 arrays
+
+    # print('im', im.shape)
+    # print('boxlist', boxlist.shape)
+    # print('boxlist_pt', boxlist_pt.shape)
+
+    # boxlist_pt = list(boxlist_pt.unbind(0))
+
+    # crops = []
+    # for b in range(B):
+    #     crops_b = ops.roi_align(im[b:b+1], [boxlist_pt[b]], output_size=(PH, PW))
+    #     crops.append(crops_b)
+    # # # crops = im
+
+    # inds = torch.arange(B).reshape(B,1)
+    # # boxlist_pt = torch.cat([inds, boxlist_
+    # boxlist_pt = list(boxlist_pt.unbind(0)
+
+    # crops = ops.roi_align(im, [boxlist_pt[b]], output_size=(PH, PW))
+    crops = ops.roi_align(im, list(boxlist_pt.unbind(0)), output_size=(PH, PW))
+    
+
+    # print('crops', crops.shape)
+    # crops = crops.reshape(B,N,C,PH,PW)
+
+    
+    # crops = []
+    # for b in range(B):
+    #     crop_b = ops.roi_align(im[b:b+1], [boxlist_pt[b]], output_size=(PH, PW))
+    #     print('crop_b', crop_b.shape)
+    #     crops.append(crop_b)
+    # crops = torch.stack(crops, dim=0)
+        
+    # print('crops', crops.shape)
+    # boxlist_list = boxlist_pt.unbind(0)
+    # print('rgb_crop', rgb_crop.shape)
+
+    return crops
+
+
+# def get_boxlist_from_centroid_and_size(cy, cx, h, w, clip=True):
+#     # cy,cx are both B,N
+#     ymin = cy - h/2
+#     ymax = cy + h/2
+#     xmin = cx - w/2
+#     xmax = cx + w/2
+
+#     box = torch.stack([ymin, xmin, ymax, xmax], dim=-1)
+#     if clip:
+#         box = torch.clamp(box, 0, 1)
+#     return box
+
+
+def get_boxlist_from_centroid_and_size(cy, cx, h, w):#, clip=False):
+    # cy,cx are the same shape
+    ymin = cy - h/2
+    ymax = cy + h/2
+    xmin = cx - w/2
+    xmax = cx + w/2
+
+    # if clip:
+    #     ymin = torch.clamp(ymin, 0, H-1)
+    #     ymax = torch.clamp(ymax, 0, H-1)
+    #     xmin = torch.clamp(xmin, 0, W-1)
+    #     xmax = torch.clamp(xmax, 0, W-1)
+    
+    box = torch.stack([ymin, xmin, ymax, xmax], dim=-1)
+    return box
+
+
+def get_box2d_from_mask(mask, normalize=False):
+    # mask is B, 1, H, W
+
+    B, C, H, W = mask.shape
+    assert(C==1)
+    xy = utils.basic.gridcloud2d(B, H, W, norm=False, device=mask.device) # B, H*W, 2
+
+    box = torch.zeros((B, 4), dtype=torch.float32, device=mask.device)
+    for b in range(B):
+        xy_b = xy[b] # H*W, 2
+        mask_b = mask[b].reshape(H*W)
+        xy_ = xy_b[mask_b > 0]
+        x_ = xy_[:,0]
+        y_ = xy_[:,1]
+        ymin = torch.min(y_)
+        ymax = torch.max(y_)
+        xmin = torch.min(x_)
+        xmax = torch.max(x_)
+        box[b] = torch.stack([ymin, xmin, ymax, xmax], dim=0)
+    if normalize:
+        box = normalize_boxlist2d(box.unsqueeze(1), H, W).squeeze(1)
+    return box
+
+def convert_box2d_to_intrinsics(box2d, pix_T_cam, H, W, use_image_aspect_ratio=True, mult_padding=1.0):
+    # box2d is B x 4, with ymin, xmin, ymax, xmax in normalized coords
+    # ymin, xmin, ymax, xmax = torch.unbind(box2d, dim=1)
+    # H, W is the original size of the image
+    # mult_padding is relative to object size in pixels
+
+    # i assume we're rendering an image the same size as the original (H, W)
+
+    if not mult_padding==1.0:
+        y, x = get_centroid_from_box2d(box2d)
+        h, w = get_size_from_box2d(box2d)
+        box2d = get_box2d_from_centroid_and_size(
+            y, x, h*mult_padding, w*mult_padding, clip=False)
+        
+    if use_image_aspect_ratio:
+        h, w = get_size_from_box2d(box2d)
+        y, x = get_centroid_from_box2d(box2d)
+
+        # note h,w are relative right now
+        # we need to undo this, to see the real ratio
+
+        h = h*float(H)
+        w = w*float(W)
+        box_ratio = h/w
+        im_ratio = H/float(W)
+
+        # print('box_ratio:', box_ratio)
+        # print('im_ratio:', im_ratio)
+
+        if box_ratio >= im_ratio:
+            w = h/im_ratio
+            # print('setting w:', h/im_ratio)
+        else:
+            h = w*im_ratio
+            # print('setting h:', w*im_ratio)
+            
+        box2d = get_box2d_from_centroid_and_size(
+            y, x, h/float(H), w/float(W), clip=False)
+
+    assert(h > 1e-4)
+    assert(w > 1e-4)
+        
+    ymin, xmin, ymax, xmax = torch.unbind(box2d, dim=1)
+
+    fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
+
+    # the topleft of the new image will now have a different offset from the center of projection
+    
+    new_x0 = x0 - xmin*W
+    new_y0 = y0 - ymin*H
+
+    pix_T_cam = pack_intrinsics(fx, fy, new_x0, new_y0)
+    # this alone will give me an image in original resolution,
+    # with its topleft at the box corner
+
+    box_h, box_w = get_size_from_box2d(box2d)
+    # these are normalized, and shaped B. (e.g., [0.4], [0.3])
+
+    # we are going to scale the image by the inverse of this,
+    # since we are zooming into this area
+
+    sy = 1./box_h
+    sx = 1./box_w
+
+    pix_T_cam = scale_intrinsics(pix_T_cam, sx, sy)
+    return pix_T_cam, box2d
+
+def generatePolygon(ctr_x, ctr_y, avg_r, irregularity, spikiness, num_verts):
+    '''
+    Start with the center of the polygon at ctr_x, ctr_y, 
+    Then creates the polygon by sampling points on a circle around the center.
+    Random noise is added by varying the angular spacing between sequential points,
+    and by varying the radial distance of each point from the centre.
+
+    Params:
+        ctr_x, ctr_y - coordinates of the "centre" of the polygon
+        avg_r - in px, the average radius of this polygon, this roughly controls how large the polygon is, really only useful for order of magnitude.
+        irregularity - [0,1] indicating how much variance there is in the angular spacing of vertices. [0,1] will map to [0, 2pi/numberOfVerts]
+        spikiness - [0,1] indicating how much variance there is in each vertex from the circle of radius avg_r. [0,1] will map to [0, avg_r]
+        num_verts
+
+    Returns:
+        np.array [num_verts, 2] - CCW order.
+    '''
+    # spikiness
+    spikiness = np.clip(spikiness, 0, 1) * avg_r
+
+    # generate n angle steps
+    irregularity = np.clip(irregularity, 0, 1) * 2 * np.pi / num_verts
+    lower = (2*np.pi / num_verts) - irregularity
+    upper = (2*np.pi / num_verts) + irregularity
+
+    # angle steps
+    angle_steps = np.random.uniform(lower, upper, num_verts)
+    sc = (2 * np.pi) / angle_steps.sum()
+    angle_steps *= sc
+
+    # get all radii
+    angle = np.random.uniform(0, 2*np.pi)
+    radii = np.clip(np.random.normal(avg_r, spikiness, num_verts), 0, 2 * avg_r)
+
+    # compute all points
+    points = []
+    for i in range(num_verts):
+        x = ctr_x + radii[i] * np.cos(angle)
+        y = ctr_y + radii[i] * np.sin(angle)
+        points.append([x, y])
+        angle += angle_steps[i]
+
+    return np.array(points).astype(int)
+
+def get_random_affine_2d(B, rot_min=-5.0, rot_max=5.0, tx_min=-0.1, tx_max=0.1, ty_min=-0.1, ty_max=0.1, sx_min=-0.05, sx_max=0.05, sy_min=-0.05, sy_max=0.05, shx_min=-0.05, shx_max=0.05, shy_min=-0.05, shy_max=0.05):
+    '''
+    Params:
+        rot_min: rotation amount min
+        rot_max: rotation amount max
+
+        tx_min: translation x min
+        tx_max: translation x max
+
+        ty_min: translation y min
+        ty_max: translation y max
+
+        sx_min: scaling x min
+        sx_max: scaling x max
+
+        sy_min: scaling y min
+        sy_max: scaling y max
+
+        shx_min: shear x min
+        shx_max: shear x max
+
+        shy_min: shear y min
+        shy_max: shear y max
+
+    Returns:
+        transformation matrix: (B, 3, 3)
+    '''
+    # rotation
+    if rot_max - rot_min != 0:
+        rot_amount = np.random.uniform(low=rot_min, high=rot_max, size=B)
+        rot_amount = np.pi/180.0*rot_amount
+    else:
+        rot_amount = rot_min
+    rotation = np.zeros((B, 3, 3)) # B, 3, 3
+    rotation[:, 2, 2] = 1
+    rotation[:, 0, 0] = np.cos(rot_amount)
+    rotation[:, 0, 1] = -np.sin(rot_amount)
+    rotation[:, 1, 0] = np.sin(rot_amount)
+    rotation[:, 1, 1] = np.cos(rot_amount)
+
+    # translation
+    translation = np.zeros((B, 3, 3)) # B, 3, 3
+    translation[:, [0,1,2], [0,1,2]] = 1 
+    if tx_max - tx_min != 0:
+        trans_x = np.random.uniform(low=tx_min, high=tx_max, size=B)
+        translation[:, 0, 2] = trans_x
+    else:
+        translation[:, 0, 2] = tx_max
+    if ty_max - ty_min != 0:
+        trans_y = np.random.uniform(low=ty_min, high=ty_max, size=B)
+        translation[:, 1, 2] = trans_y
+    else:
+        translation[:, 1, 2] = ty_max
+
+    # scaling
+    scaling = np.zeros((B, 3, 3)) # B, 3, 3
+    scaling[:, [0,1,2], [0,1,2]] = 1 
+    if sx_max - sx_min != 0:
+        scale_x = 1 + np.random.uniform(low=sx_min, high=sx_max, size=B)
+        scaling[:, 0, 0] = scale_x
+    else:
+        scaling[:, 0, 0] = sx_max
+    if sy_max - sy_min != 0:
+        scale_y = 1 + np.random.uniform(low=sy_min, high=sy_max, size=B)
+        scaling[:, 1, 1] = scale_y
+    else:
+        scaling[:, 1, 1] = sy_max
+
+    # shear
+    shear = np.zeros((B, 3, 3)) # B, 3, 3
+    shear[:, [0,1,2], [0,1,2]] = 1 
+    if shx_max - shx_min != 0:
+        shear_x = np.random.uniform(low=shx_min, high=shx_max, size=B)
+        shear[:, 0, 1] = shear_x
+    else:
+        shear[:, 0, 1] = shx_max
+    if shy_max - shy_min != 0:
+        shear_y = np.random.uniform(low=shy_min, high=shy_max, size=B)
+        shear[:, 1, 0] = shear_y
+    else:
+        shear[:, 1, 0] = shy_max
+
+    # compose all those
+    rt = np.einsum("ijk,ikl->ijl", rotation, translation)
+    ss = np.einsum("ijk,ikl->ijl", scaling, shear)
+    trans = np.einsum("ijk,ikl->ijl", rt, ss)
+
+    return trans
+
+def pixels2camera(x,y,z,fx,fy,x0,y0):
+    # x and y are locations in pixel coordinates, z is a depth in meters
+    # they can be images or pointclouds
+    # fx, fy, x0, y0 are camera intrinsics
+    # returns xyz, sized B x N x 3
+
+    B = x.shape[0]
+    
+    fx = torch.reshape(fx, [B,1])
+    fy = torch.reshape(fy, [B,1])
+    x0 = torch.reshape(x0, [B,1])
+    y0 = torch.reshape(y0, [B,1])
+
+    x = torch.reshape(x, [B,-1])
+    y = torch.reshape(y, [B,-1])
+    z = torch.reshape(z, [B,-1])
+    
+    # unproject
+    x = (z/fx)*(x-x0)
+    y = (z/fy)*(y-y0)
+    
+    xyz = torch.stack([x,y,z], dim=2)
+    # B x N x 3
+    return xyz
+
+def camera2pixels(xyz, pix_T_cam):
+    # xyz is shaped B x H*W x 3
+    # returns xy, shaped B x H*W x 2
+    
+    fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
+    x, y, z = torch.unbind(xyz, dim=-1)
+    B = list(z.shape)[0]
+
+    fx = torch.reshape(fx, [B,1])
+    fy = torch.reshape(fy, [B,1])
+    x0 = torch.reshape(x0, [B,1])
+    y0 = torch.reshape(y0, [B,1])
+    x = torch.reshape(x, [B,-1])
+    y = torch.reshape(y, [B,-1])
+    z = torch.reshape(z, [B,-1])
+
+    EPS = 1e-4
+    z = torch.clamp(z, min=EPS)
+    x = (x*fx)/z + x0
+    y = (y*fy)/z + y0
+    xy = torch.stack([x, y], dim=-1)
+    return xy
+
+def depth2pointcloud(z, pix_T_cam):
+    B, C, H, W = list(z.shape)
+    device = z.device
+    y, x = utils.basic.meshgrid2d(B, H, W, device=device)
+    z = torch.reshape(z, [B, H, W])
+    fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
+    xyz = pixels2camera(x, y, z, fx, fy, x0, y0)
+    return xyz
+
+
+
+def create_depth_image_single(xy, z, H, W, force_positive=True, max_val=100.0, serial=False, slices=20):
+    # turn the xy coordinates into image inds
+    xy = torch.round(xy).long()
+    depth = torch.zeros(H*W, dtype=torch.float32, device=xy.device)
+    depth[:] = max_val
+    
+    # lidar reports a sphere of measurements
+    # only use the inds that are within the image bounds
+    # also, only use forward-pointing depths (z > 0)
+    valid_inds = (xy[:,0] <= W-1) & (xy[:,1] <= H-1) & (xy[:,0] >= 0) & (xy[:,1] >= 0) & (z[:] > 0)
+
+    # gather these up
+    xy = xy[valid_inds]
+    z = z[valid_inds]
+
+    inds = utils.basic.sub2ind(H, W, xy[:,1], xy[:,0]).long()
+    if not serial:
+        depth[inds] = z
+    else:
+        if False:
+            for (index, replacement) in zip(inds, z):
+                if depth[index] > replacement:
+                    depth[index] = replacement
+        # ok my other idea is:
+        # sort the depths by distance
+        # create N depth maps
+        # merge them back-to-front
+
+        # actually maybe you don't even need the separate maps
+
+        sort_inds = torch.argsort(z, descending=True)
+        xy = xy[sort_inds]
+        z = z[sort_inds]
+        N = len(sort_inds)
+        def split(a, n):
+            k, m = divmod(len(a), n)
+            return (a[i*k+min(i, m):(i+1)*k+min(i+1, m)] for i in range(n))
+
+        slice_inds = split(range(N), slices)
+        for si in slice_inds:
+            mini_z = z[si]
+            mini_xy = xy[si]
+            inds = utils.basic.sub2ind(H, W, mini_xy[:,1], mini_xy[:,0]).long()
+            depth[inds] = mini_z
+        # cool; this is rougly as fast as the parallel, and as accurate as the serial
+        
+        if False:
+            print('inds', inds.shape)
+            unique, inverse, counts = torch.unique(inds, return_inverse=True, return_counts=True)
+            print('unique', unique.shape)
+
+            perm = torch.arange(inverse.size(0), dtype=inverse.dtype, device=inverse.device)
+            inverse, perm = inverse.flip([0]), perm.flip([0])
+            perm = inverse.new_empty(unique.size(0)).scatter_(0, inverse, perm)
+
+            # new_inds = inds[inverse_inds]
+            # depth[new_inds] = z[unique_inds]
+
+            depth[unique] = z[perm]
+
+            # now for the duplicates...
+
+            dup = counts > 1
+            dup_unique = unique[dup]
+            print('dup_unique', dup_unique.shape)
+            depth[dup_unique] = 0.5
+        
+    if force_positive:
+        # valid = (depth > 0.0).float()
+        depth[torch.where(depth == 0.0)] = max_val
+    # else:
+    #     valid = torch.ones_like(depth)
+
+    valid = (depth > 0.0).float() * (depth < max_val).float()
+    
+    depth = torch.reshape(depth, [1, H, W])
+    valid = torch.reshape(valid, [1, H, W])
+    return depth, valid
+
+def create_depth_image(pix_T_cam, xyz_cam, H, W, offset_amount=0, max_val=100.0, serial=False, slices=20):
+    B, N, D = list(xyz_cam.shape)
+    assert(D==3)
+    B2, E, F = list(pix_T_cam.shape)
+    assert(B==B2)
+    assert(E==4)
+    assert(F==4)
+    xy = apply_pix_T_cam(pix_T_cam, xyz_cam)
+    z = xyz_cam[:,:,2]
+
+    depth = torch.zeros(B, 1, H, W, dtype=torch.float32, device=xyz_cam.device)
+    valid = torch.zeros(B, 1, H, W, dtype=torch.float32, device=xyz_cam.device)
+    for b in list(range(B)):
+        xy_b, z_b = xy[b], z[b]
+        ind = z_b > 0
+        xy_b = xy_b[ind]
+        z_b = z_b[ind]
+        depth_b, valid_b = create_depth_image_single(xy_b, z_b, H, W, max_val=max_val, serial=serial, slices=slices)
+        if offset_amount:
+            depth_b = depth_b.reshape(-1)
+            valid_b = valid_b.reshape(-1)
+            
+            for off_x in range(offset_amount):
+                for off_y in range(offset_amount):
+                    for sign in [-1,1]:
+                        offset = np.array([sign*off_x,sign*off_y]).astype(np.float32)
+                        offset = torch.from_numpy(offset).reshape(1, 2).to(xyz_cam.device)
+                        # offsets.append(offset)
+                        depth_, valid_ = create_depth_image_single(xy_b + offset, z_b, H, W, max_val=max_val)
+                        depth_ = depth_.reshape(-1)
+                        valid_ = valid_.reshape(-1)
+                        # at invalid locations, use this new value
+                        depth_b[valid_b==0] = depth_[valid_b==0]
+                        valid_b[valid_b==0] = valid_[valid_b==0]
+                    
+            depth_b = depth_b.reshape(1, H, W)
+            valid_b = valid_b.reshape(1, H, W)
+        depth[b] = depth_b
+        valid[b] = valid_b
+    return depth, valid

utils/improc.py

@@ -0,0 +1,1645 @@
+import torch
+import numpy as np
+import utils.basic
+import utils.py
+from sklearn.decomposition import PCA
+from matplotlib import cm
+import matplotlib.pyplot as plt
+import cv2
+import torch.nn.functional as F
+import torchvision
+EPS = 1e-6
+
+from skimage.color import (
+    rgb2lab, rgb2yuv, rgb2ycbcr, lab2rgb, yuv2rgb, ycbcr2rgb,
+    rgb2hsv, hsv2rgb, rgb2xyz, xyz2rgb, rgb2hed, hed2rgb)
+
+def _convert(input_, type_):
+    return {
+        'float': input_.float(),
+        'double': input_.double(),
+    }.get(type_, input_)
+
+def _generic_transform_sk_3d(transform, in_type='', out_type=''):
+    def apply_transform_individual(input_):
+        device = input_.device
+        input_ = input_.cpu()
+        input_ = _convert(input_, in_type)
+
+        input_ = input_.permute(1, 2, 0).detach().numpy()
+        transformed = transform(input_)
+        output = torch.from_numpy(transformed).float().permute(2, 0, 1)
+        output = _convert(output, out_type)
+        return output.to(device)
+
+    def apply_transform(input_):
+        to_stack = []
+        for image in input_:
+            to_stack.append(apply_transform_individual(image))
+        return torch.stack(to_stack)
+    return apply_transform
+
+hsv_to_rgb = _generic_transform_sk_3d(hsv2rgb)
+
+def preprocess_color_tf(x):
+    import tensorflow as tf
+    return tf.cast(x,tf.float32) * 1./255 - 0.5
+
+def preprocess_color(x):
+    if isinstance(x, np.ndarray):
+        return x.astype(np.float32) * 1./255 - 0.5
+    else:
+        return x.float() * 1./255 - 0.5
+
+def pca_embed(emb, keep, valid=None):
+    ## emb -- [S,H/2,W/2,C]
+    ## keep is the number of principal components to keep
+    ## Helper function for reduce_emb.
+    emb = emb + EPS
+    #emb is B x C x H x W
+    emb = emb.permute(0, 2, 3, 1).cpu().detach().numpy() #this is B x H x W x C
+
+    if valid:
+        valid = valid.cpu().detach().numpy().reshape((H*W))
+
+    emb_reduced = list()
+
+    B, H, W, C = np.shape(emb)
+    for img in emb:
+        if np.isnan(img).any():
+            emb_reduced.append(np.zeros([H, W, keep]))
+            continue
+
+        pixels_kd = np.reshape(img, (H*W, C))
+        
+        if valid:
+            pixels_kd_pca = pixels_kd[valid]
+        else:
+            pixels_kd_pca = pixels_kd
+
+        P = PCA(keep)
+        P.fit(pixels_kd_pca)
+
+        if valid:
+            pixels3d = P.transform(pixels_kd)*valid
+        else:
+            pixels3d = P.transform(pixels_kd)
+
+        out_img = np.reshape(pixels3d, [H,W,keep]).astype(np.float32)
+        if np.isnan(out_img).any():
+            emb_reduced.append(np.zeros([H, W, keep]))
+            continue
+
+        emb_reduced.append(out_img)
+
+    emb_reduced = np.stack(emb_reduced, axis=0).astype(np.float32)
+
+    return torch.from_numpy(emb_reduced).permute(0, 3, 1, 2)
+
+def pca_embed_together(emb, keep):
+    ## emb -- [S,H/2,W/2,C]
+    ## keep is the number of principal components to keep
+    ## Helper function for reduce_emb.
+    emb = emb + EPS
+    #emb is B x C x H x W
+    emb = emb.permute(0, 2, 3, 1).cpu().detach().float().numpy() #this is B x H x W x C
+
+    B, H, W, C = np.shape(emb)
+    if np.isnan(emb).any():
+        return torch.zeros(B, keep, H, W)
+    
+    pixelskd = np.reshape(emb, (B*H*W, C))
+    P = PCA(keep)
+    P.fit(pixelskd)
+    pixels3d = P.transform(pixelskd)
+    out_img = np.reshape(pixels3d, [B,H,W,keep]).astype(np.float32)
+        
+    if np.isnan(out_img).any():
+        return torch.zeros(B, keep, H, W)
+    
+    return torch.from_numpy(out_img).permute(0, 3, 1, 2)
+
+def reduce_emb(emb, valid=None, inbound=None, together=False):
+    ## emb -- [S,C,H/2,W/2], inbound -- [S,1,H/2,W/2]
+    ## Reduce number of chans to 3 with PCA. For vis.
+    # S,H,W,C = emb.shape.as_list()
+    S, C, H, W = list(emb.size())
+    keep = 4
+
+    if together:
+        reduced_emb = pca_embed_together(emb, keep)
+    else:
+        reduced_emb = pca_embed(emb, keep, valid) #not im
+
+    reduced_emb = reduced_emb[:,1:]
+    reduced_emb = utils.basic.normalize(reduced_emb) - 0.5
+    if inbound is not None:
+        emb_inbound = emb*inbound
+    else:
+        emb_inbound = None
+
+    return reduced_emb, emb_inbound
+
+def get_feat_pca(feat, valid=None):
+    B, C, D, W = list(feat.size())
+    # feat is B x C x D x W. If 3D input, average it through Height dimension before passing into this function.
+
+    pca, _ = reduce_emb(feat, valid=valid,inbound=None, together=True)
+    # pca is B x 3 x W x D
+    return pca
+
+def gif_and_tile(ims, just_gif=False):
+    S = len(ims) 
+    # each im is B x H x W x C
+    # i want a gif in the left, and the tiled frames on the right
+    # for the gif tool, this means making a B x S x H x W tensor
+    # where the leftmost part is sequential and the rest is tiled
+    gif = torch.stack(ims, dim=1)
+    if just_gif:
+        return gif
+    til = torch.cat(ims, dim=2)
+    til = til.unsqueeze(dim=1).repeat(1, S, 1, 1, 1)
+    im = torch.cat([gif, til], dim=3)
+    return im
+
+def back2color(i, blacken_zeros=False):
+    if blacken_zeros:
+        const = torch.tensor([-0.5])
+        i = torch.where(i==0.0, const.cuda() if i.is_cuda else const, i)
+        return back2color(i)
+    else:
+        return ((i+0.5)*255).type(torch.ByteTensor)
+    
+def xy2heatmap(xy, sigma, grid_xs, grid_ys, norm=False):
+    # xy is B x N x 2, containing float x and y coordinates of N things
+    # grid_xs and grid_ys are B x N x Y x X
+
+    B, N, Y, X = list(grid_xs.shape)
+
+    mu_x = xy[:,:,0].clone()
+    mu_y = xy[:,:,1].clone()
+
+    x_valid = (mu_x>-0.5) & (mu_x<float(X+0.5))
+    y_valid = (mu_y>-0.5) & (mu_y<float(Y+0.5))
+    not_valid = ~(x_valid & y_valid)
+
+    mu_x[not_valid] = -10000
+    mu_y[not_valid] = -10000
+
+    mu_x = mu_x.reshape(B, N, 1, 1).repeat(1, 1, Y, X)
+    mu_y = mu_y.reshape(B, N, 1, 1).repeat(1, 1, Y, X)
+
+    sigma_sq = sigma*sigma
+    # sigma_sq = (sigma*sigma).reshape(B, N, 1, 1)
+    sq_diff_x = (grid_xs - mu_x)**2
+    sq_diff_y = (grid_ys - mu_y)**2
+
+    term1 = 1./2.*np.pi*sigma_sq
+    term2 = torch.exp(-(sq_diff_x+sq_diff_y)/(2.*sigma_sq))
+    gauss = term1*term2
+
+    if norm:
+        # normalize so each gaussian peaks at 1
+        gauss_ = gauss.reshape(B*N, Y, X)
+        gauss_ = utils.basic.normalize(gauss_)
+        gauss = gauss_.reshape(B, N, Y, X)
+
+    return gauss
+
+def xy2heatmaps(xy, Y, X, sigma=30.0, norm=True):
+    # xy is B x N x 2
+
+    B, N, D = list(xy.shape)
+    assert(D==2)
+
+    device = xy.device
+    
+    grid_y, grid_x = utils.basic.meshgrid2d(B, Y, X, device=device)
+    # grid_x and grid_y are B x Y x X
+    grid_xs = grid_x.unsqueeze(1).repeat(1, N, 1, 1)
+    grid_ys = grid_y.unsqueeze(1).repeat(1, N, 1, 1)
+    heat = xy2heatmap(xy, sigma, grid_xs, grid_ys, norm=norm)
+    return heat
+
+def draw_circles_at_xy(xy, Y, X, sigma=12.5, round=False):
+    B, N, D = list(xy.shape)
+    assert(D==2)
+    prior = xy2heatmaps(xy, Y, X, sigma=sigma)
+    # prior is B x N x Y x X
+    if round:
+        prior = (prior > 0.5).float()
+    return prior
+
+def seq2color(im, norm=True, colormap='coolwarm'):
+    B, S, H, W = list(im.shape)
+    # S is sequential
+
+    # prep a mask of the valid pixels, so we can blacken the invalids later
+    mask = torch.max(im, dim=1, keepdim=True)[0]
+
+    # turn the S dim into an explicit sequence
+    coeffs = np.linspace(1.0, float(S), S).astype(np.float32)/float(S)
+    
+    # # increase the spacing from the center
+    # coeffs[:int(S/2)] -= 2.0
+    # coeffs[int(S/2)+1:] += 2.0
+    
+    coeffs = torch.from_numpy(coeffs).float().cuda()
+    coeffs = coeffs.reshape(1, S, 1, 1).repeat(B, 1, H, W)
+    # scale each channel by the right coeff
+    im = im * coeffs
+    # now im is in [1/S, 1], except for the invalid parts which are 0
+    # keep the highest valid coeff at each pixel
+    im = torch.max(im, dim=1, keepdim=True)[0]
+
+    out = []
+    for b in range(B):
+        im_ = im[b]
+        # move channels out to last dim_
+        im_ = im_.detach().cpu().numpy()
+        im_ = np.squeeze(im_)
+        # im_ is H x W
+        if colormap=='coolwarm':
+            im_ = cm.coolwarm(im_)[:, :, :3]
+        elif colormap=='PiYG':
+            im_ = cm.PiYG(im_)[:, :, :3]
+        elif colormap=='winter':
+            im_ = cm.winter(im_)[:, :, :3]
+        elif colormap=='spring':
+            im_ = cm.spring(im_)[:, :, :3]
+        elif colormap=='onediff':
+            im_ = np.reshape(im_, (-1))
+            im0_ = cm.spring(im_)[:, :3]
+            im1_ = cm.winter(im_)[:, :3]
+            im1_[im_==1/float(S)] = im0_[im_==1/float(S)]
+            im_ = np.reshape(im1_, (H, W, 3))
+        else:
+            assert(False) # invalid colormap
+        # move channels into dim 0
+        im_ = np.transpose(im_, [2, 0, 1])
+        im_ = torch.from_numpy(im_).float().cuda()
+        out.append(im_)
+    out = torch.stack(out, dim=0)
+    
+    # blacken the invalid pixels, instead of using the 0-color
+    out = out*mask
+    # out = out*255.0
+
+    # put it in [-0.5, 0.5]
+    out = out - 0.5
+    
+    return out
+
+def colorize(d):
+    # this is actually just grayscale right now
+
+    if d.ndim==2:
+        d = d.unsqueeze(dim=0)
+    else:
+        assert(d.ndim==3)
+        
+    # color_map = cm.get_cmap('plasma')
+    color_map = cm.get_cmap('inferno') 
+    # S1, D = traj.shape
+
+    # print('d1', d.shape)
+    C,H,W = d.shape
+    assert(C==1)
+    d = d.reshape(-1)
+    d = d.detach().cpu().numpy()
+    # print('d2', d.shape)
+    color = np.array(color_map(d)) * 255 # rgba
+    # print('color1', color.shape)
+    color = np.reshape(color[:,:3], [H*W, 3])
+    # print('color2', color.shape)
+    color = torch.from_numpy(color).permute(1,0).reshape(3,H,W)
+    # # gather
+    # cm = matplotlib.cm.get_cmap(cmap if cmap is not None else 'gray')
+    # if cmap=='RdBu' or cmap=='RdYlGn':
+    #     colors = cm(np.arange(256))[:, :3]
+    #  else:
+    #      colors = cm.colors
+    #      colors = np.array(colors).astype(np.float32)
+    #      colors = np.reshape(colors, [-1, 3])
+    #      colors = tf.constant(colors, dtype=tf.float32)
+
+    #      value = tf.gather(colors, indices)
+    # colorize(value, normalize=True, vmin=None, vmax=None, cmap=None, vals=255)
+        
+    # copy to the three chans
+    # d = d.repeat(3, 1, 1)
+    return color
+
+
+def oned2inferno(d, norm=True, do_colorize=False):
+    # convert a 1chan input to a 3chan image output
+
+    # if it's just B x H x W, add a C dim
+    if d.ndim==3:
+        d = d.unsqueeze(dim=1)
+    # d should be B x C x H x W, where C=1
+    B, C, H, W = list(d.shape)
+    assert(C==1)
+
+    if norm:
+        d = utils.basic.normalize(d)
+        
+    if do_colorize:
+        rgb = torch.zeros(B, 3, H, W)
+        for b in list(range(B)):
+            rgb[b] = colorize(d[b])
+    else:
+        rgb = d.repeat(1, 3, 1, 1)*255.0
+    # rgb = (255.0*rgb).type(torch.ByteTensor)
+    rgb = rgb.type(torch.ByteTensor)
+
+    # rgb = tf.cast(255.0*rgb, tf.uint8)
+    # rgb = tf.reshape(rgb, [-1, hyp.H, hyp.W, 3])
+    # rgb = tf.expand_dims(rgb, axis=0)
+    return rgb
+
+def oned2gray(d, norm=True):
+    # convert a 1chan input to a 3chan image output
+
+    # if it's just B x H x W, add a C dim
+    if d.ndim==3:
+        d = d.unsqueeze(dim=1)
+    # d should be B x C x H x W, where C=1
+    B, C, H, W = list(d.shape)
+    assert(C==1)
+    
+    if norm:
+        d = utils.basic.normalize(d)
+
+    rgb = d.repeat(1,3,1,1)
+    rgb = (255.0*rgb).type(torch.ByteTensor)
+    return rgb
+
+
+def draw_frame_id_on_vis(vis, frame_id, scale=0.5, left=5, top=20, shadow=True):
+
+    rgb = vis.detach().cpu().numpy()[0]
+    rgb = np.transpose(rgb, [1, 2, 0]) # put channels last
+    rgb = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) 
+    color = (255, 255, 255)
+    # print('putting frame id', frame_id)
+
+    frame_str = utils.basic.strnum(frame_id)
+
+    text_color_bg = (0,0,0)
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    text_size, _ = cv2.getTextSize(frame_str, font, scale, 1)
+    text_w, text_h = text_size
+    if shadow:
+        cv2.rectangle(rgb, (left, top-text_h), (left + text_w, top+1), text_color_bg, -1)
+    
+    cv2.putText(
+        rgb,
+        frame_str,
+        (left, top), # from left, from top
+        font,
+        scale, # font scale (float)
+        color, 
+        1) # font thickness (int)
+    rgb = cv2.cvtColor(rgb.astype(np.uint8), cv2.COLOR_BGR2RGB)
+    vis = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)
+    return vis
+
+def draw_frame_str_on_vis(vis, frame_str, scale=0.5, left=5, top=40, shadow=True):
+
+    rgb = vis.detach().cpu().numpy()[0]
+    rgb = np.transpose(rgb, [1, 2, 0]) # put channels last
+    rgb = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) 
+    color = (255, 255, 255)
+
+    text_color_bg = (0,0,0)
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    text_size, _ = cv2.getTextSize(frame_str, font, scale, 1)
+    text_w, text_h = text_size
+    if shadow:
+        cv2.rectangle(rgb, (left, top-text_h), (left + text_w, top+1), text_color_bg, -1)
+    
+    cv2.putText(
+        rgb,
+        frame_str,
+        (left, top), # from left, from top
+        font, 
+        scale, # font scale (float)
+        color, 
+        1) # font thickness (int)
+    rgb = cv2.cvtColor(rgb.astype(np.uint8), cv2.COLOR_BGR2RGB)
+    vis = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)
+    return vis
+
+COLORMAP_FILE = "./utils/bremm.png"
+class ColorMap2d:
+    def __init__(self, filename=None):
+        self._colormap_file = filename or COLORMAP_FILE
+        self._img = plt.imread(self._colormap_file)
+        
+        self._height = self._img.shape[0]
+        self._width = self._img.shape[1]
+
+    def __call__(self, X):
+        assert len(X.shape) == 2
+        output = np.zeros((X.shape[0], 3))
+        for i in range(X.shape[0]):
+            x, y = X[i, :]
+            xp = int((self._width-1) * x)
+            yp = int((self._height-1) * y)
+            xp = np.clip(xp, 0, self._width-1)
+            yp = np.clip(yp, 0, self._height-1)
+            output[i, :] = self._img[yp, xp]
+        return output
+    
+def get_n_colors(N, sequential=False):
+    label_colors = []
+    for ii in range(N):
+        if sequential:
+            rgb = cm.winter(ii/(N-1))
+            rgb = (np.array(rgb) * 255).astype(np.uint8)[:3]
+        else:
+            # rgb = np.zeros(3)
+            # while np.sum(rgb) < 128: # ensure min brightness
+            #     rgb = np.random.randint(0,256,3)
+            rgb = cm.gist_rainbow(ii/(N-1))
+            rgb = (np.array(rgb) * 255).astype(np.uint8)[:3]
+            
+        label_colors.append(rgb)
+    return label_colors
+
+class Summ_writer(object):
+    def __init__(self, writer, global_step, log_freq=10, fps=8, scalar_freq=100, just_gif=False):
+        self.writer = writer
+        self.global_step = global_step
+        self.log_freq = log_freq
+        self.scalar_freq = scalar_freq
+        self.fps = fps
+        self.just_gif = just_gif
+        self.maxwidth = 10000
+        self.save_this = (self.global_step % self.log_freq == 0)
+        self.scalar_freq = max(scalar_freq,1)
+        self.save_scalar = (self.global_step % self.scalar_freq == 0)
+        if self.save_this:
+            self.save_scalar = True
+
+    def summ_gif(self, name, tensor, blacken_zeros=False):
+        # tensor should be in B x S x C x H x W
+        
+        assert tensor.dtype in {torch.uint8,torch.float32}
+        shape = list(tensor.shape)
+
+        if tensor.dtype == torch.float32:
+            tensor = back2color(tensor, blacken_zeros=blacken_zeros)
+
+        video_to_write = tensor[0:1]
+
+        S = video_to_write.shape[1]
+        if S==1:
+            # video_to_write is 1 x 1 x C x H x W
+            self.writer.add_image(name, video_to_write[0,0], global_step=self.global_step)
+        else:
+            self.writer.add_video(name, video_to_write, fps=self.fps, global_step=self.global_step)
+            
+        return video_to_write
+
+    def draw_boxlist2d_on_image(self, rgb, boxlist, scores=None, tids=None, linewidth=1):
+        B, C, H, W = list(rgb.shape)
+        assert(C==3)
+        B2, N, D = list(boxlist.shape)
+        assert(B2==B)
+        assert(D==4) # ymin, xmin, ymax, xmax
+
+        rgb = back2color(rgb)
+        if scores is None:
+            scores = torch.ones(B2, N).float()
+        if tids is None:
+            # tids = torch.arange(N).reshape(1,N).repeat(B2,1).long()
+            tids = torch.zeros(B2, N).long()
+        out = self.draw_boxlist2d_on_image_py(
+            rgb[0].cpu().detach().numpy(),
+            boxlist[0].cpu().detach().numpy(),
+            scores[0].cpu().detach().numpy(),
+            tids[0].cpu().detach().numpy(),
+            linewidth=linewidth)
+        out = torch.from_numpy(out).type(torch.ByteTensor).permute(2, 0, 1)
+        out = torch.unsqueeze(out, dim=0)
+        out = preprocess_color(out)
+        out = torch.reshape(out, [1, C, H, W])
+        return out
+    
+    def draw_boxlist2d_on_image_py(self, rgb, boxlist, scorelist, tidlist, linewidth=1):
+        # all inputs are numpy tensors
+        # rgb is H x W x 3
+        # boxlist is N x 4
+        # scorelist is N
+        # tidlist is N
+
+        rgb = np.transpose(rgb, [1, 2, 0]) # put channels last
+        # rgb = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
+        rgb = rgb.astype(np.uint8).copy()
+
+        H, W, C = rgb.shape
+        assert(C==3)
+        N, D = boxlist.shape
+        assert(D==4)
+        M = scorelist.shape[0]
+        assert(M==N)
+        O = tidlist.shape[0]
+        assert(M==O)
+
+        # color_map = cm.get_cmap('Accent')
+        # color_map = cm.get_cmap('Set3')
+        color_map = cm.get_cmap('tab20')
+        color_map = color_map.colors
+        # print('color_map', color_map)
+        
+        # draw
+        for (box, score, tid) in zip(boxlist, scorelist, tidlist):
+            # box is 4
+            if not np.isclose(score, 0.0, atol=1e-3):
+                # ymin, xmin, ymax, xmax = box
+                xmin, ymin, xmax, ymax = box
+                
+                color = color_map[tid]
+                color = np.array(color)*255.0
+                color = color.round()
+
+                if not np.isclose(score, 1.0, atol=1e-3):
+                    cv2.putText(rgb,
+                                # '%d (%.2f)' % (tidlist[ind], scorelist[ind]), 
+                                '%.2f' % (score), 
+                                (int(xmin), int(ymin)),
+                                cv2.FONT_HERSHEY_SIMPLEX,
+                                0.5, # font size
+                                color),
+
+                xmin = int(np.clip(xmin,0,W-1))
+                ymin = int(np.clip(ymin,0,H-1))
+                xmax = int(np.clip(xmax,0,W-1))
+                ymax = int(np.clip(ymax,0,H-1))
+
+                # print('xmin, xmax, ymin, ymax', xmin, xmax, ymin, ymax)
+                
+                cv2.line(rgb, (xmin, ymin), (xmin, ymax), color, linewidth, cv2.LINE_4)
+                cv2.line(rgb, (xmin, ymin), (xmax, ymin), color, linewidth, cv2.LINE_4)
+                cv2.line(rgb, (xmax, ymin), (xmax, ymax), color, linewidth, cv2.LINE_4)
+                cv2.line(rgb, (xmax, ymax), (xmin, ymax), color, linewidth, cv2.LINE_4)
+                
+        # rgb = cv2.cvtColor(rgb.astype(np.uint8), cv2.COLOR_BGR2RGB)
+        return rgb
+    
+    def summ_boxlist2d(self, name, rgb, boxlist, scores=None, tids=None, frame_id=None, frame_str=None, only_return=False, linewidth=1):
+        B, C, H, W = list(rgb.shape)
+        boxlist_vis = self.draw_boxlist2d_on_image(rgb, boxlist, scores=scores, tids=tids, linewidth=linewidth)
+        return self.summ_rgb(name, boxlist_vis, frame_id=frame_id, frame_str=frame_str, only_return=only_return)
+    
+    def summ_rgbs(self, name, ims, frame_ids=None, frame_strs=None, blacken_zeros=False, only_return=False):
+        if self.save_this:
+
+            ims = gif_and_tile(ims, just_gif=self.just_gif)
+            vis = ims
+
+            assert vis.dtype in {torch.uint8,torch.float32}
+
+            if vis.dtype == torch.float32:
+                vis = back2color(vis, blacken_zeros)           
+
+            B, S, C, H, W = list(vis.shape)
+
+            if frame_ids is not None:
+                assert(len(frame_ids)==S)
+                for s in range(S):
+                    vis[:,s] = draw_frame_id_on_vis(vis[:,s], frame_ids[s])
+                    
+            if frame_strs is not None:
+                assert(len(frame_strs)==S)
+                for s in range(S):
+                    vis[:,s] = draw_frame_str_on_vis(vis[:,s], frame_strs[s])
+
+            if int(W) > self.maxwidth:
+                vis = vis[:,:,:,:self.maxwidth]
+
+            if only_return:
+                return vis
+            else:
+                return self.summ_gif(name, vis, blacken_zeros)
+
+    def summ_rgb(self, name, ims, blacken_zeros=False, frame_id=None, frame_str=None, only_return=False, halfres=False, shadow=True):
+        if self.save_this:
+            assert ims.dtype in {torch.uint8,torch.float32}
+
+            if ims.dtype == torch.float32:
+                ims = back2color(ims, blacken_zeros)
+
+            #ims is B x C x H x W
+            vis = ims[0:1] # just the first one
+            B, C, H, W = list(vis.shape)
+
+            if halfres:
+                vis = F.interpolate(vis, scale_factor=0.5)
+
+            if frame_id is not None:
+                vis = draw_frame_id_on_vis(vis, frame_id, shadow=shadow)
+                
+            if frame_str is not None:
+                vis = draw_frame_str_on_vis(vis, frame_str, shadow=shadow)
+
+            if int(W) > self.maxwidth:
+                vis = vis[:,:,:,:self.maxwidth]
+
+            if only_return:
+                return vis
+            else:
+                return self.summ_gif(name, vis.unsqueeze(1), blacken_zeros)
+
+    def flow2color(self, flow, clip=0.0):
+        B, C, H, W = list(flow.size())
+        assert(C==2)
+        flow = flow[0:1].detach()
+
+        if False:
+            flow = flow[0].detach().cpu().permute(1,2,0).numpy() # H,W,2
+            if clip > 0:
+                clip_flow = clip
+            else:
+                clip_flow = None
+            im = utils.py.flow_to_image(flow, clip_flow=clip_flow, convert_to_bgr=True)
+            # im = utils.py.flow_to_image(flow, convert_to_bgr=True)
+            im = torch.from_numpy(im).permute(2,0,1).unsqueeze(0).byte() # 1,3,H,W
+            im = torch.flip(im, dims=[1]).clone() # BGR
+
+            # # i prefer black bkg
+            # white_pixels = (im == 255).all(dim=1, keepdim=True)
+            # im[white_pixels.expand(-1, 3, -1, -1)] = 0
+
+            return im
+              
+        # flow_abs = torch.abs(flow)
+        # flow_mean = flow_abs.mean(dim=[1,2,3])
+        # flow_std = flow_abs.std(dim=[1,2,3])
+        if clip==0:
+            clip = torch.max(torch.abs(flow)).item()
+
+        # if clip:
+        flow = torch.clamp(flow, -clip, clip)/clip
+        # else:
+        #     # # Apply some kind of normalization. Divide by the perceived maximum (mean + std*2)
+        #     # flow_max = flow_mean + flow_std*2 + 1e-10
+        #     # for b in range(B):
+        #     #     flow[b] = flow[b].clamp(-flow_max[b].item(), flow_max[b].item()) / flow_max[b].clamp(min=1)
+
+        #     flow_max = torch.max(flow_abs[b])
+        #     for b in range(B):
+        #         flow[b] = flow[b].clamp(-flow_max.item(), flow_max.item()) / flow_max[b].clamp(min=1)
+            
+
+        radius = torch.sqrt(torch.sum(flow**2, dim=1, keepdim=True)) #B x 1 x H x W
+        radius_clipped = torch.clamp(radius, 0.0, 1.0)
+
+        angle = torch.atan2(-flow[:, 1:2], -flow[:, 0:1]) / np.pi # B x 1 x H x W
+
+        hue = torch.clamp((angle + 1.0) / 2.0, 0.0, 1.0)
+        # hue = torch.mod(angle / (2 * np.pi) + 1.0, 1.0)
+            
+        saturation = torch.ones_like(hue) * 0.75
+        value = radius_clipped
+        hsv = torch.cat([hue, saturation, value], dim=1) #B x 3 x H x W
+
+        #flow = tf.image.hsv_to_rgb(hsv)
+        flow = hsv_to_rgb(hsv)
+        flow = (flow*255.0).type(torch.ByteTensor)
+        # flow = torch.flip(flow, dims=[1]).clone() # BGR
+        return flow
+    
+    def summ_flow(self, name, im, clip=0.0, only_return=False, frame_id=None, frame_str=None, shadow=True):
+        # flow is B x C x D x W
+        if self.save_this:
+            return self.summ_rgb(name, self.flow2color(im, clip=clip), only_return=only_return, frame_id=frame_id, frame_str=frame_str, shadow=shadow)
+        else:
+            return None
+            
+    def summ_oneds(self, name, ims, frame_ids=None, frame_strs=None, bev=False, fro=False, logvis=False, reduce_max=False, max_val=0.0, norm=True, only_return=False, do_colorize=False):
+        if self.save_this:
+            if bev: 
+                B, C, H, _, W = list(ims[0].shape)
+                if reduce_max:
+                    ims = [torch.max(im, dim=3)[0] for im in ims]
+                else:
+                    ims = [torch.mean(im, dim=3) for im in ims]
+            elif fro: 
+                B, C, _, H, W = list(ims[0].shape)
+                if reduce_max:
+                    ims = [torch.max(im, dim=2)[0] for im in ims]
+                else:
+                    ims = [torch.mean(im, dim=2) for im in ims]
+
+
+            if len(ims) != 1: # sequence
+                im = gif_and_tile(ims, just_gif=self.just_gif)
+            else:
+                im = torch.stack(ims, dim=1) # single frame
+
+            B, S, C, H, W = list(im.shape)
+            
+            if logvis and max_val:
+                max_val = np.log(max_val)
+                im = torch.log(torch.clamp(im, 0)+1.0)
+                im = torch.clamp(im, 0, max_val)
+                im = im/max_val
+                norm = False
+            elif max_val:
+                im = torch.clamp(im, 0, max_val)
+                im = im/max_val
+                norm = False
+                
+            if norm:
+                # normalize before oned2inferno,
+                # so that the ranges are similar within B across S
+                im = utils.basic.normalize(im)
+
+            im = im.view(B*S, C, H, W)
+            vis = oned2inferno(im, norm=norm, do_colorize=do_colorize)
+            vis = vis.view(B, S, 3, H, W)
+
+            if frame_ids is not None:
+                assert(len(frame_ids)==S)
+                for s in range(S):
+                    vis[:,s] = draw_frame_id_on_vis(vis[:,s], frame_ids[s])
+
+            if frame_strs is not None:
+                assert(len(frame_strs)==S)
+                for s in range(S):
+                    vis[:,s] = draw_frame_str_on_vis(vis[:,s], frame_strs[s])
+
+            if W > self.maxwidth:
+                vis = vis[...,:self.maxwidth]
+
+            if only_return:
+                return vis
+            else:
+                self.summ_gif(name, vis)
+
+    def summ_oned(self, name, im, bev=False, fro=False, logvis=False, max_val=0, max_along_y=False, norm=True, frame_id=None, frame_str=None, only_return=False, shadow=True):
+        if self.save_this:
+
+            if bev: 
+                B, C, H, _, W = list(im.shape)
+                if max_along_y:
+                    im = torch.max(im, dim=3)[0]
+                else:
+                    im = torch.mean(im, dim=3)
+            elif fro:
+                B, C, _, H, W = list(im.shape)
+                if max_along_y:
+                    im = torch.max(im, dim=2)[0]
+                else:
+                    im = torch.mean(im, dim=2)
+            else:
+                B, C, H, W = list(im.shape)
+                
+            im = im[0:1] # just the first one
+            assert(C==1)
+            
+            if logvis and max_val:
+                max_val = np.log(max_val)
+                im = torch.log(im)
+                im = torch.clamp(im, 0, max_val)
+                im = im/max_val
+                norm = False
+            elif max_val:
+                im = torch.clamp(im, 0, max_val)/max_val
+                norm = False
+
+            vis = oned2inferno(im, norm=norm)
+            if W > self.maxwidth:
+                vis = vis[...,:self.maxwidth]
+            return self.summ_rgb(name, vis, blacken_zeros=False, frame_id=frame_id, frame_str=frame_str, only_return=only_return, shadow=shadow)
+
+    def summ_4chan(self, name, im, norm=True, frame_id=None, frame_str=None, only_return=False):
+        if self.save_this:
+
+            B, C, H, W = list(im.shape)
+                
+            im = im[0:1] # just the first one
+            assert(C==4)
+
+            # d = utils.basic.normalize(d)
+            im0 = im[:,0:1]
+            im1 = im[:,1:2]
+            im2 = im[:,2:3]
+            im3 = im[:,3:4]
+
+            im0 = utils.basic.normalize(im0).round()
+            im1 = utils.basic.normalize(im1).round()
+            im2 = utils.basic.normalize(im2).round()
+            im3 = utils.basic.normalize(im3).round()
+
+            # kp_vis = sw.summ_rgbs('tff/2_kp_s%d' % s, kp.unbind(1), only_return=True)
+            # kp_any = (torch.max(kp_vis, dim=2, keepdims=True)[0]).repeat(1, 1, 3, 1, 1)
+            # kp_vis[kp_any==0] = fcp_vis[kp_any==0]
+            
+            # vis0 = oned2inferno(im0, norm=False)
+            # vis1 = oned2inferno(im1, norm=False)
+            # vis2 = oned2inferno(im2, norm=False)
+            # vis3 = oned2inferno(im3, norm=False)
+
+            # vis0 = self.summ_seg('', im0[:,0:1]*1, only_return=True, frame_id=frame_id, frame_str=frame_str, colormap='tab20')
+            # vis1 = self.summ_seg('', im1[:,0:1]*2, only_return=True, frame_id=frame_id, frame_str=frame_str, colormap='tab20')
+            # vis2 = self.summ_seg('', im2[:,0:1]*3, only_return=True, frame_id=frame_id, frame_str=frame_str, colormap='tab20')
+            # vis3 = self.summ_seg('', im3[:,0:1]*4, only_return=True, frame_id=frame_id, frame_str=frame_str, colormap='tab20')
+
+            vis0 = self.summ_seg('', im0[:,0]*1, only_return=True, colormap='tab20')
+            vis1 = self.summ_seg('', im1[:,0]*2, only_return=True, colormap='tab20')
+            vis2 = self.summ_seg('', im2[:,0]*3, only_return=True, colormap='tab20')
+            vis3 = self.summ_seg('', im3[:,0]*4, only_return=True, colormap='tab20')
+
+            # vis_any = (torch.max(vis2, dim=2, keepdims=True)[0]).repeat(1, 1, 3, 1, 1)
+            # vis3[vis_any==0] = fcp_vis[kp_any==0]
+
+            vis0_any = (torch.max(vis0, dim=1, keepdims=True)[0]).repeat(1, 3, 1, 1)
+            vis1_any = (torch.max(vis1, dim=1, keepdims=True)[0]).repeat(1, 3, 1, 1)
+            vis2_any = (torch.max(vis2, dim=1, keepdims=True)[0]).repeat(1, 3, 1, 1)
+            vis3_any = (torch.max(vis3, dim=1, keepdims=True)[0]).repeat(1, 3, 1, 1)
+            vis0[vis0_any==0] = vis1[vis0_any==0]
+            vis0[vis1_any==0] = vis2[vis1_any==0]
+            vis0[vis2_any==0] = vis3[vis2_any==0]
+
+            print('vis0', vis0.shape, vis0.device)
+
+            vis0 = vis0.cpu()
+            
+            # vis = oned2inferno(im, norm=norm)
+            # if W > self.maxwidth:
+            #     vis = vis[...,:self.maxwidth]
+            return self.summ_rgb(name, vis0, blacken_zeros=False, frame_id=frame_id, frame_str=frame_str, only_return=only_return)
+        
+            
+    def summ_feats(self, name, feats, valids=None, pca=True, fro=False, only_return=False, frame_ids=None, frame_strs=None):
+        if self.save_this:
+            if valids is not None:
+                valids = torch.stack(valids, dim=1)
+            
+            feats  = torch.stack(feats, dim=1)
+            # feats leads with B x S x C
+
+            if feats.ndim==6:
+
+                # feats is B x S x C x D x H x W
+                if fro:
+                    reduce_dim = 3
+                else:
+                    reduce_dim = 4
+                    
+                if valids is None:
+                    feats = torch.mean(feats, dim=reduce_dim)
+                else: 
+                    valids = valids.repeat(1, 1, feats.size()[2], 1, 1, 1)
+                    feats = utils.basic.reduce_masked_mean(feats, valids, dim=reduce_dim)
+
+            B, S, C, D, W = list(feats.size())
+
+            if not pca:
+                # feats leads with B x S x C
+                feats = torch.mean(torch.abs(feats), dim=2, keepdims=True)
+                # feats leads with B x S x 1
+                feats = torch.unbind(feats, dim=1)
+                return self.summ_oneds(name=name, ims=feats, norm=True, only_return=only_return, frame_ids=frame_ids, frame_strs=frame_strs)
+
+            else:
+                __p = lambda x: utils.basic.pack_seqdim(x, B)
+                __u = lambda x: utils.basic.unpack_seqdim(x, B)
+
+                feats_  = __p(feats)
+                
+                if valids is None:
+                    feats_pca_ = get_feat_pca(feats_)
+                else:
+                    valids_ = __p(valids)
+                    feats_pca_ = get_feat_pca(feats_, valids)
+
+                feats_pca = __u(feats_pca_)
+
+                return self.summ_rgbs(name=name, ims=torch.unbind(feats_pca, dim=1), only_return=only_return, frame_ids=frame_ids, frame_strs=frame_strs)
+
+    def summ_feat(self, name, feat, valid=None, pca=True, only_return=False, bev=False, fro=False, frame_id=None, frame_str=None):
+        if self.save_this:
+            if feat.ndim==5: # B x C x D x H x W
+
+                if bev:
+                    reduce_axis = 3
+                elif fro:
+                    reduce_axis = 2
+                else:
+                    # default to bev
+                    reduce_axis = 3
+                
+                if valid is None:
+                    feat = torch.mean(feat, dim=reduce_axis)
+                else:
+                    valid = valid.repeat(1, feat.size()[1], 1, 1, 1)
+                    feat = utils.basic.reduce_masked_mean(feat, valid, dim=reduce_axis)
+                    
+            B, C, D, W = list(feat.shape)
+
+            if not pca:
+                feat = torch.mean(torch.abs(feat), dim=1, keepdims=True)
+                # feat is B x 1 x D x W
+                return self.summ_oned(name=name, im=feat, norm=True, only_return=only_return, frame_id=frame_id, frame_str=frame_str)
+            else:
+                feat_pca = get_feat_pca(feat, valid)
+                return self.summ_rgb(name, feat_pca, only_return=only_return, frame_id=frame_id, frame_str=frame_str)
+            
+    def summ_scalar(self, name, value):
+        if (not (isinstance(value, int) or isinstance(value, float) or isinstance(value, np.float32))) and ('torch' in value.type()):
+            value = value.detach().cpu().numpy()
+        if not np.isnan(value):
+            if (self.log_freq == 1):
+                self.writer.add_scalar(name, value, global_step=self.global_step)
+            elif self.save_this or self.save_scalar:
+                self.writer.add_scalar(name, value, global_step=self.global_step)
+                
+    def summ_seg(self, name, seg, only_return=False, frame_id=None, frame_str=None, colormap='tab20', label_colors=None):
+        if not self.save_this:
+            return
+
+        B,H,W = seg.shape
+
+        if label_colors is None:
+            custom_label_colors = False
+            # label_colors = get_n_colors(int(torch.max(seg).item()), sequential=True)
+            label_colors = cm.get_cmap(colormap).colors
+            label_colors = [[int(i*255) for i in l] for l in label_colors]
+        else:
+            custom_label_colors = True
+        # label_colors = matplotlib.cm.get_cmap(colormap).colors
+        # label_colors = [[int(i*255) for i in l] for l in label_colors]
+        # print('label_colors', label_colors)
+        
+        # label_colors = [
+        #     (0, 0, 0),         # None
+        #     (70, 70, 70),      # Buildings
+        #     (190, 153, 153),   # Fences
+        #     (72, 0, 90),       # Other
+        #     (220, 20, 60),     # Pedestrians
+        #     (153, 153, 153),   # Poles
+        #     (157, 234, 50),    # RoadLines
+        #     (128, 64, 128),    # Roads
+        #     (244, 35, 232),    # Sidewalks
+        #     (107, 142, 35),    # Vegetation
+        #     (0, 0, 255),      # Vehicles
+        #     (102, 102, 156),  # Walls
+        #     (220, 220, 0)     # TrafficSigns
+        # ]
+
+        r = torch.zeros_like(seg,dtype=torch.uint8)
+        g = torch.zeros_like(seg,dtype=torch.uint8)
+        b = torch.zeros_like(seg,dtype=torch.uint8)
+        
+        for label in range(0,len(label_colors)):
+            if (not custom_label_colors):# and (N > 20):
+                label_ = label % 20
+            else:
+                label_ = label
+            
+            idx = (seg == label)
+            r[idx] = label_colors[label_][0]
+            g[idx] = label_colors[label_][1]
+            b[idx] = label_colors[label_][2]
+            
+        rgb = torch.stack([r,g,b],axis=1)
+        return self.summ_rgb(name,rgb,only_return=only_return, frame_id=frame_id, frame_str=frame_str)
+    
+    def summ_traj2ds_on_rgbs(self, name, trajs, rgbs, visibs=None, valids=None, frame_ids=None, frame_strs=None, only_return=False, show_dots=True, cmap='coolwarm', vals=None, linewidth=1, max_show=1024):
+        # trajs is B, S, N, 2
+        # rgbs is B, S, C, H, W
+        B, S, C, H, W = rgbs.shape
+        B, S2, N, D = trajs.shape
+        assert(S==S2)
+
+
+        rgbs = rgbs[0] # S, C, H, W
+        trajs = trajs[0] # S, N, 2
+        if valids is None:
+            valids = torch.ones_like(trajs[:,:,0]) # S, N
+        else:
+            valids = valids[0]
+            
+        if visibs is None:
+            visibs = torch.ones_like(trajs[:,:,0]) # S, N
+        else:
+            visibs = visibs[0]
+
+        if vals is not None:
+            vals = vals[0] # N
+            # print('vals', vals.shape)
+
+        if N > max_show:
+            inds = np.random.choice(N, max_show)
+            trajs = trajs[:,inds]
+            valids = valids[:,inds]
+            visibs = visibs[:,inds]
+            if vals is not None:
+                vals = vals[inds]
+            N = trajs.shape[1]
+
+        trajs = trajs.clamp(-16, W+16)
+        
+        rgbs_color = []
+        for rgb in rgbs:
+            rgb = back2color(rgb).detach().cpu().numpy() 
+            rgb = np.transpose(rgb, [1, 2, 0]) # put channels last
+            rgbs_color.append(rgb) # each element 3 x H x W
+
+        for i in range(min(N, max_show)):
+            if cmap=='onediff' and i==0:
+                cmap_ = 'spring'
+            elif cmap=='onediff':
+                cmap_ = 'winter'
+            else:
+                cmap_ = cmap
+            traj = trajs[:,i].long().detach().cpu().numpy() # S, 2
+            valid = valids[:,i].long().detach().cpu().numpy() # S
+            
+            # print('traj', traj.shape)
+            # print('valid', valid.shape)
+            
+            if vals is not None:
+                # val = vals[:,i].float().detach().cpu().numpy() # []
+                val = vals[i].float().detach().cpu().numpy() # []
+                # print('val', val.shape)
+            else:
+                val = None
+            
+            for t in range(S):
+                if valid[t]:
+                    rgbs_color[t] = self.draw_traj_on_image_py(rgbs_color[t], traj[:t+1], S=S, show_dots=show_dots, cmap=cmap_, val=val, linewidth=linewidth)
+
+        for i in range(min(N, max_show)):
+            if cmap=='onediff' and i==0:
+                cmap_ = 'spring'
+            elif cmap=='onediff':
+                cmap_ = 'winter'
+            else:
+                cmap_ = cmap
+            traj = trajs[:,i] # S,2
+            vis = visibs[:,i].round() # S
+            valid = valids[:,i] # S
+            rgbs_color = self.draw_circ_on_images_py(rgbs_color, traj, vis, S=S, show_dots=show_dots, cmap=cmap_, linewidth=linewidth)
+
+        rgbs = []
+        for rgb in rgbs_color:
+            rgb = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)
+            rgbs.append(preprocess_color(rgb))
+
+        return self.summ_rgbs(name, rgbs, only_return=only_return, frame_ids=frame_ids, frame_strs=frame_strs)
+
+    def summ_traj2ds_on_rgbs2(self, name, trajs, visibles, rgbs, valids=None, frame_ids=None, frame_strs=None, only_return=False, show_dots=True, cmap=None, linewidth=1, max_show=1024):
+        # trajs is B, S, N, 2
+        # rgbs is B, S, C, H, W
+        B, S, C, H, W = rgbs.shape
+        B, S2, N, D = trajs.shape
+        assert(S==S2)
+
+        rgbs = rgbs[0] # S, C, H, W
+        trajs = trajs[0] # S, N, 2
+        visibles = visibles[0] # S, N
+        if valids is None:
+            valids = torch.ones_like(trajs[:,:,0]) # S, N
+        else:
+            valids = valids[0]
+        
+        rgbs_color = []
+        for rgb in rgbs:
+            rgb = back2color(rgb).detach().cpu().numpy() 
+            rgb = np.transpose(rgb, [1, 2, 0]) # put channels last
+            rgbs_color.append(rgb) # each element 3 x H x W
+
+        trajs = trajs.long().detach().cpu().numpy() # S, N, 2
+        visibles = visibles.float().detach().cpu().numpy() # S, N
+        valids = valids.long().detach().cpu().numpy() # S, N
+
+        for i in range(min(N, max_show)):
+            if cmap=='onediff' and i==0:
+                cmap_ = 'spring'
+            elif cmap=='onediff':
+                cmap_ = 'winter'
+            else:
+                cmap_ = cmap
+            traj = trajs[:,i] # S,2
+            vis = visibles[:,i] # S
+            valid = valids[:,i] # S
+            rgbs_color = self.draw_traj_on_images_py(rgbs_color, traj, S=S, show_dots=show_dots, cmap=cmap_, linewidth=linewidth)
+            
+        for i in range(min(N, max_show)):
+            if cmap=='onediff' and i==0:
+                cmap_ = 'spring'
+            elif cmap=='onediff':
+                cmap_ = 'winter'
+            else:
+                cmap_ = cmap
+            traj = trajs[:,i] # S,2
+            vis = visibles[:,i] # S
+            valid = valids[:,i] # S
+            rgbs_color = self.draw_circ_on_images_py(rgbs_color, traj, vis, S=S, show_dots=show_dots, cmap=None, linewidth=linewidth)
+
+        rgbs = []
+        for rgb in rgbs_color:
+            rgb = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)
+            rgbs.append(preprocess_color(rgb))
+
+        return self.summ_rgbs(name, rgbs, only_return=only_return, frame_ids=frame_ids, frame_strs=frame_strs)
+
+    def summ_traj2ds_on_rgb(self, name, trajs, rgb, valids=None, show_dots=True, show_lines=True, frame_id=None, frame_str=None, only_return=False, cmap='coolwarm', linewidth=1, max_show=1024):
+        # trajs is B, S, N, 2
+        # rgb is B, C, H, W
+        B, C, H, W = rgb.shape
+        B, S, N, D = trajs.shape
+
+        rgb = rgb[0] # S, C, H, W
+        trajs = trajs[0] # S, N, 2
+        
+        if valids is None:
+            valids = torch.ones_like(trajs[:,:,0])
+        else:
+            valids = valids[0]
+
+        rgb_color = back2color(rgb).detach().cpu().numpy() 
+        rgb_color = np.transpose(rgb_color, [1, 2, 0]) # put channels last
+
+        # using maxdist will dampen the colors for short motions
+        # norms = torch.sqrt(1e-4 + torch.sum((trajs[-1] - trajs[0])**2, dim=1)) # N
+        # maxdist = torch.quantile(norms, 0.95).detach().cpu().numpy()
+        maxdist = None 
+        trajs = trajs.long().detach().cpu().numpy() # S, N, 2
+        valids = valids.long().detach().cpu().numpy() # S, N
+
+        if N > max_show:
+            inds = np.random.choice(N, max_show)
+            trajs = trajs[:,inds]
+            valids = valids[:,inds]
+            N = trajs.shape[1]
+        
+        for i in range(min(N, max_show)):
+            if cmap=='onediff' and i==0:
+                cmap_ = 'spring'
+            elif cmap=='onediff':
+                cmap_ = 'winter'
+            else:
+                cmap_ = cmap
+            traj = trajs[:,i] # S, 2
+            valid = valids[:,i] # S
+            if valid[0]==1:
+                traj = traj[valid>0]
+                rgb_color = self.draw_traj_on_image_py(
+                    rgb_color, traj, S=S, show_dots=show_dots, show_lines=show_lines, cmap=cmap_, maxdist=maxdist, linewidth=linewidth)
+
+        rgb_color = torch.from_numpy(rgb_color).permute(2, 0, 1).unsqueeze(0)
+        rgb = preprocess_color(rgb_color)
+        return self.summ_rgb(name, rgb, only_return=only_return, frame_id=frame_id, frame_str=frame_str)
+    
+    def draw_traj_on_image_py(self, rgb, traj, S=50, linewidth=1, show_dots=False, show_lines=True, cmap='coolwarm', val=None, maxdist=None):
+        # all inputs are numpy tensors
+        # rgb is 3 x H x W
+        # traj is S x 2
+        
+        H, W, C = rgb.shape
+        assert(C==3)
+
+        rgb = rgb.astype(np.uint8).copy()
+
+        S1, D = traj.shape
+        assert(D==2)
+
+        color_map = cm.get_cmap(cmap)
+        S1, D = traj.shape
+
+        for s in range(S1):
+            if val is not None:
+                color = np.array(color_map(val)[:3]) * 255 # rgb
+            else:
+                if maxdist is not None:
+                    val = (np.sqrt(np.sum((traj[s]-traj[0])**2))/maxdist).clip(0,1)
+                    color = np.array(color_map(val)[:3]) * 255 # rgb
+                else:
+                    color = np.array(color_map((s)/max(1,float(S-2)))[:3]) * 255 # rgb
+
+            if show_lines and s<(S1-1):
+                cv2.line(rgb,
+                         (int(traj[s,0]), int(traj[s,1])),
+                         (int(traj[s+1,0]), int(traj[s+1,1])),
+                         color,
+                         linewidth,
+                         cv2.LINE_AA)
+            if show_dots:
+                cv2.circle(rgb, (int(traj[s,0]), int(traj[s,1])), linewidth, color, -1)
+
+        # if maxdist is not None:
+        #     val = (np.sqrt(np.sum((traj[-1]-traj[0])**2))/maxdist).clip(0,1)
+        #     color = np.array(color_map(val)[:3]) * 255 # rgb
+        # else:
+        #     # draw the endpoint of traj, using the next color (which may be the last color)
+        #     color = np.array(color_map((S1-1)/max(1,float(S-2)))[:3]) * 255 # rgb
+
+        # # emphasize endpoint
+        # cv2.circle(rgb, (traj[-1,0], traj[-1,1]), linewidth*2, color, -1)
+
+        return rgb
+
+
+    def draw_traj_on_images_py(self, rgbs, traj, S=50, linewidth=1, show_dots=False, cmap='coolwarm', maxdist=None):
+        # all inputs are numpy tensors
+        # rgbs is a list of H,W,3
+        # traj is S,2
+        H, W, C = rgbs[0].shape
+        assert(C==3)
+
+        rgbs = [rgb.astype(np.uint8).copy() for rgb in rgbs]
+
+        S1, D = traj.shape
+        assert(D==2)
+        
+        x = int(np.clip(traj[0,0], 0, W-1))
+        y = int(np.clip(traj[0,1], 0, H-1))
+        color = rgbs[0][y,x]
+        color = (int(color[0]),int(color[1]),int(color[2]))
+        for s in range(S):
+            # bak_color = np.array(color_map(1.0)[:3]) * 255 # rgb
+            # cv2.circle(rgbs[s], (traj[s,0], traj[s,1]), linewidth*4, bak_color, -1)
+            cv2.polylines(rgbs[s],
+                          [traj[:s+1]],
+                          False,
+                          color,
+                          linewidth,
+                          cv2.LINE_AA)
+        return rgbs
+    
+    def draw_circs_on_image_py(self, rgb, xy, colors=None, linewidth=10, radius=3, show_dots=False, maxdist=None):
+        # all inputs are numpy tensors
+        # rgbs is a list of 3,H,W
+        # xy is N,2
+        H, W, C = rgb.shape
+        assert(C==3)
+
+        rgb = rgb.astype(np.uint8).copy()
+
+        N, D = xy.shape
+        assert(D==2)
+
+
+        xy = xy.astype(np.float32)
+        xy[:,0] = np.clip(xy[:,0], 0, W-1)
+        xy[:,1] = np.clip(xy[:,1], 0, H-1)
+        xy = xy.astype(np.int32)
+
+
+
+        if colors is None:
+            colors = get_n_colors(N)
+
+        for n in range(N):
+            color = colors[n]
+            # print('color', color)
+            # color = (color[0]*255).astype(np.uint8) 
+            color = (int(color[0]),int(color[1]),int(color[2]))
+
+            # x = int(np.clip(xy[0,0], 0, W-1))
+            # y = int(np.clip(xy[0,1], 0, H-1))
+            # color_ = rgbs[0][y,x]
+            # color_ = (int(color_[0]),int(color_[1]),int(color_[2]))
+            # color_ = (int(color_[0]),int(color_[1]),int(color_[2]))
+
+            cv2.circle(rgb, (int(xy[n,0]), int(xy[n,1])), linewidth, color, 3)
+            # vis_color = int(np.squeeze(vis[s])*255)
+            # vis_color = (vis_color,vis_color,vis_color)
+            # cv2.circle(rgbs[s], (traj[s,0], traj[s,1]), linewidth+1, vis_color, -1)
+        return rgb
+    
+    def draw_circ_on_images_py(self, rgbs, traj, vis, S=50, linewidth=1, show_dots=False, cmap=None, maxdist=None):
+        # all inputs are numpy tensors
+        # rgbs is a list of 3,H,W
+        # traj is S,2
+        H, W, C = rgbs[0].shape
+        assert(C==3)
+
+        rgbs = [rgb.astype(np.uint8).copy() for rgb in rgbs]
+
+        S1, D = traj.shape
+        assert(D==2)
+
+        if cmap is None:
+            bremm = ColorMap2d()
+            traj_ = traj[0:1].astype(np.float32)
+            traj_[:,0] /= float(W)
+            traj_[:,1] /= float(H)
+            color = bremm(traj_)
+            # print('color', color)
+            color = (color[0]*255).astype(np.uint8) 
+            color = (int(color[0]),int(color[1]),int(color[2]))
+
+        for s in range(S):
+            if cmap is not None:
+                color_map = cm.get_cmap(cmap)
+                # color = np.array(color_map(s/(S-1))[:3]) * 255 # rgb
+                color = np.array(color_map((s)/max(1,float(S-2)))[:3]) * 255 # rgb
+                # color = color.astype(np.uint8)
+                # color = (color[0], color[1], color[2])
+                # print('color', color)
+            # import ipdb; ipdb.set_trace()
+                
+            cv2.circle(rgbs[s], (int(traj[s,0]), int(traj[s,1])), linewidth+2, color, -1)
+            vis_color = int(np.squeeze(vis[s])*255)
+            vis_color = (vis_color,vis_color,vis_color)
+            cv2.circle(rgbs[s], (int(traj[s,0]), int(traj[s,1])), linewidth+1, vis_color, -1)
+                
+        return rgbs
+
+    def summ_traj_as_crops(self, name, trajs_e, rgbs, frame_id=None, frame_str=None, only_return=False, show_circ=True, trajs_g=None, valids_g=None, is_g=False, anchor_ind=None, ara=None, anchors=None, frame_ids=None, frame_strs=None):
+        B, S, N, D = trajs_e.shape
+        assert(N==1)
+        assert(D==2)
+
+        rgbs = back2color(rgbs).detach().cpu().byte().numpy()
+        
+        rgbs_vis = []
+        n = 0
+        pad_amount = 128
+        trajs_e_py = trajs_e[0].detach().cpu().numpy()
+        # trajs_e_py = np.clip(trajs_e_py, min=pad_amount/2, max=pad_amoun
+        trajs_e_py = trajs_e_py + pad_amount
+
+        if trajs_g is not None:
+            trajs_g_py = trajs_g[0].detach().cpu().numpy()
+            trajs_g_py = trajs_g_py + pad_amount
+
+        if valids_g is not None:
+            valids_g_py = valids_g[0].detach().cpu().numpy()
+        else:
+            valids_g_py = np.ones_like(trajs_g_py[:,:,:,0])
+            
+        
+        for s in range(S):
+            rgb = rgbs[0,s]
+            # print('orig rgb', rgb.shape)
+            rgb = np.transpose(rgb,(1,2,0)) # H, W, 3
+
+            rgb = np.pad(rgb, ((pad_amount,pad_amount),(pad_amount,pad_amount),(0,0)))
+            # print('pad rgb', rgb.shape)
+            H, W, C = rgb.shape
+
+            if trajs_g is not None:
+                xy_g = trajs_g_py[s,n]
+                xy_g[0] = np.clip(xy_g[0], pad_amount, W-pad_amount)
+                xy_g[1] = np.clip(xy_g[1], pad_amount, H-pad_amount)
+                if valids_g_py[s,n] > 0:
+                    rgb = self.draw_circs_on_image_py(rgb, xy_g.reshape(1,2), colors=[(0,255,0)], linewidth=2, radius=3)
+
+            xy_e = trajs_e_py[s,n]
+            xy_e[0] = np.clip(xy_e[0], pad_amount, W-pad_amount)
+            xy_e[1] = np.clip(xy_e[1], pad_amount, H-pad_amount)
+
+            if show_circ:
+
+                # if (anchors is not None) and (s==anchor_ind):
+                #     rgb = self.draw_circs_on_image_py(rgb, xy_e.reshape(1,2), colors=[(0,0,0)], linewidth=8, radius=12)
+                
+                if (anchors is not None) and (s in anchors):
+                    rgb = self.draw_circs_on_image_py(rgb, xy_e.reshape(1,2), colors=[(255,255,255)], linewidth=4, radius=8)
+                
+                if is_g:
+                    rgb = self.draw_circs_on_image_py(rgb, xy_e.reshape(1,2), colors=[(0,255,0)], linewidth=2, radius=3)
+                else:
+                    rgb = self.draw_circs_on_image_py(rgb, xy_e.reshape(1,2), colors=[(255,0,255)], linewidth=2, radius=3)
+
+
+                
+            xmin = int(xy_e[0])-pad_amount//2
+            xmax = int(xy_e[0])+pad_amount//2
+            ymin = int(xy_e[1])-pad_amount//2
+            ymax = int(xy_e[1])+pad_amount//2
+            
+            rgb_ = rgb[ymin:ymax, xmin:xmax]
+
+            H_, W_ = rgb_.shape[:2]
+            # if np.any(rgb_.shape==0):
+            #     input()
+            if H_==0 or W_==0:
+                import ipdb; ipdb.set_trace()
+
+            if (ara is not None) and (s in ara):
+                # green border
+                rgb_[0,:,0] = 0
+                rgb_[0,:,1] = 255
+                rgb_[0,:,2] = 0
+
+                rgb_[-1,:,0] = 0
+                rgb_[-1,:,1] = 255
+                rgb_[-1,:,2] = 0
+                
+                rgb_[:,0,0] = 0
+                rgb_[:,0,1] = 255
+                rgb_[:,0,2] = 0
+
+                rgb_[:,-1,0] = 0
+                rgb_[:,-1,1] = 255
+                rgb_[:,-1,2] = 0
+
+            if (anchor_ind is not None) and (s==anchor_ind):
+                # inner green border
+                pad = 4
+                
+                rgb_[:pad,:,0] = 0
+                rgb_[:pad,:,1] = 255
+                rgb_[:pad,:,2] = 0
+
+                rgb_[-pad:,:,0] = 0
+                rgb_[-pad:,:,1] = 255
+                rgb_[-pad:,:,2] = 0
+                
+                rgb_[:,:pad,0] = 0
+                rgb_[:,:pad,1] = 255
+                rgb_[:,:pad,2] = 0
+
+                rgb_[:,-pad:,0] = 0
+                rgb_[:,-pad:,1] = 255
+                rgb_[:,-pad:,2] = 0
+
+            rgb_ = rgb_.transpose(2,0,1)
+            rgb_ = torch.from_numpy(rgb_)
+
+            if frame_ids is not None:
+                # if s==anchor_ind:
+                #     frame_ids[s] = frame_ids[s] + '
+                rgb_ = draw_frame_id_on_vis(rgb_.unsqueeze(0), frame_ids[s]).squeeze(0)
+
+                if s==anchor_ind:
+                    rgb_ = draw_frame_str_on_vis(rgb_.unsqueeze(0), '(A)').squeeze(0)
+                    
+            if frame_strs is not None:
+                rgb_ = draw_frame_str_on_vis(rgb_.unsqueeze(0), frame_strs[s]).squeeze(0)
+
+            rgbs_vis.append(rgb_)
+
+        # nrow = int(np.sqrt(S)*(16.0/9)/2.0)
+        nrow = int(np.sqrt(S)*1.5)
+        grid_img = torchvision.utils.make_grid(torch.stack(rgbs_vis, dim=0), nrow=nrow).unsqueeze(0)
+        # print('grid_img', grid_img.shape)
+        return self.summ_rgb(name, grid_img.byte(), frame_id=frame_id, frame_str=frame_str, only_return=only_return)
+        
+    def summ_pts_on_rgb(self, name, trajs, rgb, visibs=None, valids=None, frame_id=None, frame_str=None, only_return=False, show_dots=True, colors=None, cmap='coolwarm', linewidth=1, max_show=1024, already_sorted=False):
+        # trajs is B, S, N, 2
+        # rgbs is B, S, C, H, W
+        B, C, H, W = rgb.shape
+        B, S, N, D = trajs.shape
+
+        rgb = rgb[0] # C, H, W
+        trajs = trajs[0] # S, N, 2
+        if valids is None:
+            valids = torch.ones_like(trajs[:,:,0]) # S, N
+        else:
+            valids = valids[0]
+        if visibs is None:
+            visibs = torch.ones_like(trajs[:,:,0]) # S, N
+        else:
+            visibs = visibs[0]
+
+        trajs = trajs.clamp(-16, W+16)
+        
+        if N > max_show:
+            inds = np.random.choice(N, max_show)
+            trajs = trajs[:,inds]
+            valids = valids[:,inds]
+            visibs = visibs[:,inds]
+            N = trajs.shape[1]
+
+        if not already_sorted:
+            inds = torch.argsort(torch.mean(trajs[:,:,1], dim=0))
+            trajs = trajs[:,inds]
+            valids = valids[:,inds]
+            visibs = visibs[:,inds]
+        
+        rgb = back2color(rgb).detach().cpu().numpy() 
+        rgb = np.transpose(rgb, [1, 2, 0]) # put channels last
+
+        trajs = trajs.long().detach().cpu().numpy() # S, N, 2
+        valids = valids.long().detach().cpu().numpy() # S, N
+        visibs = visibs.long().detach().cpu().numpy() # S, N
+
+        rgb = rgb.astype(np.uint8).copy()
+        
+        for i in range(min(N, max_show)):
+            if cmap=='onediff' and i==0:
+                cmap_ = 'spring'
+            elif cmap=='onediff':
+                cmap_ = 'winter'
+            else:
+                cmap_ = cmap
+            traj = trajs[:,i] # S,2
+            valid = valids[:,i] # S
+            visib = visibs[:,i] # S
+
+            if colors is None:
+                ii = i/(1e-4+N-1.0)
+                color_map = cm.get_cmap(cmap)
+                color = np.array(color_map(ii)[:3]) * 255 # rgb
+            else:
+                color = np.array(colors[i]).astype(np.int64)
+            color = (int(color[0]),int(color[1]),int(color[2]))
+
+            for s in range(S):
+                if valid[s]:
+                    if visib[s]:
+                        thickness = -1
+                    else:
+                        thickness = 2
+                    cv2.circle(rgb, (int(traj[s,0]), int(traj[s,1])), linewidth, color, thickness)
+        rgb = torch.from_numpy(rgb).permute(2,0,1).unsqueeze(0)
+        rgb = preprocess_color(rgb)
+        return self.summ_rgb(name, rgb, only_return=only_return, frame_id=frame_id, frame_str=frame_str)
+
+    def summ_pts_on_rgbs(self, name, trajs, rgbs, visibs=None, valids=None, frame_ids=None, only_return=False, show_dots=True, cmap='coolwarm', colors=None, linewidth=1, max_show=1024, frame_strs=None):
+        # trajs is B, S, N, 2
+        # rgbs is B, S, C, H, W
+        B, S, C, H, W = rgbs.shape
+        B, S2, N, D = trajs.shape
+        assert(S==S2)
+
+        rgbs = rgbs[0] # S, C, H, W
+        trajs = trajs[0] # S, N, 2
+        if valids is None:
+            valids = torch.ones_like(trajs[:,:,0]) # S, N
+        else:
+            valids = valids[0]
+        if visibs is None:
+            visibs = torch.ones_like(trajs[:,:,0]) # S, N
+        else:
+            visibs = visibs[0]
+
+        if N > max_show:
+            inds = np.random.choice(N, max_show)
+            trajs = trajs[:,inds]
+            valids = valids[:,inds]
+            visibs = visibs[:,inds]
+            N = trajs.shape[1]
+        inds = torch.argsort(torch.mean(trajs[:,:,1], dim=0))
+        trajs = trajs[:,inds]
+        valids = valids[:,inds]
+        visibs = visibs[:,inds]
+        
+        rgbs_color = []
+        for rgb in rgbs:
+            rgb = back2color(rgb).detach().cpu().numpy() 
+            rgb = np.transpose(rgb, [1, 2, 0]) # put channels last
+            rgbs_color.append(rgb) # each element 3 x H x W
+
+        trajs = trajs.long().detach().cpu().numpy() # S, N, 2
+        valids = valids.long().detach().cpu().numpy() # S, N
+        visibs = visibs.long().detach().cpu().numpy() # S, N
+
+        rgbs_color = [rgb.astype(np.uint8).copy() for rgb in rgbs_color]
+        
+        for i in range(min(N, max_show)):
+            traj = trajs[:,i] # S,2
+            valid = valids[:,i] # S
+            visib = visibs[:,i] # S
+            
+            if colors is None:
+                ii = i/(1e-4+N-1.0)
+                color_map = cm.get_cmap(cmap)
+                color = np.array(color_map(ii)[:3]) * 255 # rgb
+            else:
+                color = np.array(colors[i]).astype(np.int64)
+            color = (int(color[0]),int(color[1]),int(color[2]))
+
+            for s in range(S):
+                if valid[s]:
+                    if visib[s]:
+                        thickness = -1
+                    else:
+                        thickness = 2
+                    cv2.circle(rgbs_color[s], (int(traj[s,0]), int(traj[s,1])), int(linewidth), color, thickness)
+        rgbs = []
+        for rgb in rgbs_color:
+            rgb = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)
+            rgbs.append(preprocess_color(rgb))
+
+        return self.summ_rgbs(name, rgbs, only_return=only_return, frame_ids=frame_ids, frame_strs=frame_strs)
+    
+    
+def erode2d(im, times=1):
+    device = im.device
+    weights2d = torch.ones(1, 1, 3, 3, dtype=im.dtype, device=device)
+    for time in range(times):
+        im = 1.0 - F.conv2d(1.0 - im, weights2d, padding=1).clamp(0, 1)
+    return im
+
+def dilate2d(im, times=1):
+    device = im.device
+    assert(times>0)
+    dilation_kernel_size = times*2 + 1
+    padding_size = dilation_kernel_size // 2
+    dilation_kernel = torch.ones((1, 1, dilation_kernel_size, dilation_kernel_size), device=device)
+    im = F.conv2d(im, dilation_kernel, padding=padding_size, groups=im.shape[1]).clamp(0,1)
+    return im
+
+

utils/loss.py

@@ -0,0 +1,220 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from typing import List
+import utils.basic
+
+
+def sequence_loss(
+    flow_preds,
+    flow_gt,
+    valids,
+    vis=None,
+    gamma=0.8,
+    use_huber_loss=False,
+    loss_only_for_visible=False,
+):
+    """Loss function defined over sequence of flow predictions"""
+    total_flow_loss = 0.0
+    for j in range(len(flow_gt)):
+        B, S, N, D = flow_gt[j].shape
+        B, S2, N = valids[j].shape
+        assert S == S2
+        n_predictions = len(flow_preds[j])
+        flow_loss = 0.0
+        for i in range(n_predictions):
+            i_weight = gamma ** (n_predictions - i - 1)
+            flow_pred = flow_preds[j][i]
+            if use_huber_loss:
+                i_loss = huber_loss(flow_pred, flow_gt[j], delta=6.0)
+            else:
+                i_loss = (flow_pred - flow_gt[j]).abs()  # B, S, N, 2
+            i_loss = torch.mean(i_loss, dim=3)  # B, S, N
+            valid_ = valids[j].clone()
+            if loss_only_for_visible:
+                valid_ = valid_ * vis[j]
+            flow_loss += i_weight * utils.basic.reduce_masked_mean(i_loss, valid_)
+        flow_loss = flow_loss / n_predictions
+        total_flow_loss += flow_loss
+    return total_flow_loss / len(flow_gt)
+
+def sequence_loss_dense(
+    flow_preds,
+    flow_gt,
+    valids,
+    vis=None,
+    gamma=0.8,
+    use_huber_loss=False,
+    loss_only_for_visible=False,
+):
+    """Loss function defined over sequence of flow predictions"""
+    total_flow_loss = 0.0
+    for j in range(len(flow_gt)):
+        # print('flow_gt[j]', flow_gt[j].shape)
+        B, S, D, H, W = flow_gt[j].shape
+        B, S2, _, H, W = valids[j].shape
+        assert S == S2
+        n_predictions = len(flow_preds[j])
+        flow_loss = 0.0
+        # import ipdb; ipdb.set_trace()
+        for i in range(n_predictions):
+            # print('flow_e[j][i]', flow_preds[j][i].shape)
+            i_weight = gamma ** (n_predictions - i - 1)
+            flow_pred = flow_preds[j][i] # B,S,2,H,W
+            if use_huber_loss:
+                i_loss = huber_loss(flow_pred, flow_gt[j], delta=6.0) # B,S,2,H,W
+            else:
+                i_loss = (flow_pred - flow_gt[j]).abs() # B,S,2,H,W
+            i_loss_ = torch.mean(i_loss, dim=2) # B,S,H,W
+            valid_ = valids[j].reshape(B,S,H,W)
+            # print(' (%d,%d) i_loss_' % (i,j), i_loss_.shape)
+            # print(' (%d,%d) valid_' % (i,j), valid_.shape)
+            if loss_only_for_visible:
+                valid_ = valid_ * vis[j].reshape(B,-1,H,W) # usually B,S,H,W, but maybe B,1,H,W
+            flow_loss += i_weight * utils.basic.reduce_masked_mean(i_loss_, valid_, broadcast=True)
+            # import ipdb; ipdb.set_trace()
+        flow_loss = flow_loss / n_predictions
+        total_flow_loss += flow_loss
+    return total_flow_loss / len(flow_gt)
+
+
+def huber_loss(x, y, delta=1.0):
+    """Calculate element-wise Huber loss between x and y"""
+    diff = x - y
+    abs_diff = diff.abs()
+    flag = (abs_diff <= delta).float()
+    return flag * 0.5 * diff**2 + (1 - flag) * delta * (abs_diff - 0.5 * delta)
+
+
+def sequence_BCE_loss(vis_preds, vis_gts, valids=None, use_logits=False):
+    total_bce_loss = 0.0
+    # all_vis_preds = [torch.stack(vp) for vp in vis_preds]
+    # all_vis_preds = torch.stack(all_vis_preds)
+    # utils.basic.print_stats('all_vis_preds', all_vis_preds)
+    for j in range(len(vis_preds)):
+        n_predictions = len(vis_preds[j])
+        bce_loss = 0.0
+        for i in range(n_predictions):
+            # utils.basic.print_stats('vis_preds[%d][%d]' % (j,i), vis_preds[j][i])
+            # utils.basic.print_stats('vis_gts[%d]' % (i), vis_gts[i])
+            if use_logits:
+                loss = F.binary_cross_entropy_with_logits(vis_preds[j][i], vis_gts[j], reduction='none')
+            else:
+                loss = F.binary_cross_entropy(vis_preds[j][i], vis_gts[j], reduction='none')
+            if valids is None:
+                bce_loss += loss.mean()
+            else:
+                bce_loss += (loss * valids[j]).mean()
+        bce_loss = bce_loss / n_predictions
+        total_bce_loss += bce_loss
+    return total_bce_loss / len(vis_preds)
+
+
+# def sequence_BCE_loss_dense(vis_preds, vis_gts):
+#     total_bce_loss = 0.0
+#     for j in range(len(vis_preds)):
+#         n_predictions = len(vis_preds[j])
+#         bce_loss = 0.0
+#         for i in range(n_predictions):
+#             vis_e = vis_preds[j][i]
+#             vis_g = vis_gts[j]
+#             print('vis_e', vis_e.shape, 'vis_g', vis_g.shape)
+#             vis_loss = F.binary_cross_entropy(vis_e, vis_g)
+#             bce_loss += vis_loss
+#         bce_loss = bce_loss / n_predictions
+#         total_bce_loss += bce_loss
+#     return total_bce_loss / len(vis_preds)
+
+
+def sequence_prob_loss(
+        tracks: torch.Tensor,
+        confidence: torch.Tensor,
+        target_points: torch.Tensor,
+        visibility: torch.Tensor,
+        expected_dist_thresh: float = 12.0,
+        use_logits=False,
+):
+    """Loss for classifying if a point is within pixel threshold of its target."""
+    # Points with an error larger than 12 pixels are likely to be useless; marking
+    # them as occluded will actually improve Jaccard metrics and give
+    # qualitatively better results.
+    total_logprob_loss = 0.0
+    for j in range(len(tracks)):
+        n_predictions = len(tracks[j])
+        logprob_loss = 0.0
+        for i in range(n_predictions):
+            err = torch.sum((tracks[j][i].detach() - target_points[j]) ** 2, dim=-1)
+            valid = (err <= expected_dist_thresh**2).float()
+            if use_logits:
+                loss = F.binary_cross_entropy_with_logits(confidence[j][i], valid, reduction="none")
+            else:
+                loss = F.binary_cross_entropy(confidence[j][i], valid, reduction="none")
+            loss *= visibility[j]
+            loss = torch.mean(loss, dim=[1, 2])
+            logprob_loss += loss
+        logprob_loss = logprob_loss / n_predictions
+        total_logprob_loss += logprob_loss
+    return total_logprob_loss / len(tracks)
+
+def sequence_prob_loss_dense(
+    tracks: torch.Tensor,
+    confidence: torch.Tensor,
+    target_points: torch.Tensor,
+    visibility: torch.Tensor,
+    expected_dist_thresh: float = 12.0,
+    use_logits=False,
+):
+    """Loss for classifying if a point is within pixel threshold of its target."""
+    # Points with an error larger than 12 pixels are likely to be useless; marking
+    # them as occluded will actually improve Jaccard metrics and give
+    # qualitatively better results.
+
+    # all_confidence = [torch.stack(vp) for vp in confidence]
+    # all_confidence = torch.stack(all_confidence)
+    # utils.basic.print_stats('all_confidence', all_confidence)
+    
+    total_logprob_loss = 0.0
+    for j in range(len(tracks)):
+        n_predictions = len(tracks[j])
+        logprob_loss = 0.0
+        for i in range(n_predictions):
+            # print('trajs_e', tracks[j][i].shape)
+            # print('trajs_g', target_points[j].shape)
+            err = torch.sum((tracks[j][i].detach() - target_points[j]) ** 2, dim=2)
+            positive = (err <= expected_dist_thresh**2).float()
+            # print('conf', confidence[j][i].shape, 'positive', positive.shape)
+            if use_logits:
+                loss = F.binary_cross_entropy_with_logits(confidence[j][i].squeeze(2), positive, reduction="none")
+            else:
+                loss = F.binary_cross_entropy(confidence[j][i].squeeze(2), positive, reduction="none")
+            loss *= visibility[j].squeeze(2) # B,S,H,W
+            loss = torch.mean(loss, dim=[1,2,3])
+            logprob_loss += loss
+        logprob_loss = logprob_loss / n_predictions
+        total_logprob_loss += logprob_loss
+    return total_logprob_loss / len(tracks)
+
+
+def masked_mean(data, mask, dim):
+    if mask is None:
+        return data.mean(dim=dim, keepdim=True)
+    mask = mask.float()
+    mask_sum = torch.sum(mask, dim=dim, keepdim=True)
+    mask_mean = torch.sum(data * mask, dim=dim, keepdim=True) / torch.clamp(
+        mask_sum, min=1.0
+    )
+    return mask_mean
+
+
+def masked_mean_var(data: torch.Tensor, mask: torch.Tensor, dim: List[int]):
+    if mask is None:
+        return data.mean(dim=dim, keepdim=True), data.var(dim=dim, keepdim=True)
+    mask = mask.float()
+    mask_sum = torch.sum(mask, dim=dim, keepdim=True)
+    mask_mean = torch.sum(data * mask, dim=dim, keepdim=True) / torch.clamp(
+        mask_sum, min=1.0
+    )
+    mask_var = torch.sum(
+        mask * (data - mask_mean) ** 2, dim=dim, keepdim=True
+    ) / torch.clamp(mask_sum, min=1.0)
+    return mask_mean.squeeze(dim), mask_var.squeeze(dim)

utils/metric.py

@@ -0,0 +1,176 @@
+import math
+import numpy as np
+import cv2
+
+def _seg2bmap(seg, width=None, height=None):
+    """
+    From a segmentation, compute a binary boundary map with 1 pixel wide
+    boundaries.  The boundary pixels are offset by 1/2 pixel towards the
+    origin from the actual segment boundary.
+    Arguments:
+        seg     : Segments labeled from 1..k.
+        width	  :	Width of desired bmap  <= seg.shape[1]
+        height  :	Height of desired bmap <= seg.shape[0]
+    Returns:
+        bmap (ndarray):	Binary boundary map.
+     David Martin <[email protected]>
+     January 2003
+    """
+
+    seg = seg.astype(bool)
+    seg[seg > 0] = 1
+
+    assert np.atleast_3d(seg).shape[2] == 1
+
+    width = seg.shape[1] if width is None else width
+    height = seg.shape[0] if height is None else height
+
+    h, w = seg.shape[:2]
+
+    ar1 = float(width) / float(height)
+    ar2 = float(w) / float(h)
+
+    assert not (width > w | height > h | abs(ar1 - ar2) >
+                0.01), "Can" "t convert %dx%d seg to %dx%d bmap." % (w, h, width, height)
+
+    e = np.zeros_like(seg)
+    s = np.zeros_like(seg)
+    se = np.zeros_like(seg)
+
+    e[:, :-1] = seg[:, 1:]
+    s[:-1, :] = seg[1:, :]
+    se[:-1, :-1] = seg[1:, 1:]
+
+    b = seg ^ e | seg ^ s | seg ^ se
+    b[-1, :] = seg[-1, :] ^ e[-1, :]
+    b[:, -1] = seg[:, -1] ^ s[:, -1]
+    b[-1, -1] = 0
+
+    if w == width and h == height:
+        bmap = b
+    else:
+        bmap = np.zeros((height, width))
+        for x in range(w):
+            for y in range(h):
+                if b[y, x]:
+                    j = 1 + math.floor((y - 1) + height / h)
+                    i = 1 + math.floor((x - 1) + width / h)
+                    bmap[j, i] = 1
+
+    return bmap
+
+# https://github.com/davisvideochallenge/davis2017-evaluation/blob/master/davis2017/metrics.py#L6
+def db_eval_iou(annotation, segmentation, void_pixels=None):
+    """ Compute region similarity as the Jaccard Index.
+    Arguments:
+        annotation   (ndarray): binary annotation   map.
+        segmentation (ndarray): binary segmentation map.
+        void_pixels  (ndarray): optional mask with void pixels
+
+    Return:
+        jaccard (float): region similarity
+    """
+    assert annotation.shape == segmentation.shape, \
+        f'Annotation({annotation.shape}) and segmentation:{segmentation.shape} dimensions do not match.'
+    annotation = annotation.astype(bool)
+    segmentation = segmentation.astype(bool)
+
+    if void_pixels is not None:
+        assert annotation.shape == void_pixels.shape, \
+            f'Annotation({annotation.shape}) and void pixels:{void_pixels.shape} dimensions do not match.'
+        void_pixels = void_pixels.astype(bool)
+    else:
+        void_pixels = np.zeros_like(segmentation)
+
+    # Intersection between all sets
+    inters = np.sum((segmentation & annotation) & np.logical_not(void_pixels), axis=(-2, -1))
+    union = np.sum((segmentation | annotation) & np.logical_not(void_pixels), axis=(-2, -1))
+
+    j = inters / union
+    if j.ndim == 0:
+        j = 0 if np.isclose(union, 0) else j
+    else:
+        j[np.isclose(union, 0)] = 0
+    return j
+
+
+def db_eval_boundary(annotation, segmentation, void_pixels=None, bound_th=0.008):
+    assert annotation.shape == segmentation.shape
+    if void_pixels is not None:
+        assert annotation.shape == void_pixels.shape
+    if annotation.ndim == 3:
+        n_frames = annotation.shape[0]
+        f_res = np.zeros(n_frames)
+        for frame_id in range(n_frames):
+            void_pixels_frame = None if void_pixels is None else void_pixels[frame_id, :, :, ]
+            f_res[frame_id] = f_measure(segmentation[frame_id, :, :, ], annotation[frame_id, :, :], void_pixels_frame, bound_th=bound_th)
+    elif annotation.ndim == 2:
+        f_res = f_measure(segmentation, annotation, void_pixels, bound_th=bound_th)
+    else:
+        raise ValueError(f'db_eval_boundary does not support tensors with {annotation.ndim} dimensions')
+    return f_res
+
+
+def f_measure(foreground_mask, gt_mask, void_pixels=None, bound_th=0.008):
+    """
+    Compute mean,recall and decay from per-frame evaluation.
+    Calculates precision/recall for boundaries between foreground_mask and
+    gt_mask using morphological operators to speed it up.
+
+    Arguments:
+        foreground_mask (ndarray): binary segmentation image.
+        gt_mask         (ndarray): binary annotated image.
+        void_pixels     (ndarray): optional mask with void pixels
+
+    Returns:
+        F (float): boundaries F-measure
+    """
+    assert np.atleast_3d(foreground_mask).shape[2] == 1
+    if void_pixels is not None:
+        void_pixels = void_pixels.astype(bool)
+    else:
+        void_pixels = np.zeros_like(foreground_mask).astype(bool)
+
+    bound_pix = bound_th if bound_th >= 1 else \
+        np.ceil(bound_th * np.linalg.norm(foreground_mask.shape))
+
+    # Get the pixel boundaries of both masks
+    fg_boundary = _seg2bmap(foreground_mask * np.logical_not(void_pixels))
+    gt_boundary = _seg2bmap(gt_mask * np.logical_not(void_pixels))
+
+    from skimage.morphology import disk
+
+    # fg_dil = binary_dilation(fg_boundary, disk(bound_pix))
+    fg_dil = cv2.dilate(fg_boundary.astype(np.uint8), disk(bound_pix).astype(np.uint8))
+    # gt_dil = binary_dilation(gt_boundary, disk(bound_pix))
+    gt_dil = cv2.dilate(gt_boundary.astype(np.uint8), disk(bound_pix).astype(np.uint8))
+
+    # Get the intersection
+    gt_match = gt_boundary * fg_dil
+    fg_match = fg_boundary * gt_dil
+
+    # Area of the intersection
+    n_fg = np.sum(fg_boundary)
+    n_gt = np.sum(gt_boundary)
+
+    # % Compute precision and recall
+    if n_fg == 0 and n_gt > 0:
+        precision = 1
+        recall = 0
+    elif n_fg > 0 and n_gt == 0:
+        precision = 0
+        recall = 1
+    elif n_fg == 0 and n_gt == 0:
+        precision = 1
+        recall = 1
+    else:
+        precision = np.sum(fg_match) / float(n_fg)
+        recall = np.sum(gt_match) / float(n_gt)
+
+    # Compute F measure
+    if precision + recall == 0:
+        F = 0
+    else:
+        F = 2 * precision * recall / (precision + recall)
+
+    return F

utils/misc.py

@@ -0,0 +1,1062 @@
+import torch
+import numpy as np
+import math
+import torch.nn.functional as F
+import utils.basic
+from typing import Tuple, Union
+
+def standardize_test_data(rgbs, trajs, visibs, valids, S_cap=600, only_first=False, seq_len=None):
+    trajs = trajs.astype(np.float32) # S,N,2
+    visibs = visibs.astype(np.float32) # S,N
+    valids = valids.astype(np.float32) # S,N
+    
+    visval_ok = np.sum(valids*visibs, axis=0) > 1
+    trajs = trajs[:,visval_ok]
+    visibs = visibs[:,visval_ok]
+    valids = valids[:,visval_ok]
+
+    # fill in missing data
+    N = trajs.shape[1]
+    for ni in range(N):
+        trajs[:,ni] = utils.misc.data_replace_with_nearest(trajs[:,ni], valids[:,ni])
+
+    # use cap or seq_len
+    if seq_len is not None:
+        S = min(len(rgbs), seq_len)
+    else:
+        S = len(rgbs)
+    S = min(S, S_cap)
+
+    if only_first:
+        # we'll find the best frame to start on
+        best_count = 0
+        best_ind = 0
+
+        for si in range(0,len(rgbs)-64):
+            # try this slice
+            visibs_ = visibs[si:min(si+S,len(rgbs)+1)] # S,N
+            valids_ = valids[si:min(si+S,len(rgbs)+1)] # S,N
+            visval_ok0 = (visibs_[0]*valids_[0]) > 0 # N
+            visval_okA = np.sum(visibs_*valids_, axis=0) > 1 # N
+            all_ok = visval_ok0 & visval_okA
+            # print('- slicing %d to %d; sum(ok) %d' % (si, min(si+S,len(rgbs)+1), np.sum(all_ok)))
+            if np.sum(all_ok) > best_count:
+                best_count = np.sum(all_ok)
+                best_ind = si
+        si = best_ind
+        rgbs = rgbs[si:si+S]
+        trajs = trajs[si:si+S]
+        visibs = visibs[si:si+S]
+        valids = valids[si:si+S]
+        vis_ok0 = visibs[0] > 0  # N
+        trajs = trajs[:,vis_ok0]
+        visibs = visibs[:,vis_ok0]
+        valids = valids[:,vis_ok0]
+        # print('- best_count', best_count, 'best_ind', best_ind)
+
+    if seq_len is not None:
+        rgbs = rgbs[:seq_len]
+        trajs = trajs[:seq_len]
+        valids = valids[:seq_len]
+    
+    # req two timesteps valid (after seqlen trim)
+    visval_ok = np.sum(visibs*valids, axis=0) > 1
+    trajs = trajs[:,visval_ok]
+    valids = valids[:,visval_ok]
+    visibs = visibs[:,visval_ok]
+
+    return rgbs, trajs, visibs, valids
+
+    
+
+def get_2d_sincos_pos_embed(
+    embed_dim: int, grid_size: Union[int, Tuple[int, int]]
+) -> torch.Tensor:
+    """
+    This function initializes a grid and generates a 2D positional embedding using sine and cosine functions.
+    It is a wrapper of get_2d_sincos_pos_embed_from_grid.
+    Args:
+    - embed_dim: The embedding dimension.
+    - grid_size: The grid size.
+    Returns:
+    - pos_embed: The generated 2D positional embedding.
+    """
+    if isinstance(grid_size, tuple):
+        grid_size_h, grid_size_w = grid_size
+    else:
+        grid_size_h = grid_size_w = grid_size
+    grid_h = torch.arange(grid_size_h, dtype=torch.float)
+    grid_w = torch.arange(grid_size_w, dtype=torch.float)
+    grid = torch.meshgrid(grid_w, grid_h, indexing="xy")
+    grid = torch.stack(grid, dim=0)
+    grid = grid.reshape([2, 1, grid_size_h, grid_size_w])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    return pos_embed.reshape(1, grid_size_h, grid_size_w, -1).permute(0, 3, 1, 2)
+
+
+def get_2d_sincos_pos_embed_from_grid(
+    embed_dim: int, grid: torch.Tensor
+) -> torch.Tensor:
+    """
+    This function generates a 2D positional embedding from a given grid using sine and cosine functions.
+
+    Args:
+    - embed_dim: The embedding dimension.
+    - grid: The grid to generate the embedding from.
+
+    Returns:
+    - emb: The generated 2D positional embedding.
+    """
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = torch.cat([emb_h, emb_w], dim=2)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(
+    embed_dim: int, pos: torch.Tensor
+) -> torch.Tensor:
+    """
+    This function generates a 1D positional embedding from a given grid using sine and cosine functions.
+
+    Args:
+    - embed_dim: The embedding dimension.
+    - pos: The position to generate the embedding from.
+
+    Returns:
+    - emb: The generated 1D positional embedding.
+    """
+    assert embed_dim % 2 == 0
+    omega = torch.arange(embed_dim // 2, dtype=torch.double)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = torch.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = torch.sin(out)  # (M, D/2)
+    emb_cos = torch.cos(out)  # (M, D/2)
+
+    emb = torch.cat([emb_sin, emb_cos], dim=1)  # (M, D)
+    return emb[None].float()
+
+
+def get_2d_embedding(xy: torch.Tensor, C: int, cat_coords: bool = True) -> torch.Tensor:
+    """
+    This function generates a 2D positional embedding from given coordinates using sine and cosine functions.
+
+    Args:
+    - xy: The coordinates to generate the embedding from.
+    - C: The size of the embedding.
+    - cat_coords: A flag to indicate whether to concatenate the original coordinates to the embedding.
+
+    Returns:
+    - pe: The generated 2D positional embedding.
+    """
+    B, N, D = xy.shape
+    assert D == 2
+
+    x = xy[:, :, 0:1]
+    y = xy[:, :, 1:2]
+    div_term = (
+        torch.arange(0, C, 2, device=xy.device, dtype=torch.float32) * (1000.0 / C)
+    ).reshape(1, 1, int(C / 2))
+
+    pe_x = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+    pe_y = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+
+    pe_x[:, :, 0::2] = torch.sin(x * div_term)
+    pe_x[:, :, 1::2] = torch.cos(x * div_term)
+
+    pe_y[:, :, 0::2] = torch.sin(y * div_term)
+    pe_y[:, :, 1::2] = torch.cos(y * div_term)
+
+    pe = torch.cat([pe_x, pe_y], dim=2)  # (B, N, C*3)
+    if cat_coords:
+        pe = torch.cat([xy, pe], dim=2)  # (B, N, C*3+3)
+    return pe
+
+# from datasets.dataset import mask2bbox
+
+# from pips2
+def posemb_sincos_2d_xy(xy, C, temperature=10000, dtype=torch.float32, cat_coords=False):
+    device = xy.device
+    dtype = xy.dtype
+    B, S, D = xy.shape
+    assert(D==2)
+    x = xy[:,:,0]
+    y = xy[:,:,1]
+    assert (C % 4) == 0, 'feature dimension must be multiple of 4 for sincos emb'
+    omega = torch.arange(C // 4, device=device) / (C // 4 - 1)
+    omega = 1. / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :] 
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1) 
+    pe = pe.reshape(B,S,C).type(dtype)
+    if cat_coords:
+        pe = torch.cat([pe, xy], dim=2) # B,N,C+2
+    return pe
+
+
+# # prevent circular imports
+# def mask2bbox(mask):
+#     if mask.ndim == 3:
+#         mask = mask[..., 0]
+#     ys, xs = np.where(mask > 0.4)
+#     if ys.size == 0 or xs.size==0:
+#         return np.array((0, 0, 0, 0), dtype=int)
+#     lt = np.array([np.min(xs), np.min(ys)])
+#     rb = np.array([np.max(xs), np.max(ys)]) + 1
+#     return np.concatenate([lt, rb])
+
+# def get_stark_2d_embedding(H, W, C=64, device='cuda:0', temperature=10000, normalize=True):
+#     scale = 2*math.pi
+#     mask = torch.ones((1,H,W), dtype=torch.float32, device=device)
+#     y_embed = mask.cumsum(1, dtype=torch.float32)  # cumulative sum along axis 1 (h axis) --> (b, h, w)
+#     x_embed = mask.cumsum(2, dtype=torch.float32)  # cumulative sum along axis 2 (w axis) --> (b, h, w)
+#     if normalize:
+#         eps = 1e-6
+#         y_embed = y_embed / (y_embed[:, -1:, :] + eps) * scale  # 2pi * (y / sigma(y))
+#         x_embed = x_embed / (x_embed[:, :, -1:] + eps) * scale  # 2pi * (x / sigma(x))
+
+#     dim_t = torch.arange(C, dtype=torch.float32, device=device)  # (0,1,2,...,d/2)
+#     dim_t = temperature ** (2 * (dim_t // 2) / C)
+
+#     pos_x = x_embed[:, :, :, None] / dim_t # (b,h,w,d/2)
+#     pos_y = y_embed[:, :, :, None] / dim_t # (b,h,w,d/2)
+#     pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) # (b,h,w,d/2)
+#     pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) # (b,h,w,d/2)
+#     pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)  # (b,h,w,d)
+#     return pos
+
+# def get_1d_embedding(x, C, cat_coords=False):
+#     B, N, D = x.shape
+#     assert(D==1)
+
+#     div_term = (torch.arange(0, C, 2, device=x.device, dtype=torch.float32) * (10000.0 / C)).reshape(1, 1, int(C/2))
+    
+#     pe_x = torch.zeros(B, N, C, device=x.device, dtype=torch.float32)
+    
+#     pe_x[:, :, 0::2] = torch.sin(x * div_term)
+#     pe_x[:, :, 1::2] = torch.cos(x * div_term)
+    
+#     if cat_coords:
+#         pe_x = torch.cat([pe, x], dim=2) # B,N,C*2+2
+#     return pe_x
+
+# def posemb_sincos_2d(h, w, dim, temperature=10000, dtype=torch.float32, device='cuda:0'):
+
+#     y, x = torch.meshgrid(torch.arange(h, device = device), torch.arange(w, device = device), indexing = 'ij')
+#     assert (dim % 4) == 0, 'feature dimension must be multiple of 4 for sincos emb'
+#     omega = torch.arange(dim // 4, device = device) / (dim // 4 - 1)
+#     omega = 1. / (temperature ** omega)
+
+#     y = y.flatten()[:, None] * omega[None, :]
+#     x = x.flatten()[:, None] * omega[None, :] 
+#     pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1) # B,C,H,W
+#     return pe.type(dtype)
+
+def iou(bbox1, bbox2):
+    # bbox1, bbox2: [x1, y1, x2, y2]
+    x1, y1, x2, y2 = bbox1
+    x1_, y1_, x2_, y2_ = bbox2
+    inter_x1 = max(x1, x1_)
+    inter_y1 = max(y1, y1_)
+    inter_x2 = min(x2, x2_)
+    inter_y2 = min(y2, y2_)
+    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
+    area1 = (x2 - x1) * (y2 - y1)
+    area2 = (x2_ - x1_) * (y2_ - y1_)
+    iou = inter_area / (area1 + area2 - inter_area)
+    return iou
+
+# def get_2d_embedding(xy, C, cat_coords=False):
+#     B, N, D = xy.shape
+#     assert(D==2)
+
+#     x = xy[:,:,0:1]
+#     y = xy[:,:,1:2]
+#     div_term = (torch.arange(0, C, 2, device=xy.device, dtype=torch.float32) * (10000.0 / C)).reshape(1, 1, int(C/2))
+    
+#     pe_x = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+#     pe_y = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+    
+#     pe_x[:, :, 0::2] = torch.sin(x * div_term)
+#     pe_x[:, :, 1::2] = torch.cos(x * div_term)
+    
+#     pe_y[:, :, 0::2] = torch.sin(y * div_term)
+#     pe_y[:, :, 1::2] = torch.cos(y * div_term)
+    
+#     pe = torch.cat([pe_x, pe_y], dim=2) # B,N,C*2
+#     if cat_coords:
+#         pe = torch.cat([pe, xy], dim=2) # B,N,C*2+2
+#     return pe
+
+# def get_3d_embedding(xyz, C, cat_coords=False):
+#     B, N, D = xyz.shape
+#     assert(D==3)
+
+#     x = xyz[:,:,0:1]
+#     y = xyz[:,:,1:2]
+#     z = xyz[:,:,2:3]
+#     div_term = (torch.arange(0, C, 2, device=xyz.device, dtype=torch.float32) * (10000.0 / C)).reshape(1, 1, int(C/2))
+    
+#     pe_x = torch.zeros(B, N, C, device=xyz.device, dtype=torch.float32)
+#     pe_y = torch.zeros(B, N, C, device=xyz.device, dtype=torch.float32)
+#     pe_z = torch.zeros(B, N, C, device=xyz.device, dtype=torch.float32)
+    
+#     pe_x[:, :, 0::2] = torch.sin(x * div_term)
+#     pe_x[:, :, 1::2] = torch.cos(x * div_term)
+    
+#     pe_y[:, :, 0::2] = torch.sin(y * div_term)
+#     pe_y[:, :, 1::2] = torch.cos(y * div_term)
+    
+#     pe_z[:, :, 0::2] = torch.sin(z * div_term)
+#     pe_z[:, :, 1::2] = torch.cos(z * div_term)
+    
+#     pe = torch.cat([pe_x, pe_y, pe_z], dim=2) # B, N, C*3
+#     if cat_coords:
+#         pe = torch.cat([pe, xyz], dim=2) # B, N, C*3+3
+#     return pe
+    
+class SimplePool():
+    def __init__(self, pool_size, version='pt', min_size=1):
+        self.pool_size = pool_size
+        self.version = version
+        self.items = []
+        self.min_size = min_size
+        
+        if not (version=='pt' or version=='np'):
+            print('version = %s; please choose pt or np')
+            assert(False) # please choose pt or np
+            
+    def __len__(self):
+        return len(self.items)
+    
+    def mean(self, min_size=None):
+        if min_size is None:
+            pool_size_thresh = self.min_size
+        elif min_size=='half':
+            pool_size_thresh = self.pool_size/2
+        else:
+            pool_size_thresh = min_size
+            
+        if self.version=='np':
+            if len(self.items) >= pool_size_thresh:
+                return np.sum(self.items)/float(len(self.items))
+            else:
+                return np.nan
+        if self.version=='pt':
+            if len(self.items) >= pool_size_thresh:
+                return torch.sum(self.items)/float(len(self.items))
+            else:
+                return torch.from_numpy(np.nan)
+    
+    def sample(self, with_replacement=True):
+        idx = np.random.randint(len(self.items))
+        if with_replacement:
+            return self.items[idx]
+        else:
+            return self.items.pop(idx)
+    
+    def fetch(self, num=None):
+        if self.version=='pt':
+            item_array = torch.stack(self.items)
+        elif self.version=='np':
+            item_array = np.stack(self.items)
+        if num is not None:
+            # there better be some items
+            assert(len(self.items) >= num)
+                
+            # if there are not that many elements just return however many there are
+            if len(self.items) < num:
+                return item_array
+            else:
+                idxs = np.random.randint(len(self.items), size=num)
+                return item_array[idxs]
+        else:
+            return item_array
+            
+    def is_full(self):
+        full = len(self.items)==self.pool_size
+        return full
+    
+    def empty(self):
+        self.items = []
+
+    def have_min_size(self):
+        return len(self.items) >= self.min_size
+        
+            
+    def update(self, items):
+        for item in items:
+            if len(self.items) < self.pool_size:
+                # the pool is not full, so let's add this in
+                self.items.append(item)
+            else:
+                # the pool is full
+                # pop from the front
+                self.items.pop(0)
+                # add to the back
+                self.items.append(item)
+        return self.items
+
+
+class SimpleHeap():
+    def __init__(self, pool_size, version='pt'):
+        self.pool_size = pool_size
+        self.version = version
+        self.items = []
+        self.vals = []
+        
+        if not (version=='pt' or version=='np'):
+            print('version = %s; please choose pt or np')
+            assert(False) # please choose pt or np
+            
+    def __len__(self):
+        return len(self.items)
+    
+    def sample(self, random=True, with_replacement=True, semirandom=False):
+        vals_arr = np.stack(self.vals)
+        if random:
+            ind = np.random.randint(len(self.items))
+        else:
+            if semirandom and len(vals_arr)>1:
+                # choose from the harder half
+                inds = np.argsort(vals_arr) # ascending
+                inds = inds[len(vals_arr)//2:]
+                ind = np.random.choice(inds)
+            else:
+                # find the most valuable element
+                ind = np.argmax(vals_arr)
+
+            
+        if with_replacement:
+            return self.items[ind]
+        else:
+            item = self.items.pop(ind)
+            val = self.vals.pop(ind)
+            return item
+    
+    def fetch(self, num=None):
+        if self.version=='pt':
+            item_array = torch.stack(self.items)
+        elif self.version=='np':
+            item_array = np.stack(self.items)
+        if num is not None:
+            # there better be some items
+            assert(len(self.items) >= num)
+                
+            # if there are not that many elements just return however many there are
+            if len(self.items) < num:
+                return item_array
+            else:
+                idxs = np.random.randint(len(self.items), size=num)
+                return item_array[idxs]
+        else:
+            return item_array
+            
+    def is_full(self):
+        full = len(self.items)==self.pool_size
+        return full
+    
+    def empty(self):
+        self.items = []
+            
+    def update(self, vals, items):
+        for val,item in zip(vals, items):
+            if len(self.items) < self.pool_size:
+                # the pool is not full, so let's add this in
+                self.items.append(item)
+                self.vals.append(val)
+            else:
+                # the pool is full
+                # find our least-valuable element
+                # and see if we should replace it
+                vals_arr = np.stack(self.vals)
+                ind = np.argmin(vals_arr)
+                if vals_arr[ind] < val:
+                    # pop the min
+                    self.items.pop(ind)
+                    self.vals.pop(ind)
+
+                    # add to the back
+                    self.items.append(item)
+                    self.vals.append(val)
+        return self.items
+    
+
+def farthest_point_sample(xyz, npoint, include_ends=False, deterministic=False):
+    """
+    Input:
+        xyz: pointcloud data, [B, N, C], where C is probably 3
+        npoint: number of samples
+    Return:
+        inds: sampled pointcloud index, [B, npoint]
+    """
+    device = xyz.device
+    B, N, C = xyz.shape
+    xyz = xyz.float()
+    inds = torch.zeros(B, npoint, dtype=torch.long, device=device)
+    distance = torch.ones((B, N), dtype=torch.float32, device=device) * 1e10
+    if deterministic:
+        farthest = torch.randint(0, 1, (B,), dtype=torch.long, device=device)
+    else:
+        farthest = torch.randint(0, N, (B,), dtype=torch.long, device=device)
+    batch_indices = torch.arange(B, dtype=torch.long, device=device)
+    for i in range(npoint):
+        if include_ends:
+            if i==0:
+                farthest = 0
+            elif i==1:
+                farthest = N-1
+        inds[:, i] = farthest
+        centroid = xyz[batch_indices, farthest, :].view(B, 1, C)
+        dist = torch.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = torch.max(distance, -1)[1]
+
+        if npoint > N:
+            # if we need more samples, make them random
+            distance += torch.randn_like(distance)
+    return inds
+
+def farthest_point_sample_py(xyz, npoint, deterministic=False):
+    N,C = xyz.shape
+    inds = np.zeros(npoint, dtype=np.int32)
+    distance = np.ones(N) * 1e10
+    if deterministic:
+        farthest = 0
+    else:
+        farthest = np.random.randint(0, N, dtype=np.int32)
+    for i in range(npoint):
+        inds[i] = farthest
+        centroid = xyz[farthest, :].reshape(1,C)
+        dist = np.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = np.argmax(distance, -1)
+        if npoint > N:
+            # if we need more samples, make them random
+            distance += np.random.randn(*distance.shape)
+    return inds
+    
+def balanced_ce_loss(pred, gt, pos_weight=0.5, valid=None, dim=None, return_both=False, use_halfmask=False, H=64, W=64):
+    # # pred and gt are the same shape
+    # for (a,b) in zip(pred.size(), gt.size()):
+    #     if not a==b:
+    #         print('mismatch: pred, gt', pred.shape, gt.shape)
+    #     assert(a==b) # some shape mismatch!
+
+    pred = pred.reshape(-1)
+    gt = gt.reshape(-1)
+    device = pred.device
+    
+    if valid is not None:
+        valid = valid.reshape(-1)
+        for (a,b) in zip(pred.size(), valid.size()):
+            assert(a==b) # some shape mismatch!
+    else:
+        valid = torch.ones_like(gt)
+
+    pos = (gt > 0.95).float()
+    if use_halfmask:
+        pos_wide = (gt >= 0.5).float()
+        halfmask = (gt == 0.5).float()
+    else:
+        pos_wide = pos
+
+    neg = (gt < 0.05).float()
+
+    label = pos_wide*2.0 - 1.0
+    a = -label * pred
+    b = F.relu(a)
+    loss = b + torch.log(torch.exp(-b)+torch.exp(a-b))
+
+    if torch.sum(pos*valid)>0:
+        pos_loss = loss[(pos*valid) > 0].mean()
+    else:
+        pos_loss = torch.tensor(0.0, requires_grad=True, device=device)
+        
+    if torch.sum(neg*valid)>0:
+        neg_loss = loss[(neg*valid) > 0].mean()
+    else:
+        neg_loss = torch.tensor(0.0, requires_grad=True, device=device)
+    balanced_loss = pos_weight*pos_loss + (1-pos_weight)*neg_loss
+    return balanced_loss
+
+    # pos_loss = utils.basic.reduce_masked_mean(loss, pos*valid, dim=dim)
+    # neg_loss = utils.basic.reduce_masked_mean(loss, neg*valid, dim=dim)
+
+    if use_halfmask:
+        # here we will find the pixels which are already leaning positive,
+        # and encourage them to be more positive
+        B = loss.shape[0]
+        loss_ = loss.reshape(B,-1)
+        mask_ = halfmask.reshape(B,-1) * valid.reshape(B,-1)
+
+        # to avoid the issue where spikes become spikier,
+        # we will only apply this loss on batch els where we predicted zero positives
+        pred_sig_ = torch.sigmoid(pred).reshape(B,-1)
+        no_pred_ = torch.max(pred_sig_.round(), axis=1)[0] < 1 # B
+        # and only on batch els where we have negatives available
+        have_neg_ = torch.sum(neg, dim=1)>0 # B
+
+        loss_ = loss_[no_pred_ & have_neg_] # N,H*W
+        mask_ = mask_[no_pred_ & have_neg_] # N,H*W
+        N = loss_.shape[0]
+
+        if N > 0:
+            # we want: 
+            # in the neg pixels, 
+            # set them to the max loss of the pos pixels,
+            # so that they do not contribute to the min
+            loss__ = loss_.reshape(-1)
+            mask__ = mask_.reshape(-1)
+            if torch.sum(mask__)>0:
+                # print('loss_', loss_.shape, 'mask_', mask_.shape, 'loss__', loss__.shape, 'mask__', mask__.shape)
+                mloss__ = loss__.detach()
+                mloss__[mask__==0] = torch.max(loss__[mask__==1])
+                mloss_ = mloss__.reshape(N,H*W)
+
+                # now, in each batch el, take a tiny region around the argmin, so we can boost this region
+                minloss_mask_ = torch.zeros_like(mloss_).scatter(1,mloss_.argmin(1,True),value=1)
+                minloss_mask_ = utils.improc.dilate2d(minloss_mask_.view(N,1,H,W), times=3).reshape(N,H*W)
+
+                loss__ = loss_.reshape(-1)
+                minloss_mask__ = minloss_mask_.reshape(-1)
+                half_loss = loss__[minloss_mask__>0].mean()
+
+                # print('N', N, 'half_loss', half_loss)
+                pos_loss = pos_loss + half_loss
+
+        # if False:
+        #     min_pos = 8
+
+        #     # only apply the loss when we have some negatives available,
+        #     # otherwise it's a whole "ignore" frame, which may mean
+        #     # we are unsure if the target is even there
+        #     if torch.sum(mask__==0) > 0: # negatives available
+        #         # only apply the loss when the halfmask is larger area than
+        #         # min_pos (the number of pixels we want to boost),
+        #         # so that indexing will work 
+        #         if torch.all(torch.sum(mask_==1, dim=1) >= min_pos): # topk indexing will work
+        #             # in the pixels we will not use,
+        #             # set them to the max of the pixels we may use,
+        #             # so that they do not contribute to the min
+        #             loss__[mask__==0] = torch.max(loss__[mask__==1])
+        #             loss_ = loss__.reshape(B,-1)
+
+        #             half_loss = torch.mean(torch.topk(loss_, min_pos, dim=1, largest=False)[0], dim=1) # B
+
+        #             have_neg = (torch.sum(neg, dim=1)>0).float() # B
+        #             pos_loss = pos_loss + half_loss*have_neg
+
+                
+
+
+            
+        
+        # half_loss = []
+        # for b in range(B):
+        #     loss_b = loss_[b]
+        #     mask_b = mask_[b]
+        #     if torch.sum(mask_b):
+        #         inds = torch.nonzero(mask_b).reshape(-1)
+        #         half_loss.append(torch.min(loss_b[inds]))
+        # if len(half_loss):
+        #     # # half_loss_ = half_loss.reshape(-1)
+        #     # half_loss = torch.min(half_loss, dim=1)[0] # B
+        #     # half_loss = torch.mean(torch.topk(half_loss, 4, dim=1, largest=False)[0], dim=1) # B
+        #     pos_loss = pos_loss + torch.stack(half_loss).mean()
+
+    if return_both:
+        return pos_loss, neg_loss
+    balanced_loss = pos_weight*pos_loss + (1-pos_weight)*neg_loss
+
+    return balanced_loss
+
+def dice_loss(pred, gt):
+    # gt has ignores at 0.5
+    # pred and gt are the same shape
+    for (a,b) in zip(pred.size(), gt.size()):
+        assert(a==b) # some shape mismatch!
+    
+    prob = pred.sigmoid()
+
+    # flatten everything except batch
+    prob = prob.flatten(1)
+    gt = gt.flatten(1)
+    
+    pos = (gt > 0.95).float()
+    neg = (gt < 0.05).float()
+    valid = (pos+neg).float().clamp(0,1)
+    
+    numerator = 2 * (prob * pos * valid).sum(1)
+    denominator = (prob*valid).sum(1) + (pos*valid).sum(1)
+    loss = 1 - (numerator + 1) / (denominator + 1)
+    return loss
+
+def sigmoid_focal_loss(pred, gt, alpha=0.25, gamma=2):#, use_halfmask=False):
+    # gt has ignores at 0.5
+    # pred and gt are the same shape
+    for (a,b) in zip(pred.size(), gt.size()):
+        assert(a==b) # some shape mismatch!
+
+    # flatten everything except batch
+    pred = pred.flatten(1)
+    gt = gt.flatten(1)
+        
+    pos = (gt > 0.95).float()
+    neg = (gt < 0.05).float()
+    # if use_halfmask:
+    #     pos_wide = (gt >= 0.5).float()
+    #     halfmask = (gt == 0.5).float()
+    # else:
+    #     pos_wide = pos
+    valid = (pos+neg).float().clamp(0,1)
+    
+    prob = pred.sigmoid()
+    ce_loss = F.binary_cross_entropy_with_logits(pred, pos, reduction="none")
+    p_t = prob * pos + (1 - prob) * (1 - pos)
+    loss = ce_loss * ((1 - p_t) ** gamma)
+
+    if alpha >= 0:
+        alpha_t = alpha * pos + (1 - alpha) * (1 - pos)
+        loss = alpha_t * loss
+
+    loss = (loss*valid).sum(1) / (1 + valid.sum(1))
+    return loss 
+    
+# def dice_loss(inputs, targets, normalizer=1):
+#     inputs = inputs.sigmoid()
+#     inputs = inputs.flatten(1)
+#     numerator = 2 * (inputs * targets).sum(1)
+#     denominator = inputs.sum(-1) + targets.sum(-1)
+#     loss = 1 - (numerator + 1) / (denominator + 1)
+#     return loss.sum() / normalizer
+
+# def sigmoid_focal_loss(inputs, targets, normalizer=1, alpha=0.25, gamma=2):
+#     prob = inputs.sigmoid()
+#     ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
+#     p_t = prob * targets + (1 - prob) * (1 - targets)
+#     loss = ce_loss * ((1 - p_t) ** gamma)
+
+#     if alpha >= 0:
+#         alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
+#         loss = alpha_t * loss
+
+#     return loss.mean(1).sum() / normalizer
+                                        
+def data_replace_with_nearest(xys, valids):
+    # replace invalid xys with nearby ones
+    invalid_idx = np.where(valids==0)[0]
+    valid_idx = np.where(valids==1)[0]
+    for idx in invalid_idx:
+        nearest = valid_idx[np.argmin(np.abs(valid_idx - idx))]
+        xys[idx] = xys[nearest]
+    return xys
+
+def data_get_traj_from_masks(masks):
+    if masks.ndim==4:
+        masks = masks[...,0]
+    S, H, W = masks.shape
+    masks = (masks > 0.1).astype(np.float32)
+    fills = np.zeros((S))
+    xy_means = np.zeros((S,2))
+    xy_rands = np.zeros((S,2))
+    valids = np.zeros((S))
+    for si, mask in enumerate(masks):
+        if np.sum(mask) > 0:
+            ys, xs = np.where(mask)
+            inds = np.random.permutation(len(xs))
+            xs, ys = xs[inds], ys[inds]
+            x0, x1 = np.min(xs), np.max(xs)+1
+            y0, y1 = np.min(ys), np.max(ys)+1
+            # if (x1-x0)>0 and (y1-y0)>0:
+            xy_means[si] = np.array([xs.mean(), ys.mean()])
+            xy_rands[si] = np.array([xs[0], ys[0]])
+            valids[si] = 1
+            crop = mask[y0:y1, x0:x1]
+            fill = np.mean(crop)
+            fills[si] = fill
+        # print('fills', fills)
+    return xy_means, xy_rands, valids, fills
+
+def data_zoom(zoom, xys, visibs, rgbs, valids=None, masks=None, masks2=None, masks3=None, masks4=None):
+    S, H, W, C = rgbs.shape
+    S,N,D = xys.shape
+    
+    _, H, W, C = rgbs.shape
+    assert(C==3)
+    crop_W = int(W//zoom)
+    crop_H = int(H//zoom)
+
+    if np.random.rand() < 0.25: # follow-crop
+        # start with xy traj
+        # smooth_xys = xys.copy()
+        smooth_xys = xys[:,np.random.randint(N)].reshape(S,1,2)
+        # make it inbounds
+        smooth_xys = np.clip(smooth_xys, [crop_W // 2, crop_H // 2], [W - crop_W // 2, H - crop_H // 2])
+        
+        # smooth it out, to remove info about the traj, and simulate camera motion
+        for _ in range(S*3):
+            for ii in range(S):
+                if ii==0:
+                    smooth_xys[ii] = (smooth_xys[ii] + smooth_xys[ii+1])/2.0
+                elif ii==S-1:
+                    smooth_xys[ii] = (smooth_xys[ii-1] + smooth_xys[ii])/2.0
+                else:
+                    smooth_xys[ii] = (smooth_xys[ii-1] + smooth_xys[ii] + smooth_xys[ii+1])/3.0
+    else: # static (no-hint) crop
+        # zero-vel on random available coordinate
+
+        if valids is not None:
+            visval = visibs*valids # S,N
+            visval = np.sum(visval, axis=1) # S
+        else:
+            visval = np.sum(visibs, axis=1) # S
+            
+        anchor_inds = np.nonzero(visval >= np.mean(visval))[0]
+        ind = anchor_inds[np.random.randint(len(anchor_inds))]
+        # print('ind', ind)
+        smooth_xys = xys[ind:ind+1].repeat(S,axis=0)
+        smooth_xys = smooth_xys.mean(axis=1, keepdims=True)
+        # xmid = np.random.randint(crop_W//2, W-crop_W//2)
+        # ymid = np.random.randint(crop_H//2, H-crop_H//2)
+        # smooth_xys = np.stack([xmid, ymid], axis=-1).reshape(1,1,2).repeat(S, axis=0) # S,1,2
+        smooth_xys = np.clip(smooth_xys, [crop_W // 2, crop_H // 2], [W - crop_W // 2, H - crop_H // 2])
+                
+    if np.random.rand() < 0.5:
+        # add a random alternate trajectory, to help push us off center
+        alt_xys = np.random.randint(-crop_H//8, crop_H//8, (S,1,2))
+        for _ in range(4): # smooth out
+            for ii in range(S):
+                if ii==0:
+                    alt_xys[ii] = (alt_xys[ii] + alt_xys[ii+1])/2.0
+                elif ii==S-1:
+                    alt_xys[ii] = (alt_xys[ii-1] + alt_xys[ii])/2.0
+                else:
+                    alt_xys[ii] = (alt_xys[ii-1] + alt_xys[ii] + alt_xys[ii+1])/3.0
+        smooth_xys = smooth_xys + alt_xys
+        
+    smooth_xys = np.clip(smooth_xys, [crop_W // 2, crop_H // 2], [W - crop_W // 2, H - crop_H // 2])
+
+    rgbs_crop = []
+    if masks is not None:
+        masks_crop = []
+    if masks2 is not None:
+        masks2_crop = []
+    if masks3 is not None:
+        masks3_crop = []
+    if masks4 is not None:
+        masks4_crop = []
+
+    offsets = []
+    for si in range(S):
+        xy_mid = smooth_xys[si].squeeze(0).round().astype(np.int32) # 2
+
+        xmid, ymid = xy_mid[0], xy_mid[1]
+
+        x0, x1 = np.clip(xmid-crop_W//2, 0, W), np.clip(xmid+crop_W//2, 0, W)
+        y0, y1 = np.clip(ymid-crop_H//2, 0, H), np.clip(ymid+crop_H//2, 0, H)
+        offset = np.array([x0, y0]).reshape(1,2)
+
+        rgbs_crop.append(rgbs[si,y0:y1,x0:x1])
+        if masks is not None:
+            masks_crop.append(masks[si,y0:y1,x0:x1])
+        if masks2 is not None:
+            masks2_crop.append(masks2[si,y0:y1,x0:x1])
+        if masks3 is not None:
+            masks3_crop.append(masks3[si,y0:y1,x0:x1])
+        if masks4 is not None:
+            masks4_crop.append(masks4[si,y0:y1,x0:x1])
+        xys[si] -= offset
+
+        offsets.append(offset)
+
+    rgbs = np.stack(rgbs_crop, axis=0)
+    if masks is not None:
+        masks = np.stack(masks_crop, axis=0)
+    if masks2 is not None:
+        masks2 = np.stack(masks2_crop, axis=0)
+    if masks3 is not None:
+        masks3 = np.stack(masks3_crop, axis=0)
+    if masks4 is not None:
+        masks4 = np.stack(masks4_crop, axis=0)
+
+    # update visibility annotations
+    for si in range(S):
+        oob_inds = np.logical_or(
+            np.logical_or(xys[si,:,0] < 0, xys[si,:,0] > crop_W-1),
+            np.logical_or(xys[si,:,1] < 0, xys[si,:,1] > crop_H-1))
+        visibs[si,oob_inds] = 0
+    
+    # if masks4 is not None:
+    #     return xys, visibs, valids, rgbs, masks, masks2, masks3, masks4
+    # if masks3 is not None:
+    #     return xys, visibs, valids, rgbs, masks, masks2, masks3
+    # if masks2 is not None:
+    #     return xys, visibs, valids, rgbs, masks, masks2
+    # if masks is not None:
+    #     return xys, visibs, valids, rgbs, masks
+    # else:
+    # return xys, visibs, valids, rgbs
+    if valids is not None:
+        return xys, visibs, rgbs, valids
+    else:
+        return xys, visibs, rgbs
+
+
+def data_zoom_bbox(zoom, bboxes, visibs, rgbs):#, valids=None):
+    S, H, W, C = rgbs.shape
+
+    _, H, W, C = rgbs.shape
+    assert(C==3)
+    crop_W = int(W//zoom)
+    crop_H = int(H//zoom)
+
+    xys = bboxes[:,0:2]*0.5 + bboxes[:,2:4]*0.5
+
+    if np.random.rand() < 0.25: # follow-crop
+        # start with xy traj
+        smooth_xys = xys.copy()
+        # make it inbounds
+        smooth_xys = np.clip(smooth_xys, [crop_W // 2, crop_H // 2], [W - crop_W // 2, H - crop_H // 2])
+        # smooth it out, to remove info about the traj, and simulate camera motion
+        for _ in range(S*3):
+            for ii in range(1,S-1):
+                smooth_xys[ii] = (smooth_xys[ii-1] + smooth_xys[ii] + smooth_xys[ii+1])/3.0
+    else: # static (no-hint) crop
+        # zero-vel on random available coordinate
+        anchor_inds = np.nonzero(visibs.reshape(-1)>0.5)[0]
+        ind = anchor_inds[np.random.randint(len(anchor_inds))]
+        smooth_xys = xys[ind:ind+1].repeat(S,axis=0)
+        # xmid = np.random.randint(crop_W//2, W-crop_W//2)
+        # ymid = np.random.randint(crop_H//2, H-crop_H//2)
+        # smooth_xys = np.stack([xmid, ymid], axis=-1).reshape(1,1,2).repeat(S, axis=0) # S,1,2
+        smooth_xys = np.clip(smooth_xys, [crop_W // 2, crop_H // 2], [W - crop_W // 2, H - crop_H // 2])
+    # print('xys', xys)
+    # print('smooth_xys', smooth_xys)
+                
+    if np.random.rand() < 0.5:
+        # add a random alternate trajectory, to help push us off center
+        alt_xys = np.random.randint(-crop_H//8, crop_H//8, (S,2))
+        for _ in range(3):
+            for ii in range(1,S-1):
+                alt_xys[ii] = (alt_xys[ii-1] + alt_xys[ii] + alt_xys[ii+1])/3.0
+        smooth_xys = smooth_xys + alt_xys
+        
+    smooth_xys = np.clip(smooth_xys, [crop_W // 2, crop_H // 2], [W - crop_W // 2, H - crop_H // 2])
+    
+    rgbs_crop = []
+
+    offsets = []
+    for si in range(S):
+        xy_mid = smooth_xys[si].round().astype(np.int32)
+        xmid, ymid = xy_mid[0], xy_mid[1]
+
+        x0, x1 = np.clip(xmid-crop_W//2, 0, W), np.clip(xmid+crop_W//2, 0, W)
+        y0, y1 = np.clip(ymid-crop_H//2, 0, H), np.clip(ymid+crop_H//2, 0, H)
+        offset = np.array([x0, y0]).reshape(2)
+
+        rgbs_crop.append(rgbs[si,y0:y1,x0:x1])
+        xys[si] -= offset
+        bboxes[si,0:2] -= offset
+        bboxes[si,2:4] -= offset
+        
+        offsets.append(offset)
+
+    rgbs = np.stack(rgbs_crop, axis=0)
+
+    # update visibility annotations
+    for si in range(S):
+        # avoid 1px edge
+        oob_inds = np.logical_or(
+            np.logical_or(xys[si,0] < 1, xys[si,0] > W-2),
+            np.logical_or(xys[si,1] < 1, xys[si,1] > H-2))
+        visibs[si,oob_inds] = 0
+
+    # clamp to image bounds
+    xys0 = np.minimum(np.maximum(bboxes[:,0:2], np.zeros((2,), dtype=int)), np.array([W, H]) - 1) # S,2
+    xys1 = np.minimum(np.maximum(bboxes[:,2:4], np.zeros((2,), dtype=int)), np.array([W, H]) - 1) # S,2
+    bboxes = np.concatenate([xys0, xys1], axis=1)
+    return bboxes, visibs, rgbs 
+        
+
+def data_pad_if_necessary(rgbs, masks, masks2=None):
+    S,H,W,C = rgbs.shape
+    
+    mask_areas = (masks > 0).reshape(S,-1).sum(axis=1)
+    mask_areas_norm = mask_areas / np.max(mask_areas)
+    visibs = mask_areas_norm
+    
+    bboxes = np.stack([mask2bbox(mask) for mask in masks])
+    whs = bboxes[:,2:4] - bboxes[:,0:2]
+    whs = whs[visibs > 0.5]
+    # print('mean wh', np.mean(whs[:,0]), np.mean(whs[:,1]))
+    if np.mean(whs[:,0]) >= W/2:
+        # print('padding w')
+        pad = ((0,0),(0,0),(W//4,W//4),(0,0))
+        rgbs = np.pad(rgbs, pad, mode="constant")
+        masks = np.pad(masks, pad[:3], mode="constant")
+        if masks2 is not None:
+            masks2 = np.pad(masks2, pad[:3], mode="constant")
+    # print('rgbs', rgbs.shape)
+    # print('masks', masks.shape)
+    if np.mean(whs[:,1]) >= H/2:
+        # print('padding h')
+        pad = ((0,0),(H//4,H//4),(0,0),(0,0))
+        rgbs = np.pad(rgbs, pad, mode="constant")
+        masks = np.pad(masks, pad[:3], mode="constant")
+        if masks2 is not None:
+            masks2 = np.pad(masks2, pad[:3], mode="constant", constant_values=0.5)
+
+    if masks2 is not None:
+        return rgbs, masks, masks2
+    return rgbs, masks
+
+def data_pad_if_necessary_b(rgbs, bboxes, visibs):
+    S,H,W,C = rgbs.shape
+    whs = bboxes[:,2:4] - bboxes[:,0:2]
+    whs = whs[visibs > 0.5]
+    if np.mean(whs[:,0]) >= W/2:
+        pad = ((0,0),(0,0),(W//4,W//4),(0,0))
+        rgbs = np.pad(rgbs, pad, mode="constant")
+        bboxes[:,0] += W//4
+        bboxes[:,2] += W//4
+    if np.mean(whs[:,1]) >= H/2:
+        pad = ((0,0),(H//4,H//4),(0,0),(0,0))
+        rgbs = np.pad(rgbs, pad, mode="constant")
+        bboxes[:,1] += H//4
+        bboxes[:,3] += H//4
+    return rgbs, bboxes
+
+def posenc(x, min_deg, max_deg):
+    """Cat x with a positional encoding of x with scales 2^[min_deg, max_deg-1].
+    Instead of computing [sin(x), cos(x)], we use the trig identity
+    cos(x) = sin(x + pi/2) and do one vectorized call to sin([x, x+pi/2]).
+    Args:
+      x: torch.Tensor, variables to be encoded. Note that x should be in [-pi, pi].
+      min_deg: int, the minimum (inclusive) degree of the encoding.
+      max_deg: int, the maximum (exclusive) degree of the encoding.
+      legacy_posenc_order: bool, keep the same ordering as the original tf code.
+    Returns:
+      encoded: torch.Tensor, encoded variables.
+    """
+    if min_deg == max_deg:
+        return x
+    scales = torch.tensor(
+        [2**i for i in range(min_deg, max_deg)], dtype=x.dtype, device=x.device
+    )
+
+    xb = (x[..., None, :] * scales[:, None]).reshape(list(x.shape[:-1]) + [-1])
+    four_feat = torch.sin(torch.cat([xb, xb + 0.5 * torch.pi], dim=-1))
+    return torch.cat([x] + [four_feat], dim=-1)
+

utils/py.py

@@ -0,0 +1,755 @@
+import glob, math
+import numpy as np
+# from scipy import misc
+# from scipy import linalg
+from PIL import Image
+import io
+import matplotlib.pyplot as plt
+EPS = 1e-6
+
+
+XMIN = -64.0 # right (neg is left)
+XMAX = 64.0 # right
+YMIN = -64.0 # down (neg is up)
+YMAX = 64.0 # down
+ZMIN = -64.0 # forward
+ZMAX = 64.0 # forward
+
+def print_stats(name, tensor):
+    tensor = tensor.astype(np.float32)
+    print('%s min = %.2f, mean = %.2f, max = %.2f' % (name, np.min(tensor), np.mean(tensor), np.max(tensor)), tensor.shape)
+    
+def reduce_masked_mean(x, mask, axis=None, keepdims=False):
+    # x and mask are the same shape
+    # returns shape-1
+    # axis can be a list of axes
+    prod = x*mask
+    numer = np.sum(prod, axis=axis, keepdims=keepdims)
+    denom = EPS+np.sum(mask, axis=axis, keepdims=keepdims)
+    mean = numer/denom
+    return mean
+
+def reduce_masked_sum(x, mask, axis=None, keepdims=False):
+    # x and mask are the same shape
+    # returns shape-1
+    # axis can be a list of axes
+    prod = x*mask
+    numer = np.sum(prod, axis=axis, keepdims=keepdims)
+    return numer
+
+def reduce_masked_median(x, mask, keep_batch=False):
+    # x and mask are the same shape
+    # returns shape-1
+    # axis can be a list of axes
+
+    if not (x.shape == mask.shape):
+        print('reduce_masked_median: these shapes should match:', x.shape, mask.shape)
+        assert(False)
+    # assert(x.shape == mask.shape)
+
+    B = list(x.shape)[0]
+
+    if keep_batch:
+        x = np.reshape(x, [B, -1])
+        mask = np.reshape(mask, [B, -1])
+        meds = np.zeros([B], np.float32)
+        for b in list(range(B)):
+            xb = x[b]
+            mb = mask[b]
+            if np.sum(mb) > 0:
+                xb = xb[mb > 0]
+                meds[b] = np.median(xb)
+            else:
+                meds[b] = np.nan
+        return meds
+    else:
+        x = np.reshape(x, [-1])
+        mask = np.reshape(mask, [-1])
+        if np.sum(mask) > 0:
+            x = x[mask > 0]
+            med = np.median(x)
+        else:
+            med = np.nan
+        med = np.array([med], np.float32)
+        return med
+    
+def get_nFiles(path):
+    return len(glob.glob(path))
+
+def get_file_list(path):
+    return glob.glob(path)
+
+def rotm2eul(R):
+    # R is 3x3
+    sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])
+    if sy > 1e-6: # singular
+        x = math.atan2(R[2,1] , R[2,2])
+        y = math.atan2(-R[2,0], sy)
+        z = math.atan2(R[1,0], R[0,0])
+    else:
+        x = math.atan2(-R[1,2], R[1,1])
+        y = math.atan2(-R[2,0], sy)
+        z = 0
+    return x, y, z
+            
+def rad2deg(rad):
+    return rad*180.0/np.pi
+
+def deg2rad(deg):
+    return deg/180.0*np.pi
+            
+def eul2rotm(rx, ry, rz):
+    # copy of matlab, but order of inputs is different
+    # R = [  cy*cz   sy*sx*cz-sz*cx    sy*cx*cz+sz*sx
+    #        cy*sz   sy*sx*sz+cz*cx    sy*cx*sz-cz*sx
+    #        -sy            cy*sx             cy*cx]
+    sinz = np.sin(rz)
+    siny = np.sin(ry)
+    sinx = np.sin(rx)
+    cosz = np.cos(rz)
+    cosy = np.cos(ry)
+    cosx = np.cos(rx)
+    r11 = cosy*cosz
+    r12 = sinx*siny*cosz - cosx*sinz
+    r13 = cosx*siny*cosz + sinx*sinz
+    r21 = cosy*sinz
+    r22 = sinx*siny*sinz + cosx*cosz
+    r23 = cosx*siny*sinz - sinx*cosz
+    r31 = -siny
+    r32 = sinx*cosy
+    r33 = cosx*cosy
+    r1 = np.stack([r11,r12,r13],axis=-1)
+    r2 = np.stack([r21,r22,r23],axis=-1)
+    r3 = np.stack([r31,r32,r33],axis=-1)
+    r = np.stack([r1,r2,r3],axis=0)
+    return r
+
+def wrap2pi(rad_angle):
+    # puts the angle into the range [-pi, pi]
+    return np.arctan2(np.sin(rad_angle), np.cos(rad_angle))
+
+def rot2view(rx,ry,rz,x,y,z):
+    # takes rot angles and 3d position as input
+    # returns viewpoint angles as output
+    # (all in radians)
+    # it will perform strangely if z <= 0
+    az = wrap2pi(ry - (-np.arctan2(z, x) - 1.5*np.pi))
+    el = -wrap2pi(rx - (-np.arctan2(z, y) - 1.5*np.pi))
+    th = -rz
+    return az, el, th
+
+def invAxB(a,b):
+    """
+    Compute the relative 3d transformation between a and b.
+    
+    Input:
+    a -- first pose (homogeneous 4x4 matrix)
+    b -- second pose (homogeneous 4x4 matrix)
+    
+    Output:
+    Relative 3d transformation from a to b.
+    """
+    return np.dot(np.linalg.inv(a),b)
+        
+def merge_rt(r, t):
+    # r is 3 x 3
+    # t is 3 or maybe 3 x 1
+    t = np.reshape(t, [3, 1])
+    rt = np.concatenate((r,t), axis=1)
+    # rt is 3 x 4
+    br = np.reshape(np.array([0,0,0,1], np.float32), [1, 4])
+    # br is 1 x 4 
+    rt = np.concatenate((rt, br), axis=0)
+    # rt is 4 x 4
+    return rt
+
+def split_rt(rt):
+    r = rt[:3,:3]
+    t = rt[:3,3]
+    r = np.reshape(r, [3, 3])
+    t = np.reshape(t, [3, 1])
+    return r, t
+
+def split_intrinsics(K):
+    # K is 3 x 4 or 4 x 4
+    fx = K[0,0]
+    fy = K[1,1]
+    x0 = K[0,2]
+    y0 = K[1,2]
+    return fx, fy, x0, y0
+                    
+def merge_intrinsics(fx, fy, x0, y0):
+    # inputs are shaped []
+    K = np.eye(4)
+    K[0,0] = fx
+    K[1,1] = fy
+    K[0,2] = x0
+    K[1,2] = y0
+    # K is shaped 4 x 4
+    return K
+                            
+def scale_intrinsics(K, sx, sy):
+    fx, fy, x0, y0 = split_intrinsics(K)
+    fx *= sx
+    fy *= sy
+    x0 *= sx
+    y0 *= sy
+    return merge_intrinsics(fx, fy, x0, y0)
+
+# def meshgrid(H, W):
+#     x = np.linspace(0, W-1, W)
+#     y = np.linspace(0, H-1, H)
+#     xv, yv = np.meshgrid(x, y)
+#     return xv, yv
+
+def compute_distance(transform):
+    """
+    Compute the distance of the translational component of a 4x4 homogeneous matrix.
+    """
+    return numpy.linalg.norm(transform[0:3,3])
+
+def radian_l1_dist(e, g):
+    # if our angles are in [0, 360] we can follow this stack overflow answer:
+    # https://gamedev.stackexchange.com/questions/4467/comparing-angles-and-working-out-the-difference
+    # wrap2pi brings the angles to [-180, 180]; adding pi puts them in [0, 360]
+    e = wrap2pi(e)+np.pi
+    g = wrap2pi(g)+np.pi
+    l = np.abs(np.pi - np.abs(np.abs(e-g) - np.pi))
+    return l
+
+def apply_pix_T_cam(pix_T_cam, xyz):
+    fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
+    # xyz is shaped B x H*W x 3
+    # returns xy, shaped B x H*W x 2
+    N, C = xyz.shape
+    x, y, z = np.split(xyz, 3, axis=-1)
+    EPS = 1e-4
+    z = np.clip(z, EPS, None)
+    x = (x*fx)/(z)+x0
+    y = (y*fy)/(z)+y0
+    xy = np.concatenate([x, y], axis=-1)
+    return xy
+
+def apply_4x4(RT, XYZ):
+    # RT is 4 x 4
+    # XYZ is N x 3
+
+    # put into homogeneous coords
+    X, Y, Z = np.split(XYZ, 3, axis=1)
+    ones = np.ones_like(X)
+    XYZ1 = np.concatenate([X, Y, Z, ones], axis=1)
+    # XYZ1 is N x 4
+
+    XYZ1_t = np.transpose(XYZ1)
+    # this is 4 x N
+
+    XYZ2_t = np.dot(RT, XYZ1_t)
+    # this is 4 x N
+    
+    XYZ2 = np.transpose(XYZ2_t)
+    # this is N x 4
+    
+    XYZ2 = XYZ2[:,:3]
+    # this is N x 3
+    
+    return XYZ2
+
+def Ref2Mem(xyz, Z, Y, X):
+    # xyz is N x 3, in ref coordinates
+    # transforms ref coordinates into mem coordinates
+    N, C = xyz.shape
+    assert(C==3)
+    mem_T_ref = get_mem_T_ref(Z, Y, X)
+    xyz = apply_4x4(mem_T_ref, xyz)
+    return xyz
+
+# def Mem2Ref(xyz_mem, MH, MW, MD):
+#     # xyz is B x N x 3, in mem coordinates
+#     # transforms mem coordinates into ref coordinates
+#     B, N, C = xyz_mem.get_shape().as_list()
+#     ref_T_mem = get_ref_T_mem(B, MH, MW, MD)
+#     xyz_ref = utils_geom.apply_4x4(ref_T_mem, xyz_mem)
+#     return xyz_ref
+
+def get_mem_T_ref(Z, Y, X):
+    # sometimes we want the mat itself
+    # note this is not a rigid transform
+    
+    # for interpretability, let's construct this in two steps...
+
+    # translation
+    center_T_ref = np.eye(4, dtype=np.float32)
+    center_T_ref[0,3] = -XMIN
+    center_T_ref[1,3] = -YMIN
+    center_T_ref[2,3] = -ZMIN
+
+    VOX_SIZE_X = (XMAX-XMIN)/float(X)
+    VOX_SIZE_Y = (YMAX-YMIN)/float(Y)
+    VOX_SIZE_Z = (ZMAX-ZMIN)/float(Z)
+    
+    # scaling
+    mem_T_center = np.eye(4, dtype=np.float32)
+    mem_T_center[0,0] = 1./VOX_SIZE_X
+    mem_T_center[1,1] = 1./VOX_SIZE_Y
+    mem_T_center[2,2] = 1./VOX_SIZE_Z
+    
+    mem_T_ref = np.dot(mem_T_center, center_T_ref)
+    return mem_T_ref
+
+def safe_inverse(a):
+    r, t = split_rt(a)
+    t = np.reshape(t, [3, 1])
+    r_transpose = r.T
+    inv = np.concatenate([r_transpose, -np.matmul(r_transpose, t)], 1)
+    bottom_row = a[3:4, :] # this is [0, 0, 0, 1]
+    inv = np.concatenate([inv, bottom_row], 0)
+    return inv
+
+def get_ref_T_mem(Z, Y, X):
+    mem_T_ref = get_mem_T_ref(X, Y, X)
+    # note safe_inverse is inapplicable here,
+    # since the transform is nonrigid
+    ref_T_mem = np.linalg.inv(mem_T_ref)
+    return ref_T_mem
+
+def voxelize_xyz(xyz_ref, Z, Y, X):
+    # xyz_ref is N x 3
+    xyz_mem = Ref2Mem(xyz_ref, Z, Y, X)
+    # this is N x 3
+    voxels = get_occupancy(xyz_mem, Z, Y, X)
+    voxels = np.reshape(voxels, [Z, Y, X, 1])
+    return voxels
+
+def get_inbounds(xyz, Z, Y, X, already_mem=False):
+    # xyz is H*W x 3
+    
+    if not already_mem:
+        xyz = Ref2Mem(xyz, Z, Y, X)
+    
+    x_valid = np.logical_and(
+        np.greater_equal(xyz[:,0], -0.5), 
+        np.less(xyz[:,0], float(X)-0.5))
+    y_valid = np.logical_and(
+        np.greater_equal(xyz[:,1], -0.5), 
+        np.less(xyz[:,1], float(Y)-0.5))
+    z_valid = np.logical_and(
+        np.greater_equal(xyz[:,2], -0.5), 
+        np.less(xyz[:,2], float(Z)-0.5))
+    inbounds = np.logical_and(np.logical_and(x_valid, y_valid), z_valid)
+    return inbounds
+
+def sub2ind3d_zyx(depth, height, width, d, h, w):
+    # same as sub2ind3d, but inputs in zyx order
+    # when gathering/scattering with these inds, the tensor should be Z x Y x X
+    return d*height*width + h*width + w
+
+def sub2ind3d_yxz(height, width, depth, h, w, d):
+    return h*width*depth + w*depth + d
+
+# def ind2sub(height, width, ind):
+#     # int input
+#     y = int(ind / height)
+#     x = ind % height
+#     return y, x
+
+def get_occupancy(xyz_mem, Z, Y, X):
+    # xyz_mem is N x 3
+    # we want to fill a voxel tensor with 1's at these inds
+
+    inbounds = get_inbounds(xyz_mem, Z, Y, X, already_mem=True)
+    inds = np.where(inbounds)
+
+    xyz_mem = np.reshape(xyz_mem[inds], [-1, 3])
+    # xyz_mem is N x 3
+
+    # this is more accurate than a cast/floor, but runs into issues when Y==0
+    xyz_mem = np.round(xyz_mem).astype(np.int32)
+    x = xyz_mem[:,0]
+    y = xyz_mem[:,1]
+    z = xyz_mem[:,2]
+
+    voxels = np.zeros([Z, Y, X], np.float32)
+    voxels[z, y, x] = 1.0
+
+    return voxels
+
+def pixels2camera(x,y,z,fx,fy,x0,y0):
+    # x and y are locations in pixel coordinates, z is a depth image in meters
+    # their shapes are H x W
+    # fx, fy, x0, y0 are scalar camera intrinsics
+    # returns xyz, sized [B,H*W,3]
+    
+    H, W = z.shape
+    
+    fx = np.reshape(fx, [1,1])
+    fy = np.reshape(fy, [1,1])
+    x0 = np.reshape(x0, [1,1])
+    y0 = np.reshape(y0, [1,1])
+    
+    # unproject
+    x = ((z+EPS)/fx)*(x-x0)
+    y = ((z+EPS)/fy)*(y-y0)
+    
+    x = np.reshape(x, [-1])
+    y = np.reshape(y, [-1])
+    z = np.reshape(z, [-1])
+    xyz = np.stack([x,y,z], axis=1)
+    return xyz
+
+def depth2pointcloud(z, pix_T_cam):
+    H = z.shape[0]
+    W = z.shape[1]
+    y, x = meshgrid2d(H, W)
+    z = np.reshape(z, [H, W])
+    
+    fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
+    xyz = pixels2camera(x, y, z, fx, fy, x0, y0)
+    return xyz
+
+def meshgrid2d(Y, X):
+    grid_y = np.linspace(0.0, Y-1, Y)
+    grid_y = np.reshape(grid_y, [Y, 1])
+    grid_y = np.tile(grid_y, [1, X])
+
+    grid_x = np.linspace(0.0, X-1, X)
+    grid_x = np.reshape(grid_x, [1, X])
+    grid_x = np.tile(grid_x, [Y, 1])
+
+    # outputs are Y x X
+    return grid_y, grid_x
+
+def gridcloud3d(Y, X, Z):
+    x_ = np.linspace(0, X-1, X)
+    y_ = np.linspace(0, Y-1, Y)
+    z_ = np.linspace(0, Z-1, Z)
+    y, x, z = np.meshgrid(y_, x_, z_, indexing='ij')
+    x = np.reshape(x, [-1])
+    y = np.reshape(y, [-1])
+    z = np.reshape(z, [-1])
+    xyz = np.stack([x,y,z], axis=1).astype(np.float32)
+    return xyz
+
+def gridcloud2d(Y, X):
+    x_ = np.linspace(0, X-1, X)
+    y_ = np.linspace(0, Y-1, Y)
+    y, x = np.meshgrid(y_, x_, indexing='ij')
+    x = np.reshape(x, [-1])
+    y = np.reshape(y, [-1])
+    xy = np.stack([x,y], axis=1).astype(np.float32)
+    return xy
+
+def normalize(im):
+    im = im - np.min(im)
+    im = im / np.max(im)
+    return im
+
+def wrap2pi(rad_angle):
+    # rad_angle can be any shape
+    # puts the angle into the range [-pi, pi]
+    return np.arctan2(np.sin(rad_angle), np.cos(rad_angle))
+
+def convert_occ_to_height(occ):
+    Z, Y, X, C = occ.shape
+    assert(C==1)
+    
+    height = np.linspace(float(Y), 1.0, Y)
+    height = np.reshape(height, [1, Y, 1, 1])
+    height = np.max(occ*height, axis=1)/float(Y)
+    height = np.reshape(height, [Z, X, C])
+    return height
+
+def create_depth_image(xy, Z, H, W):
+
+    # turn the xy coordinates into image inds
+    xy = np.round(xy)
+
+    # lidar reports a sphere of measurements
+    # only use the inds that are within the image bounds
+    # also, only use forward-pointing depths (Z > 0)
+    valid = (xy[:,0] < W-1) & (xy[:,1] < H-1) & (xy[:,0] >= 0) & (xy[:,1] >= 0) & (Z[:] > 0)
+
+    # gather these up
+    xy = xy[valid]
+    Z = Z[valid]
+    
+    inds = sub2ind(H,W,xy[:,1],xy[:,0])
+    depth = np.zeros((H*W), np.float32)
+
+    for (index, replacement) in zip(inds, Z):
+        depth[index] = replacement
+    depth[np.where(depth == 0.0)] = 70.0
+    depth = np.reshape(depth, [H, W])
+
+    return depth
+
+def vis_depth(depth, maxdepth=80.0, log_vis=True):
+    depth[depth<=0.0] = maxdepth
+    if log_vis:
+        depth = np.log(depth)
+        depth = np.clip(depth, 0, np.log(maxdepth))
+    else:
+        depth = np.clip(depth, 0, maxdepth)
+    depth = (depth*255.0).astype(np.uint8)
+    return depth
+
+def preprocess_color(x):
+    return x.astype(np.float32) * 1./255 - 0.5
+
+def convert_box_to_ref_T_obj(boxes):
+    shape = boxes.shape
+    boxes = boxes.reshape(-1,9)
+    rots = [eul2rotm(rx,ry,rz)
+            for rx,ry,rz in boxes[:,6:]]
+    rots = np.stack(rots,axis=0)
+    trans = boxes[:,:3]
+    ref_T_objs = [merge_rt(rot,tran)
+                  for rot,tran in zip(rots,trans)]
+    ref_T_objs = np.stack(ref_T_objs,axis=0)
+    ref_T_objs = ref_T_objs.reshape(shape[:-1]+(4,4))
+    ref_T_objs = ref_T_objs.astype(np.float32)
+    return ref_T_objs
+
+def get_rot_from_delta(delta, yaw_only=False):
+    dx = delta[:,0]
+    dy = delta[:,1]
+    dz = delta[:,2]
+
+    bot_hyp = np.sqrt(dz**2 + dx**2)
+    # top_hyp = np.sqrt(bot_hyp**2 + dy**2)
+
+    pitch = -np.arctan2(dy, bot_hyp)
+    yaw = np.arctan2(dz, dx)
+
+    if yaw_only:
+        rot = [eul2rotm(0,y,0) for y in yaw]
+    else:
+        rot = [eul2rotm(0,y,p) for (p,y) in zip(pitch,yaw)]
+        
+    rot = np.stack(rot)
+    # rot is B x 3 x 3
+    return rot
+
+def im2col(im, psize):
+    n_channels = 1 if len(im.shape) == 2 else im.shape[0]
+    (n_channels, rows, cols) = (1,) * (3 - len(im.shape)) + im.shape
+
+    im_pad = np.zeros((n_channels,
+                       int(math.ceil(1.0 * rows / psize) * psize),
+                       int(math.ceil(1.0 * cols / psize) * psize)))
+    im_pad[:, 0:rows, 0:cols] = im
+
+    final = np.zeros((im_pad.shape[1], im_pad.shape[2], n_channels,
+                      psize, psize))
+    for c in np.arange(n_channels):
+        for x in np.arange(psize):
+            for y in np.arange(psize):
+                im_shift = np.vstack(
+                    (im_pad[c, x:], im_pad[c, :x]))
+                im_shift = np.column_stack(
+                    (im_shift[:, y:], im_shift[:, :y]))
+                final[x::psize, y::psize, c] = np.swapaxes(
+                    im_shift.reshape(int(im_pad.shape[1] / psize), psize,
+                                     int(im_pad.shape[2] / psize), psize), 1, 2)
+
+    return np.squeeze(final[0:rows - psize + 1, 0:cols - psize + 1])
+
+def filter_discontinuities(depth, filter_size=9, thresh=10):
+    H, W = list(depth.shape)
+
+    # Ensure that filter sizes are okay
+    assert filter_size % 2 == 1, "Can only use odd filter sizes."
+
+    # Compute discontinuities
+    offset = int((filter_size - 1) / 2)
+    patches = 1.0 * im2col(depth, filter_size)
+    mids = patches[:, :, offset, offset]
+    mins = np.min(patches, axis=(2, 3))
+    maxes = np.max(patches, axis=(2, 3))
+
+    discont = np.maximum(np.abs(mins - mids),
+                         np.abs(maxes - mids))
+    mark = discont > thresh
+
+    # Account for offsets
+    final_mark = np.zeros((H, W), dtype=np.uint16)
+    final_mark[offset:offset + mark.shape[0],
+               offset:offset + mark.shape[1]] = mark
+
+    return depth * (1 - final_mark)
+
+def argmax2d(tensor):
+    Y, X = list(tensor.shape)
+    # flatten the Tensor along the height and width axes
+    flat_tensor = tensor.reshape(-1)
+    # argmax of the flat tensor
+    argmax = np.argmax(flat_tensor)
+    
+    # convert the indices into 2d coordinates
+    argmax_y = argmax // X # row
+    argmax_x = argmax % X # col
+
+    return argmax_y, argmax_x
+
+def plot_traj_3d(traj):
+    # traj is S x 3
+    
+    # print('traj', traj.shape)
+    S, C = list(traj.shape)
+    assert(C==3)
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection='3d')
+
+    colors = [plt.cm.RdYlBu(i) for i in np.linspace(0,1,S)]
+    # print('colors', colors)
+    
+    xs = traj[:,0]
+    ys = -traj[:,1]
+    zs = traj[:,2]
+
+    ax.scatter(xs, zs, ys, s=30, c=colors, marker='o', alpha=1.0, edgecolors=(0,0,0))#, color=color_map[n])
+        
+    ax.set_xlabel('X')
+    ax.set_ylabel('Z')
+    ax.set_zlabel('Y')
+
+    ax.set_xlim(0,1)
+    ax.set_ylim(0,1) # this is really Z
+    ax.set_zlim(-1,0) # this is really Y
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    image = np.array(Image.open(buf)) # H x W x 4
+    image = image[:,:,:3]
+
+    plt.close()
+    return image
+
+def camera2pixels(xyz, pix_T_cam):
+    # xyz is shaped N x 3
+    # returns xy, shaped N x 2
+    
+    fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
+    x, y, z = xyz[:,0], xyz[:,1], xyz[:,2]
+
+    EPS = 1e-4
+    z = np.clip(z, EPS, None)
+    x = (x*fx)/z + x0
+    y = (y*fy)/z + y0
+    xy = np.stack([x, y], axis=-1)
+    return xy
+
+def make_colorwheel():
+    """
+    Generates a color wheel for optical flow visualization as presented in:
+        Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
+        URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf
+
+    Code follows the original C++ source code of Daniel Scharstein.
+    Code follows the the Matlab source code of Deqing Sun.
+
+    Returns:
+        np.ndarray: Color wheel
+    """
+
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = np.zeros((ncols, 3))
+    col = 0
+
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.floor(255*np.arange(0,RY)/RY)
+    col = col+RY
+    # YG
+    colorwheel[col:col+YG, 0] = 255 - np.floor(255*np.arange(0,YG)/YG)
+    colorwheel[col:col+YG, 1] = 255
+    col = col+YG
+    # GC
+    colorwheel[col:col+GC, 1] = 255
+    colorwheel[col:col+GC, 2] = np.floor(255*np.arange(0,GC)/GC)
+    col = col+GC
+    # CB
+    colorwheel[col:col+CB, 1] = 255 - np.floor(255*np.arange(CB)/CB)
+    colorwheel[col:col+CB, 2] = 255
+    col = col+CB
+    # BM
+    colorwheel[col:col+BM, 2] = 255
+    colorwheel[col:col+BM, 0] = np.floor(255*np.arange(0,BM)/BM)
+    col = col+BM
+    # MR
+    colorwheel[col:col+MR, 2] = 255 - np.floor(255*np.arange(MR)/MR)
+    colorwheel[col:col+MR, 0] = 255
+    return colorwheel
+
+def flow_uv_to_colors(u, v, convert_to_bgr=False):
+    """
+    Applies the flow color wheel to (possibly clipped) flow components u and v.
+
+    According to the C++ source code of Daniel Scharstein
+    According to the Matlab source code of Deqing Sun
+
+    Args:
+        u (np.ndarray): Input horizontal flow of shape [H,W]
+        v (np.ndarray): Input vertical flow of shape [H,W]
+        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
+
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
+    colorwheel = make_colorwheel()  # shape [55x3]
+    ncols = colorwheel.shape[0]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    a = np.arctan2(-v, -u)/np.pi
+    fk = (a+1) / 2*(ncols-1)
+    k0 = np.floor(fk).astype(np.int32)
+    k1 = k0 + 1
+    k1[k1 == ncols] = 0
+    f = fk - k0
+    for i in range(colorwheel.shape[1]):
+        tmp = colorwheel[:,i]
+        col0 = tmp[k0] / 255.0
+        col1 = tmp[k1] / 255.0
+        col = (1-f)*col0 + f*col1
+        idx = (rad <= 1)
+        col[idx]  = 1 - rad[idx] * (1-col[idx])
+        col[~idx] = col[~idx] * 0.75   # out of range
+        # Note the 2-i => BGR instead of RGB
+        ch_idx = 2-i if convert_to_bgr else i
+        flow_image[:,:,ch_idx] = np.floor(255 * col)
+    return flow_image
+
+def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False):
+    """
+    Expects a two dimensional flow image of shape.
+
+    Args:
+        flow_uv (np.ndarray): Flow UV image of shape [H,W,2]
+        clip_flow (float, optional): Clip maximum of flow values. Defaults to None.
+        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
+
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    assert flow_uv.ndim == 3, 'input flow must have three dimensions'
+    assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
+    if clip_flow is not None:
+        flow_uv = np.clip(flow_uv, -clip_flow, clip_flow) / clip_flow
+        # flow_uv = np.clamp(flow, -clip, clip)/clip
+        
+    u = flow_uv[:,:,0]
+    v = flow_uv[:,:,1]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    rad_max = np.max(rad)
+    epsilon = 1e-5
+    u = u / (rad_max + epsilon)
+    v = v / (rad_max + epsilon)
+    return flow_uv_to_colors(u, v, convert_to_bgr)

utils/samp.py

@@ -0,0 +1,213 @@
+import torch
+import utils.basic
+import torch.nn.functional as F
+
+def bilinear_sampler(input, coords, align_corners=True, padding_mode="border"):
+    r"""Sample a tensor using bilinear interpolation
+
+    `bilinear_sampler(input, coords)` samples a tensor :attr:`input` at
+    coordinates :attr:`coords` using bilinear interpolation. It is the same
+    as `torch.nn.functional.grid_sample()` but with a different coordinate
+    convention.
+
+    The input tensor is assumed to be of shape :math:`(B, C, H, W)`, where
+    :math:`B` is the batch size, :math:`C` is the number of channels,
+    :math:`H` is the height of the image, and :math:`W` is the width of the
+    image. The tensor :attr:`coords` of shape :math:`(B, H_o, W_o, 2)` is
+    interpreted as an array of 2D point coordinates :math:`(x_i,y_i)`.
+
+    Alternatively, the input tensor can be of size :math:`(B, C, T, H, W)`,
+    in which case sample points are triplets :math:`(t_i,x_i,y_i)`. Note
+    that in this case the order of the components is slightly different
+    from `grid_sample()`, which would expect :math:`(x_i,y_i,t_i)`.
+
+    If `align_corners` is `True`, the coordinate :math:`x` is assumed to be
+    in the range :math:`[0,W-1]`, with 0 corresponding to the center of the
+    left-most image pixel :math:`W-1` to the center of the right-most
+    pixel.
+
+    If `align_corners` is `False`, the coordinate :math:`x` is assumed to
+    be in the range :math:`[0,W]`, with 0 corresponding to the left edge of
+    the left-most pixel :math:`W` to the right edge of the right-most
+    pixel.
+
+    Similar conventions apply to the :math:`y` for the range
+    :math:`[0,H-1]` and :math:`[0,H]` and to :math:`t` for the range
+    :math:`[0,T-1]` and :math:`[0,T]`.
+
+    Args:
+        input (Tensor): batch of input images.
+        coords (Tensor): batch of coordinates.
+        align_corners (bool, optional): Coordinate convention. Defaults to `True`.
+        padding_mode (str, optional): Padding mode. Defaults to `"border"`.
+
+    Returns:
+        Tensor: sampled points.
+    """
+
+    sizes = input.shape[2:]
+
+    assert len(sizes) in [2, 3]
+
+    if len(sizes) == 3:
+        # t x y -> x y t to match dimensions T H W in grid_sample
+        coords = coords[..., [1, 2, 0]]
+
+    if align_corners:
+        coords = coords * torch.tensor(
+            [2 / max(size - 1, 1) for size in reversed(sizes)], device=coords.device
+        )
+    else:
+        coords = coords * torch.tensor(
+            [2 / size for size in reversed(sizes)], device=coords.device
+        )
+
+    coords -= 1
+
+    return F.grid_sample(
+        input, coords, align_corners=align_corners, padding_mode=padding_mode
+    )
+
+
+def sample_features4d(input, coords):
+    r"""Sample spatial features
+
+    `sample_features4d(input, coords)` samples the spatial features
+    :attr:`input` represented by a 4D tensor :math:`(B, C, H, W)`.
+
+    The field is sampled at coordinates :attr:`coords` using bilinear
+    interpolation. :attr:`coords` is assumed to be of shape :math:`(B, R,
+    3)`, where each sample has the format :math:`(x_i, y_i)`. This uses the
+    same convention as :func:`bilinear_sampler` with `align_corners=True`.
+
+    The output tensor has one feature per point, and has shape :math:`(B,
+    R, C)`.
+
+    Args:
+        input (Tensor): spatial features.
+        coords (Tensor): points.
+
+    Returns:
+        Tensor: sampled features.
+    """
+
+    B, _, _, _ = input.shape
+
+    # B R 2 -> B R 1 2
+    coords = coords.unsqueeze(2)
+
+    # B C R 1
+    feats = bilinear_sampler(input, coords)
+
+    return feats.permute(0, 2, 1, 3).view(
+        B, -1, feats.shape[1] * feats.shape[3]
+    )  # B C R 1 -> B R C
+
+
+def sample_features5d(input, coords):
+    r"""Sample spatio-temporal features
+
+    `sample_features5d(input, coords)` works in the same way as
+    :func:`sample_features4d` but for spatio-temporal features and points:
+    :attr:`input` is a 5D tensor :math:`(B, T, C, H, W)`, :attr:`coords` is
+    a :math:`(B, R1, R2, 3)` tensor of spatio-temporal point :math:`(t_i,
+    x_i, y_i)`. The output tensor has shape :math:`(B, R1, R2, C)`.
+
+    Args:
+        input (Tensor): spatio-temporal features.
+        coords (Tensor): spatio-temporal points.
+
+    Returns:
+        Tensor: sampled features.
+    """
+
+    B, T, _, _, _ = input.shape
+
+    # B T C H W -> B C T H W
+    input = input.permute(0, 2, 1, 3, 4)
+
+    # B R1 R2 3 -> B R1 R2 1 3
+    coords = coords.unsqueeze(3)
+
+    # B C R1 R2 1
+    feats = bilinear_sampler(input, coords)
+
+    return feats.permute(0, 2, 3, 1, 4).view(
+        B, feats.shape[2], feats.shape[3], feats.shape[1]
+    )  # B C R1 R2 1 -> B R1 R2 C
+
+
+def bilinear_sample2d(im, x, y, return_inbounds=False):
+    # x and y are each B, N
+    # output is B, C, N
+    B, C, H, W = list(im.shape)
+    N = list(x.shape)[1]
+
+    x = x.float()
+    y = y.float()
+    H_f = torch.tensor(H, dtype=torch.float32)
+    W_f = torch.tensor(W, dtype=torch.float32)
+    
+    # inbound_mask = (x>-0.5).float()*(y>-0.5).float()*(x<W_f+0.5).float()*(y<H_f+0.5).float()
+
+    max_y = (H_f - 1).int()
+    max_x = (W_f - 1).int()
+
+    x0 = torch.floor(x).int()
+    x1 = x0 + 1
+    y0 = torch.floor(y).int()
+    y1 = y0 + 1
+    
+    x0_clip = torch.clamp(x0, 0, max_x)
+    x1_clip = torch.clamp(x1, 0, max_x)
+    y0_clip = torch.clamp(y0, 0, max_y)
+    y1_clip = torch.clamp(y1, 0, max_y)
+    dim2 = W
+    dim1 = W * H
+
+    base = torch.arange(0, B, dtype=torch.int64, device=x.device)*dim1
+    base = torch.reshape(base, [B, 1]).repeat([1, N])
+
+    base_y0 = base + y0_clip * dim2
+    base_y1 = base + y1_clip * dim2
+
+    idx_y0_x0 = base_y0 + x0_clip
+    idx_y0_x1 = base_y0 + x1_clip
+    idx_y1_x0 = base_y1 + x0_clip
+    idx_y1_x1 = base_y1 + x1_clip
+
+    # use the indices to lookup pixels in the flat image
+    # im is B x C x H x W
+    # move C out to last dim
+    im_flat = (im.permute(0, 2, 3, 1)).reshape(B*H*W, C)
+    i_y0_x0 = im_flat[idx_y0_x0.long()]
+    i_y0_x1 = im_flat[idx_y0_x1.long()]
+    i_y1_x0 = im_flat[idx_y1_x0.long()]
+    i_y1_x1 = im_flat[idx_y1_x1.long()]
+
+    # Finally calculate interpolated values.
+    x0_f = x0.float()
+    x1_f = x1.float()
+    y0_f = y0.float()
+    y1_f = y1.float()
+
+    w_y0_x0 = ((x1_f - x) * (y1_f - y)).unsqueeze(2)
+    w_y0_x1 = ((x - x0_f) * (y1_f - y)).unsqueeze(2)
+    w_y1_x0 = ((x1_f - x) * (y - y0_f)).unsqueeze(2)
+    w_y1_x1 = ((x - x0_f) * (y - y0_f)).unsqueeze(2)
+
+    output = w_y0_x0 * i_y0_x0 + w_y0_x1 * i_y0_x1 + \
+             w_y1_x0 * i_y1_x0 + w_y1_x1 * i_y1_x1
+    # output is B*N x C
+    output = output.view(B, -1, C)
+    output = output.permute(0, 2, 1)
+    # output is B x C x N
+
+    if return_inbounds:
+        x_valid = (x > -0.5).byte() & (x < float(W_f - 0.5)).byte()
+        y_valid = (y > -0.5).byte() & (y < float(H_f - 0.5)).byte()
+        inbounds = (x_valid & y_valid).float()
+        inbounds = inbounds.reshape(B, N) # something seems wrong here for B>1; i'm getting an error here (or downstream if i put -1)
+        return output, inbounds
+
+    return output # B, C, N

utils/saveload.py

@@ -0,0 +1,59 @@
+import pathlib
+import os
+import torch
+
+def save(ckpt_dir, module, optimizer, scheduler, global_step, keep_latest=2, model_name='model'):
+    pathlib.Path(ckpt_dir).mkdir(exist_ok=True, parents=True)
+    prev_ckpts = list(pathlib.Path(ckpt_dir).glob('%s-*pth' % model_name))
+    prev_ckpts.sort(key=lambda p: p.stat().st_mtime,reverse=True)
+    if len(prev_ckpts) > keep_latest-1:
+        for f in prev_ckpts[keep_latest-1:]:
+            f.unlink()
+    save_path = '%s/%s-%09d.pth' % (ckpt_dir, model_name, global_step)
+    save_dict = {
+        "model": module.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "global_step": global_step,
+    }
+    if scheduler is not None:
+        save_dict['scheduler'] = scheduler.state_dict()
+    print(f"saving {save_path}")
+    torch.save(save_dict, save_path)
+    return False
+
+def load(fabric, ckpt_dir, model, optimizer=None, scheduler=None, model_ema=None, step=0, model_name='model', ignore_load=None, strict=True, verbose=True, weights_only=False):
+    if verbose:
+        print('reading ckpt from %s' % ckpt_dir)
+    if not os.path.exists(ckpt_dir):
+        print('...there is no full checkpoint in %s' % ckpt_dir)
+        print('-- note this function no longer appends "saved_checkpoints/" before the ckpt_dir --')
+        assert(False)
+    else:
+        prev_ckpts = list(pathlib.Path(ckpt_dir).glob('%s-*pth' % model_name))
+        prev_ckpts.sort(key=lambda p: p.stat().st_mtime,reverse=True)
+        if len(prev_ckpts):
+            path = prev_ckpts[0]
+            # e.g., './checkpoints/2Ai4_5e-4_base18_1539/model-000050000.pth'
+            step = int(str(path).split('-')[-1].split('.')[0])
+            if verbose:
+                print('...found checkpoint %s; (parsed step %d from path)' % (path, step))
+            if fabric is not None:
+                checkpoint = fabric.load(path)
+            else:
+                checkpoint = torch.load(path, weights_only=weights_only)
+            if optimizer is not None:
+                optimizer.load_state_dict(checkpoint['optimizer'])
+            if scheduler is not None:
+                scheduler.load_state_dict(checkpoint['scheduler'])
+            assert ignore_load is None # not ready yet
+            if 'model' in checkpoint:
+                state_dict = checkpoint['model']
+            else:
+                state_dict = checkpoint
+            model.load_state_dict(state_dict, strict=strict)
+        else:
+            print('...there is no full checkpoint here!')
+    return step
+                        
+                                            
+                                                                                                                                                                                    

utils/test.py

@@ -0,0 +1,194 @@
+import torch
+import cv2
+import torch.nn.functional as F
+import numpy as np
+
+def prep_frame_for_dino(img, scale_size=[192]):
+    """
+    read a single frame & preprocess
+    """
+    ori_h, ori_w, _ = img.shape
+    if len(scale_size) == 1:
+        if(ori_h > ori_w):
+            tw = scale_size[0]
+            th = (tw * ori_h) / ori_w
+            th = int((th // 64) * 64)
+        else:
+            th = scale_size[0]
+            tw = (th * ori_w) / ori_h
+            tw = int((tw // 64) * 64)
+    else:
+        th, tw = scale_size
+    img = cv2.resize(img, (tw, th))
+    img = img.astype(np.float32)
+    img = img / 255.0
+    img = img[:, :, ::-1]
+    img = np.transpose(img.copy(), (2, 0, 1))
+    img = torch.from_numpy(img).float()
+
+    def color_normalize(x, mean=[0.485, 0.456, 0.406], std=[0.228, 0.224, 0.225]):
+        for t, m, s in zip(x, mean, std):
+            t.sub_(m)
+            t.div_(s)
+        return x
+    
+    img = color_normalize(img)
+    return img, ori_h, ori_w
+
+def get_feats_from_dino(model, frame):
+    # batch version of the other func
+    B = frame.shape[0]
+    patch_size = model.patch_embed.patch_size
+    h, w = int(frame.shape[2] / patch_size), int(frame.shape[3] / patch_size)
+    out = model.get_intermediate_layers(frame.cuda(), n=1)[0] # B, 1+h*w, dim
+    dim = out.shape[-1]
+    out = out[:, 1:, :]  # discard the [CLS] token
+    outmap = out.permute(0, 2, 1).reshape(B, dim, h, w)
+    return out, outmap, h, w
+
+def restrict_neighborhood(h, w):
+    size_mask_neighborhood = 12
+    # We restrict the set of source nodes considered to a spatial neighborhood of the query node (i.e. ``local attention'')
+    mask = torch.zeros(h, w, h, w)
+    for i in range(h):
+        for j in range(w):
+            for p in range(2 * size_mask_neighborhood + 1):
+                for q in range(2 * size_mask_neighborhood + 1):
+                    if i - size_mask_neighborhood + p < 0 or i - size_mask_neighborhood + p >= h:
+                        continue
+                    if j - size_mask_neighborhood + q < 0 or j - size_mask_neighborhood + q >= w:
+                        continue
+                    mask[i, j, i - size_mask_neighborhood + p, j - size_mask_neighborhood + q] = 1
+
+    mask = mask.reshape(h * w, h * w)
+    return mask.cuda(non_blocking=True)
+
+def label_propagation(h, w, feat_tar, list_frame_feats, list_segs, mask_neighborhood=None):
+    ncontext = len(list_frame_feats)
+    feat_sources = torch.stack(list_frame_feats) # nmb_context x dim x h*w
+
+    feat_tar = F.normalize(feat_tar, dim=1, p=2)
+    feat_sources = F.normalize(feat_sources, dim=1, p=2)
+
+    # print('feat_tar', feat_tar.shape)
+    # print('feat_sources', feat_sources.shape)
+
+    feat_tar = feat_tar.unsqueeze(0).repeat(ncontext, 1, 1)
+    aff = torch.exp(torch.bmm(feat_tar, feat_sources) / 0.1)
+
+    size_mask_neighborhood = 12
+    if size_mask_neighborhood > 0:
+        if mask_neighborhood is None:
+            mask_neighborhood = restrict_neighborhood(h, w)
+            mask_neighborhood = mask_neighborhood.unsqueeze(0).repeat(ncontext, 1, 1)
+        aff *= mask_neighborhood
+
+    aff = aff.transpose(2, 1).reshape(-1, h*w) # nmb_context*h*w (source: keys) x h*w (tar: queries)
+    topk = 5
+    tk_val, _ = torch.topk(aff, dim=0, k=topk)
+    tk_val_min, _ = torch.min(tk_val, dim=0)
+    aff[aff < tk_val_min] = 0
+
+    aff = aff / torch.sum(aff, keepdim=True, axis=0)
+
+    list_segs = [s.cuda() for s in list_segs]
+    segs = torch.cat(list_segs)
+    nmb_context, C, h, w = segs.shape
+    segs = segs.reshape(nmb_context, C, -1).transpose(2, 1).reshape(-1, C).T # C x nmb_context*h*w
+    seg_tar = torch.mm(segs, aff)
+    seg_tar = seg_tar.reshape(1, C, h, w)
+    
+    return seg_tar, mask_neighborhood
+
+def norm_mask(mask):
+    c, h, w = mask.size()
+    for cnt in range(c):
+        mask_cnt = mask[cnt,:,:]
+        if(mask_cnt.max() > 0):
+            mask_cnt = (mask_cnt - mask_cnt.min())
+            mask_cnt = mask_cnt/mask_cnt.max()
+            mask[cnt,:,:] = mask_cnt
+    return mask
+
+
+def get_dino_output(dino, rgbs, trajs_g, vis_g):
+    B, S, C, H, W = rgbs.shape
+
+    B1, S1, N, D = trajs_g.shape
+    assert(B1==B)
+    assert(S1==S)
+    assert(D==2)
+    
+    assert(B==1)
+    xy0 = trajs_g[:,0] # B, N, 2
+
+    # The queue stores the n preceeding frames
+    import queue
+    import copy
+    n_last_frames = 7
+    que = queue.Queue(n_last_frames)
+
+    # run dino
+    prep_rgbs = []
+    for s in range(S):
+        prep_rgb, ori_h, ori_w = prep_frame_for_dino(rgbs[0, s].permute(1,2,0).detach().cpu().numpy(), scale_size=[H])
+        prep_rgbs.append(prep_rgb)
+    prep_rgbs = torch.stack(prep_rgbs, dim=0) # S, 3, H, W
+    with torch.no_grad():
+        bs = 8
+        idx = 0 
+        featmaps = []
+        while idx < S:
+            end_id = min(S, idx+bs)
+            _, featmaps_cur, h, w = get_feats_from_dino(dino, prep_rgbs[idx:end_id]) # S, C, h, w
+            idx = end_id
+            featmaps.append(featmaps_cur)
+        featmaps = torch.cat(featmaps, dim=0)
+    C = featmaps.shape[1]
+    featmaps = featmaps.unsqueeze(0) # 1, S, C, h, w
+    # featmaps = F.normalize(featmaps, dim=2, p=2)
+
+    xy0 = trajs_g[:, 0, :] # B, N, 2
+    patch_size = dino.patch_embed.patch_size
+    first_seg = torch.zeros((1, N, H//patch_size, W//patch_size))
+    for n in range(N):
+        first_seg[0, n, (xy0[0, n, 1]/patch_size).long(), (xy0[0, n, 0]/patch_size).long()] = 1
+
+    frame1_feat = featmaps[0, 0].reshape(C, h*w) # dim x h*w
+    mask_neighborhood = None
+    accs = []
+    trajs_e = torch.zeros_like(trajs_g)
+    trajs_e[0,0] = trajs_g[0,0]
+    for cnt in range(1, S):
+        used_frame_feats = [frame1_feat] + [pair[0] for pair in list(que.queue)]
+        used_segs = [first_seg] + [pair[1] for pair in list(que.queue)]
+
+        feat_tar = featmaps[0, cnt].reshape(C, h*w)
+
+        frame_tar_avg, mask_neighborhood = label_propagation(h, w, feat_tar.T, used_frame_feats, used_segs, mask_neighborhood)
+
+        # pop out oldest frame if neccessary
+        if que.qsize() == n_last_frames:
+            que.get()
+        # push current results into queue
+        seg = copy.deepcopy(frame_tar_avg)
+        que.put([feat_tar, seg])
+
+        # upsampling & argmax
+        frame_tar_avg = F.interpolate(frame_tar_avg, scale_factor=patch_size, mode='bilinear', align_corners=False, recompute_scale_factor=False)[0]
+        frame_tar_avg = norm_mask(frame_tar_avg)
+        _, frame_tar_seg = torch.max(frame_tar_avg, dim=0)
+
+        for n in range(N):
+            vis = vis_g[0,cnt,n]
+            if len(torch.nonzero(frame_tar_avg[n])) > 0:
+                # weighted average
+                nz = torch.nonzero(frame_tar_avg[n])
+                coord_e = torch.sum(frame_tar_avg[n][nz[:,0], nz[:,1]].reshape(-1,1) * nz.float(), 0) / frame_tar_avg[n][nz[:,0], nz[:,1]].sum() # 2
+                coord_e = coord_e[[1,0]]
+            else:
+                # stay where it was
+                coord_e = trajs_e[0,cnt-1,n]
+                
+            trajs_e[0, cnt, n] = coord_e
+    return trajs_e

utils/visualizer.py

@@ -0,0 +1,365 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import numpy as np
+import imageio
+import torch
+
+from matplotlib import cm
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import matplotlib.pyplot as plt
+from PIL import Image, ImageDraw
+
+
+def read_video_from_path(path):
+    try:
+        reader = imageio.get_reader(path)
+    except Exception as e:
+        print("Error opening video file: ", e)
+        return None
+    frames = []
+    for i, im in enumerate(reader):
+        frames.append(np.array(im))
+    return np.stack(frames)
+
+
+def draw_circle(rgb, coord, radius, color=(255, 0, 0), visible=True, color_alpha=None):
+    # Create a draw object
+    draw = ImageDraw.Draw(rgb)
+    # Calculate the bounding box of the circle
+    left_up_point = (coord[0] - radius, coord[1] - radius)
+    right_down_point = (coord[0] + radius, coord[1] + radius)
+    # Draw the circle
+    color = tuple(list(color) + [color_alpha if color_alpha is not None else 255])
+
+    draw.ellipse(
+        [left_up_point, right_down_point],
+        fill=tuple(color) if visible else None,
+        outline=tuple(color),
+    )
+    return rgb
+
+
+def draw_line(rgb, coord_y, coord_x, color, linewidth):
+    draw = ImageDraw.Draw(rgb)
+    draw.line(
+        (coord_y[0], coord_y[1], coord_x[0], coord_x[1]),
+        fill=tuple(color),
+        width=linewidth,
+    )
+    return rgb
+
+
+def add_weighted(rgb, alpha, original, beta, gamma):
+    return (rgb * alpha + original * beta + gamma).astype("uint8")
+
+
+class Visualizer:
+    def __init__(
+        self,
+        save_dir: str = "./results",
+        grayscale: bool = False,
+        pad_value: int = 0,
+        fps: int = 10,
+        mode: str = "rainbow",  # 'cool', 'optical_flow'
+        linewidth: int = 2,
+        show_first_frame: int = 10,
+        tracks_leave_trace: int = 0,  # -1 for infinite
+    ):
+        self.mode = mode
+        self.save_dir = save_dir
+        if mode == "rainbow":
+            self.color_map = cm.get_cmap("gist_rainbow")
+        elif mode == "cool":
+            self.color_map = cm.get_cmap(mode)
+        self.show_first_frame = show_first_frame
+        self.grayscale = grayscale
+        self.tracks_leave_trace = tracks_leave_trace
+        self.pad_value = pad_value
+        self.linewidth = linewidth
+        self.fps = fps
+
+    def visualize(
+        self,
+        video: torch.Tensor,  # (B,T,C,H,W)
+        tracks: torch.Tensor,  # (B,T,N,2)
+        visibility: torch.Tensor = None,  # (B, T, N, 1) bool
+        gt_tracks: torch.Tensor = None,  # (B,T,N,2)
+        segm_mask: torch.Tensor = None,  # (B,1,H,W)
+        filename: str = "video",
+        writer=None,  # tensorboard Summary Writer, used for visualization during training
+        step: int = 0,
+        query_frame=0,
+        save_video: bool = True,
+        compensate_for_camera_motion: bool = False,
+        opacity: float = 1.0,
+    ):
+        if compensate_for_camera_motion:
+            assert segm_mask is not None
+        if segm_mask is not None:
+            coords = tracks[0, query_frame].round().long()
+            segm_mask = segm_mask[0, query_frame][coords[:, 1], coords[:, 0]].long()
+            if (segm_mask <= 0).all():
+                segm_mask = None
+
+        video = F.pad(
+            video,
+            (self.pad_value, self.pad_value, self.pad_value, self.pad_value),
+            "constant",
+            255,
+        )
+        color_alpha = int(opacity * 255)
+        tracks = tracks + self.pad_value
+
+        if self.grayscale:
+            transform = transforms.Grayscale()
+            video = transform(video)
+            video = video.repeat(1, 1, 3, 1, 1)
+
+        res_video = self.draw_tracks_on_video(
+            video=video,
+            tracks=tracks,
+            visibility=visibility,
+            segm_mask=segm_mask,
+            gt_tracks=gt_tracks,
+            query_frame=query_frame,
+            compensate_for_camera_motion=compensate_for_camera_motion,
+            color_alpha=color_alpha,
+        )
+        if save_video:
+            self.save_video(res_video, filename=filename, writer=writer, step=step)
+        return res_video
+
+    def save_video(self, video, filename, writer=None, step=0):
+        if writer is not None:
+            writer.add_video(
+                filename,
+                video.to(torch.uint8),
+                global_step=step,
+                fps=self.fps,
+            )
+        else:
+            os.makedirs(self.save_dir, exist_ok=True)
+            wide_list = list(video.unbind(1))
+            wide_list = [wide[0].permute(1, 2, 0).cpu().numpy() for wide in wide_list]
+
+            # Prepare the video file path
+            save_path = os.path.join(self.save_dir, f"{filename}.mp4")
+
+            # Create a writer object
+            video_writer = imageio.get_writer(save_path, fps=self.fps)
+
+            # Write frames to the video file
+            for frame in wide_list[2:-1]:
+                video_writer.append_data(frame)
+
+            video_writer.close()
+
+            print(f"Video saved to {save_path}")
+
+    def draw_tracks_on_video(
+        self,
+        video: torch.Tensor,
+        tracks: torch.Tensor,
+        visibility: torch.Tensor = None,
+        segm_mask: torch.Tensor = None,
+        gt_tracks=None,
+        query_frame=0,
+        compensate_for_camera_motion=False,
+        color_alpha: int = 255,
+    ):
+        B, T, C, H, W = video.shape
+        _, _, N, D = tracks.shape
+
+        assert D == 2
+        assert C == 3
+        video = video[0].permute(0, 2, 3, 1).byte().detach().cpu().numpy()  # S, H, W, C
+        tracks = tracks[0].long().detach().cpu().numpy()  # S, N, 2
+        if gt_tracks is not None:
+            gt_tracks = gt_tracks[0].detach().cpu().numpy()
+
+        res_video = []
+
+        # process input video
+        for rgb in video:
+            res_video.append(rgb.copy())
+        vector_colors = np.zeros((T, N, 3))
+
+        if self.mode == "optical_flow":
+            import flow_vis
+
+            vector_colors = flow_vis.flow_to_color(tracks - tracks[query_frame][None])
+        elif segm_mask is None:
+            if self.mode == "rainbow":
+                y_min, y_max = (
+                    tracks[query_frame, :, 1].min(),
+                    tracks[query_frame, :, 1].max(),
+                )
+                norm = plt.Normalize(y_min, y_max)
+                for n in range(N):
+                    if isinstance(query_frame, torch.Tensor):
+                        query_frame_ = query_frame[n]
+                    else:
+                        query_frame_ = query_frame
+                    color = self.color_map(norm(tracks[query_frame_, n, 1]))
+                    color = np.array(color[:3])[None] * 255
+                    vector_colors[:, n] = np.repeat(color, T, axis=0)
+            else:
+                # color changes with time
+                for t in range(T):
+                    color = np.array(self.color_map(t / T)[:3])[None] * 255
+                    vector_colors[t] = np.repeat(color, N, axis=0)
+        else:
+            if self.mode == "rainbow":
+                vector_colors[:, segm_mask <= 0, :] = 255
+
+                y_min, y_max = (
+                    tracks[0, segm_mask > 0, 1].min(),
+                    tracks[0, segm_mask > 0, 1].max(),
+                )
+                norm = plt.Normalize(y_min, y_max)
+                for n in range(N):
+                    if segm_mask[n] > 0:
+                        color = self.color_map(norm(tracks[0, n, 1]))
+                        color = np.array(color[:3])[None] * 255
+                        vector_colors[:, n] = np.repeat(color, T, axis=0)
+
+            else:
+                # color changes with segm class
+                segm_mask = segm_mask.cpu()
+                color = np.zeros((segm_mask.shape[0], 3), dtype=np.float32)
+                color[segm_mask > 0] = np.array(self.color_map(1.0)[:3]) * 255.0
+                color[segm_mask <= 0] = np.array(self.color_map(0.0)[:3]) * 255.0
+                vector_colors = np.repeat(color[None], T, axis=0)
+
+        #  draw tracks
+        if self.tracks_leave_trace != 0:
+            for t in range(query_frame + 1, T):
+                first_ind = (
+                    max(0, t - self.tracks_leave_trace)
+                    if self.tracks_leave_trace >= 0
+                    else 0
+                )
+                curr_tracks = tracks[first_ind : t + 1]
+                curr_colors = vector_colors[first_ind : t + 1]
+                if compensate_for_camera_motion:
+                    diff = (
+                        tracks[first_ind : t + 1, segm_mask <= 0]
+                        - tracks[t : t + 1, segm_mask <= 0]
+                    ).mean(1)[:, None]
+
+                    curr_tracks = curr_tracks - diff
+                    curr_tracks = curr_tracks[:, segm_mask > 0]
+                    curr_colors = curr_colors[:, segm_mask > 0]
+
+                res_video[t] = self._draw_pred_tracks(
+                    res_video[t],
+                    curr_tracks,
+                    curr_colors,
+                )
+                if gt_tracks is not None:
+                    res_video[t] = self._draw_gt_tracks(
+                        res_video[t], gt_tracks[first_ind : t + 1]
+                    )
+
+        #  draw points
+        for t in range(T):
+            img = Image.fromarray(np.uint8(res_video[t]))
+            for i in range(N):
+                coord = (tracks[t, i, 0], tracks[t, i, 1])
+                visibile = True
+                if visibility is not None:
+                    visibile = visibility[0, t, i]
+                if coord[0] != 0 and coord[1] != 0:
+                    if not compensate_for_camera_motion or (
+                        compensate_for_camera_motion and segm_mask[i] > 0
+                    ):
+                        img = draw_circle(
+                            img,
+                            coord=coord,
+                            radius=int(self.linewidth * 2),
+                            color=vector_colors[t, i].astype(int),
+                            visible=visibile,
+                            color_alpha=color_alpha,
+                        )
+            res_video[t] = np.array(img)
+
+        #  construct the final rgb sequence
+        if self.show_first_frame > 0:
+            res_video = [res_video[0]] * self.show_first_frame + res_video[1:]
+        return torch.from_numpy(np.stack(res_video)).permute(0, 3, 1, 2)[None].byte()
+
+    def _draw_pred_tracks(
+        self,
+        rgb: np.ndarray,  # H x W x 3
+        tracks: np.ndarray,  # T x 2
+        vector_colors: np.ndarray,
+        alpha: float = 0.5,
+    ):
+        T, N, _ = tracks.shape
+        rgb = Image.fromarray(np.uint8(rgb))
+        for s in range(T - 1):
+            vector_color = vector_colors[s]
+            original = rgb.copy()
+            alpha = (s / T) ** 2
+            for i in range(N):
+                coord_y = (int(tracks[s, i, 0]), int(tracks[s, i, 1]))
+                coord_x = (int(tracks[s + 1, i, 0]), int(tracks[s + 1, i, 1]))
+                if coord_y[0] != 0 and coord_y[1] != 0:
+                    rgb = draw_line(
+                        rgb,
+                        coord_y,
+                        coord_x,
+                        vector_color[i].astype(int),
+                        self.linewidth,
+                    )
+            if self.tracks_leave_trace > 0:
+                rgb = Image.fromarray(
+                    np.uint8(
+                        add_weighted(
+                            np.array(rgb), alpha, np.array(original), 1 - alpha, 0
+                        )
+                    )
+                )
+        rgb = np.array(rgb)
+        return rgb
+
+    def _draw_gt_tracks(
+        self,
+        rgb: np.ndarray,  # H x W x 3,
+        gt_tracks: np.ndarray,  # T x 2
+    ):
+        T, N, _ = gt_tracks.shape
+        color = np.array((211, 0, 0))
+        rgb = Image.fromarray(np.uint8(rgb))
+        for t in range(T):
+            for i in range(N):
+                gt_tracks = gt_tracks[t][i]
+                #  draw a red cross
+                if gt_tracks[0] > 0 and gt_tracks[1] > 0:
+                    length = self.linewidth * 3
+                    coord_y = (int(gt_tracks[0]) + length, int(gt_tracks[1]) + length)
+                    coord_x = (int(gt_tracks[0]) - length, int(gt_tracks[1]) - length)
+                    rgb = draw_line(
+                        rgb,
+                        coord_y,
+                        coord_x,
+                        color,
+                        self.linewidth,
+                    )
+                    coord_y = (int(gt_tracks[0]) - length, int(gt_tracks[1]) + length)
+                    coord_x = (int(gt_tracks[0]) + length, int(gt_tracks[1]) - length)
+                    rgb = draw_line(
+                        rgb,
+                        coord_y,
+                        coord_x,
+                        color,
+                        self.linewidth,
+                    )
+        rgb = np.array(rgb)
+        return rgb