PIPs++
import torch
import numpy as np
import os
import scipy.ndimage
import torchvision.transforms as transforms
import torch.nn.functional as F
from PIL import Image
import random
import glob
import json
import imageio
import cv2

class CrohdDataset(torch.utils.data.Dataset):
    def __init__(self,
                 dataset_location,
                 S=1000,
                 S_min=128):
        self.S = S
        
        dataset_dir = '%s/HT21' % dataset_location
        label_location = '%s/HT21Labels' % dataset_location
        subfolders = []

        dataset_dir = os.path.join(dataset_dir, "train")
        label_location = os.path.join(label_location, "train")
        subfolders = ['HT21-01', 'HT21-02', 'HT21-03', 'HT21-04']
        
        print("loading data from {0}".format(dataset_dir))
        # read gt for subfolders
        self.dataset_dir = dataset_dir
        # self.seqlen = seqlen
        self.subfolders = subfolders
        self.folder_to_gt = {} # key: folder name, value: dict with fields boxlist, scorelist, vislist
        # self.subfolder_lens = []

        self.all_folders = []
        self.all_start_frames = []
        self.all_end_frames = []
        self.all_dicts = []

        # we'll take each length-1k subseq as a video
        
        for fid, subfolder in enumerate(subfolders):
            print("loading labels for folder {0}/{1}".format(fid+1, len(subfolders)))
            label_path = os.path.join(dataset_dir, subfolder, 'gt/gt.txt')
            labels = np.loadtxt(label_path, delimiter=',')

            n_frames = int(labels[-1,0])
            # print('n_frames', n_frames)
            # print('n_frames//seqlen', n_frames//seqlen)
            
            n_heads = int(labels[:,1].max())

            boxlist = np.zeros((n_frames, n_heads, 4))
            scorelist = -1 * np.ones((n_frames, n_heads))
            vislist = np.zeros((n_frames, n_heads))

            # print('labels', labels.shape)

            for i in range(labels.shape[0]):
                frame_id, head_id, bb_left, bb_top, bb_width, bb_height, conf, cid, vis = labels[i]
                frame_id = int(frame_id) - 1 # convert 1 indexed to 0 indexed
                head_id = int(head_id) - 1 # convert 1 indexed to 0 indexed

                scorelist[frame_id, head_id] = 1
                vislist[frame_id, head_id] = vis
                box_cur = np.array([bb_left, bb_top, bb_left+bb_width, bb_top+bb_height]) # convert xywh to x1, y1, x2, y2
                boxlist[frame_id, head_id] = box_cur

            d = {
                'boxlist': np.copy(boxlist),
                'scorelist': np.copy(scorelist),
                'vislist': np.copy(vislist)
            }

            for start_frame in range(0, n_frames, self.S//2):
                end_frame = min(start_frame+self.S, n_frames)
                
                if end_frame-start_frame >= S_min: 
                    print('%d:%d' % (start_frame, end_frame))
                    self.all_folders.append(subfolder)
                    self.all_start_frames.append(start_frame)
                    self.all_end_frames.append(end_frame)
                    self.all_dicts.append(d)

        print('found %d samples' % len(self.all_folders))

    def __getitem__(self, index):
        print('index', index)

        folder = self.all_folders[index]
        d = self.all_dicts[index]
        start_frame = self.all_start_frames[index]
        end_frame = self.all_end_frames[index]

        S = self.S
        boxlist = d['boxlist'][start_frame:end_frame]
        valids = d['scorelist'][start_frame:end_frame]
        visibs = d['vislist'][start_frame:end_frame]
        # print('boxlist', boxlist.shape)
        
        # # get gt
        # boxlist = self.folder_to_gt[subfolder]['boxlist'][start_frame:start_frame+S] / 2 # S, n_head, 4
        # scorelist = self.folder_to_gt[subfolder]['scorelist'][start_frame:start_frame+S] # S, n_head
        # vislist = self.folder_to_gt[subfolder]['vislist'][start_frame:start_frame+S] # S, n_head
        # # print("boxlist", boxlist.shape, "scorelist", scorelist.shape, "vislist", vislist.shape)
        # self.rgb_paths = []

        # rgb_paths = []
        # for ii in range(start_frame, end_frame):
        #     rgb_path = os.path.join(self.dataset_dir, folder, 'img1', str(start_frame+1).zfill(6)+'.jpg')
        #     rgb_paths.append(rgb_path)

        # rgbs = []
        # for i in range(len(boxlist)):
        #     # read image
        #     image_name = os.path.join(self.dataset_dir, subfolder, 'img1', str(start_frame+i+1).zfill(6)+'.jpg')
        #     rgb = Image.open(image_name)
        #     # downsample
        #     rgb = rgb.resize((int(rgb.size[0]/2), int(rgb.size[1]/2)), Image.BILINEAR)
        #     rgb = np.array(rgb)
        #     rgbs.append(rgb)
        # rgbs = np.stack(rgbs, axis=0)

        print('N orig', boxlist.shape[1])
        
        trajs = np.stack([boxlist[:, :, [0,2]].mean(2), boxlist[:, :, [1,3]].mean(2)], axis=2) # S,N,2
        print('trajs', trajs.shape)
        print('visibs', visibs.shape)
        print('valids', valids.shape)

        # assert(trajs.shape[0]==self.S)

        S, N = trajs.shape[:2]
        H, W = 1080, 1920
        # update visibility annotations
        for si in range(S):
            # crohd annotations get noisy/wrong near edges
            oob_inds = np.logical_or(
                np.logical_or(trajs[si,:,0] < 32, trajs[si,:,0] > W-32),
                np.logical_or(trajs[si,:,1] < 32, trajs[si,:,1] > H-32))
            visibs[si,oob_inds] = 0
            # exclude oob from eval
            valids[si,oob_inds] = 0
        
        # _, _, _, H, W = rgbs.shape
        # print('H', H, 'W', W)
        # for s in range(S):
        #     for n in range(kp_xys.shape[2]):
        #         xy = kp_xys[0, s, n]
        #         if xy[0] <= 0 or xy[0] >= W or xy[1] <= 0 or xy[1] >= H:
        #             valids[0, s, n] = 0

        # N = trajs.shape[1]
        # H, W = 1080, 1920
        # for s in range(S):
        #     for n in range(N):
        #         x, y = trajs[s,n,0], trajs[s,n,1]
        #         if x <= 0 or x >= W-1 or y <= 0 or y >= H-1:
        #             valids[s,n] = 0

        vis_ok0 = visibs[0] > 0 # N
        vis_ok1 = visibs[1] > 0 # N
        score_ok0 = valids[0] > 0 # N
        score_ok1 = valids[1] > 0 # N
        mot_ok = np.sum(np.linalg.norm(trajs[1:] - trajs[:-1], axis=-1), axis=0) > 150 # N
        all_ok = vis_ok0 * vis_ok1 * score_ok0 * score_ok1 * mot_ok
        print('np.sum(all_ok)', np.sum(all_ok))

        trajs = trajs[:,all_ok]
        visibs = visibs[:,all_ok]
        valids = valids[:,all_ok]

        S, N = trajs.shape[:2]
        # to improve the visuals, let's avoid shooting the traj oob
        for ni in range(N):
            for si in range(1,S):
                if visibs[si,ni]==0:
                    trajs[si,ni] = trajs[si-1,ni]

        sample = {
            # 'rgb_paths': rgb_paths,
            'folder': os.path.join(self.dataset_dir, folder),
            'start_frame': start_frame,
            'end_frame': end_frame,
            'trajs': trajs,
            'visibs': visibs,
            'valids': valids,
        }
        return sample

    def __len__(self):
        # print("subfolder_lens", self.subfolder_lens)
        # return sum(self.subfolder_lens)
        return len(self.all_folders)

if __name__ == "__main__":
    B = 1
    S = 1200
    shuffle=False
    dataset = CrohdDataset('/Users/yangzheng/code/project/long-term-tracking/data/HT21', seqlen=S, dset='t')

    from torch.utils.data import Dataset, DataLoader
    train_dataloader = DataLoader(
        dataset,
        batch_size=B,
        shuffle=shuffle,
        num_workers=0)

    print(len(train_dataloader))
    train_iterloader = iter(train_dataloader)

    sample = next(train_iterloader)
    print(sample['rgbs'].shape)


datasets/crohddataset.py
aharley/pips2
from numpy import random
from numpy.core.numeric import full
import torch
import numpy as np
import os
import scipy.ndimage
import torchvision.transforms as transforms
import torch.nn.functional as F
import random
from torch._C import dtype, set_flush_denormal
import utils.geom
import glob
import cv2
from pathlib import Path
import sys
import albumentations as A
from functools import partial

def augment_video(augmenter, **kwargs):
    assert isinstance(augmenter, A.ReplayCompose)
    keys = kwargs.keys()
    for i in range(len(next(iter(kwargs.values())))):
        data = augmenter(**{
            key: kwargs[key][i] if key not in ['bboxes', 'keypoints'] else [kwargs[key][i]] for key in keys
        })
        if i == 0:
            augmenter = partial(A.ReplayCompose.replay, data['replay'])
        for key in keys:
            if key == 'bboxes':
                kwargs[key][i] = np.array(data[key]).reshape(4)
            elif key == 'keypoints':
                kwargs[key][i] = np.array(data[key]).reshape(2)
            else:
                kwargs[key][i] = data[key]

def read_mp4(fn):
    vidcap = cv2.VideoCapture(fn)
    frames = []
    while(vidcap.isOpened()):
        ret, frame = vidcap.read()
        if ret == False:
            break
        frames.append(frame)
    vidcap.release()
    return frames

class ExportDataset(torch.utils.data.Dataset):
    def __init__(self,
                 dataset_location='./pod_export',
                 dataset_version='aa',
                 S=36,
                 N=128,
                 crop_size=(384,512), 
                 use_augs=False,
    ):
        print('loading export...')

        self.dataset_location = dataset_location
        self.S = S
        self.N = N
        self.H, self.W = crop_size
        self.crop_size = crop_size
        self.use_augs = use_augs
        
        self.dataset_location = Path('%s/%s' % (self.dataset_location, dataset_version))

        folder_names = self.dataset_location.glob('*/')
        folder_names = [fn for fn in folder_names]
        folder_names = sorted(list(folder_names))
        print('found %d folders in %s' % (len(folder_names), self.dataset_location))

        self.all_folder_names = []
        rgbs = read_mp4(str(folder_names[0] / 'rgb.mp4'))
        rgbs = np.stack(rgbs, axis=0) # S,H,W,3
        S_local,H,W,C = rgbs.shape
        assert(H==self.H)
        assert(W==self.W)
        assert(S_local==self.S)
        
        self.all_folder_names = folder_names

        self.color_augmenter = A.ReplayCompose([
            A.GaussNoise(p=0.2),
            A.OneOf([
                A.MotionBlur(p=0.2),
                A.MedianBlur(blur_limit=3, p=0.1),
                A.Blur(blur_limit=3, p=0.1),
            ], p=0.2),
            A.OneOf([
                A.CLAHE(clip_limit=2),
                A.Sharpen(),
                A.Emboss(),
                A.RandomBrightnessContrast(),
            ], p=0.2),
            A.RGBShift(p=0.5),
            A.RandomBrightnessContrast(p=0.5),
            A.RandomGamma(p=0.5),
            A.HueSaturationValue(p=0.3),
            A.ImageCompression(quality_lower=50, quality_upper=100, p=0.3),
        ], p=0.8)
        
    
    def __getitem__(self, index):
        folder = self.all_folder_names[index]
        # print('folder', folder)
        
        rgbs = read_mp4(str(folder / 'rgb.mp4'))

        if len(rgbs)==0:
            print('corrupted mp4 in %s; returning fake' % folder)
            fake_sample = {
                'rgbs': np.zeros((self.S,3,self.H,self.W), dtype=np.uint8), 
                'track_g': np.zeros((self.S,self.N,4), dtype=np.float32)
            }
            return fake_sample
        rgbs = np.stack(rgbs, axis=0) # S,H,W,3

        d = dict(np.load(folder / 'track.npz', allow_pickle=True))
        track = d['track_g']

        H,W,C = rgbs[0].shape
        S,N,D = track.shape

        assert(N >= self.N)
        assert(H==self.H)
        assert(W==self.W)
        assert(S==self.S)
        assert(D==4) # xy, vis, valid
        assert(C==3)
        
        if N > self.N:
            inds = np.random.choice(N, self.N, replace=False)
            track = track[:,inds]
        
        if self.use_augs:
            augment_video(self.color_augmenter, image=rgbs)

        rgbs = rgbs.transpose(0,3,1,2)
        rgbs = rgbs[:,::-1].copy() # BGR->RGB

        return {
            'rgbs': rgbs,
            'track_g': track,
        }

    def __len__(self):
        return len(self.all_folder_names)

datasets/exportdataset.py
aharley/pips2
from numpy import random
import torch
import numpy as np
import os
import torchvision.transforms as transforms
import torch.nn.functional as F
from PIL import Image
import random
from torch._C import dtype, set_flush_denormal
import utils.basic
import utils.misc
import utils.improc
import glob
import cv2
import albumentations as A
from functools import partial
import sys

def augment_video(augmenter, **kwargs):
    assert isinstance(augmenter, A.ReplayCompose)
    keys = kwargs.keys()
    for i in range(len(next(iter(kwargs.values())))):
        data = augmenter(**{
            key: kwargs[key][i] if key not in ['bboxes', 'keypoints'] else [kwargs[key][i]] for key in keys
        })
        if i == 0:
            augmenter = partial(A.ReplayCompose.replay, data['replay'])
        for key in keys:
            if key == 'bboxes':
                kwargs[key][i] = np.array(data[key]).reshape(4)
            elif key == 'keypoints':
                kwargs[key][i] = np.array(data[key]).reshape(2)
            else:
                kwargs[key][i] = data[key]
                
class PointOdysseyDataset(torch.utils.data.Dataset):
    def __init__(self,
                 dataset_location='/orion/group/point_odyssey',
                 dset='train',
                 use_augs=False,
                 S=8,
                 N=32,
                 strides=[1,2,3,4],
                 crop_size=(368, 496),
                 req_full=False,
                 quick=False,
                 verbose=False,
    ):
        print('loading pointodyssey dataset...')

        self.S = S
        self.N = N
        self.req_full = False
        self.verbose = verbose

        self.use_augs = use_augs
        self.dset = dset

        self.rgb_paths = []
        self.depth_paths = []
        self.traj_paths = []
        self.annotation_paths = []
        self.full_idxs = []

        self.subdirs = []
        self.sequences = []
        
        assert(dset in ['train', 'val', 'test'])
        self.subdirs.append(os.path.join(dataset_location, dset))

        for subdir in self.subdirs:
            for seq in glob.glob(os.path.join(subdir, "*")):
                if os.path.isdir(seq):
                    seq_name = seq.split('/')[-1]
                    self.sequences.append(seq)

        self.sequences = sorted(self.sequences)
        if self.verbose:
            print(self.sequences)
        print('found %d unique videos in %s (dset=%s)' % (len(self.sequences), dataset_location, dset))
        
        print('loading trajectories...')
        if quick:
           self.sequences = self.sequences[:1] 
        
        for seq in self.sequences:
            
            rgb_path = os.path.join(seq, 'rgbs')

            annotations_path = os.path.join(seq, 'annotations.npz')
            if os.path.isfile(annotations_path):

                if self.verbose: 
                    print('seq', seq)
                    
                for stride in strides:
                    for ii in range(0,len(os.listdir(rgb_path))-self.S*stride+1, 4):
                        full_idx = ii + np.arange(self.S)*stride
                        self.rgb_paths.append([os.path.join(seq, 'rgbs', 'rgb_%05d.jpg' % idx) for idx in full_idx])
                        self.annotation_paths.append(os.path.join(seq, 'annotations.npz'))
                        self.full_idxs.append(full_idx)
                        if self.verbose:
                            sys.stdout.write('.')
                            sys.stdout.flush()
                        else:
                            if self.verbose:
                                sys.stdout.write('v')
                                sys.stdout.flush()

        print('collected %d clips of length %d in %s (dset=%s)' % (
            len(self.rgb_paths), self.S, dataset_location, dset))

        self.spatial_aug_prob = 0.7
        self.reverse_prob = 0.5

        # occlusion augmentation
        self.eraser_aug_prob = 0.2
        self.eraser_bounds = [20, 300]

        # spatial augmentations
        self.pad_bounds = [0, 64]
        self.crop_size = crop_size
        self.resize_lim = [0.25, 1.5] # sample resizes from here
        self.resize_delta = 0.1
        self.max_crop_offset = 20
        
        self.do_flip = True
        self.h_flip_prob = 0.5
        self.v_flip_prob = 0.5

        self.color_augmenter = A.ReplayCompose([
            A.GaussNoise(p=0.2),
            A.OneOf([
                A.MotionBlur(p=0.2),
                A.MedianBlur(blur_limit=3, p=0.1),
                A.Blur(blur_limit=3, p=0.1),
            ], p=0.2),
            A.OneOf([
                A.CLAHE(clip_limit=2),
                A.Sharpen(),
                A.Emboss(),
            ], p=0.2),
            A.RGBShift(p=0.5),
            A.RandomBrightnessContrast(p=0.5),
            A.RandomGamma(p=0.5),
            A.HueSaturationValue(p=0.3),
            A.ImageCompression(quality_lower=50, quality_upper=100, p=0.3),
        ], p=0.8)
        

    def getitem_helper(self, index):
        sample = None
        gotit = False

        rgb_paths = self.rgb_paths[index]
        full_idx = self.full_idxs[index]
        annotations_path = self.annotation_paths[index]
        annotations = np.load(annotations_path, allow_pickle=True)
        trajs = annotations['trajs_2d'][full_idx].astype(np.float32)
        visibs = annotations['visibilities'][full_idx].astype(np.float32)
        valids = (visibs<2).astype(np.float32) # S,N
        visibs = (visibs==1).astype(np.float32) # S,N

        # ensure that the point is good in frame0
        vis_and_val = valids * visibs
        vis0 = vis_and_val[0] > 0
        trajs = trajs[:,vis0]
        visibs = visibs[:,vis0]
        valids = valids[:,vis0]

        S,N,D = trajs.shape
        assert(D==2)
        assert(S==self.S)

        if self.req_full:
            min_N = self.N//2
        else:
            min_N = self.N

        if N < min_N:
            if self.verbose:
                print('returning before cropping: N=%d; need at least N=%d' % (N, min_N))
            return None, False

        # get rid of infs and nans
        valids_xy = np.ones_like(trajs)
        inf_idx = np.where(np.isinf(trajs))
        trajs[inf_idx] = 0
        valids_xy[inf_idx] = 0
        nan_idx = np.where(np.isnan(trajs))
        trajs[nan_idx] = 0
        valids_xy[nan_idx] = 0
        inv_idx = np.where(np.sum(valids_xy, axis=2)<2) # S,N
        visibs[inv_idx] = 0
        valids[inv_idx] = 0

        rgbs = []
        for rgb_path in rgb_paths:
            with Image.open(rgb_path) as im:
                rgbs.append(np.array(im)[:, :, :3])
                
        if self.use_augs:
            rgbs = np.stack(rgbs, 0)
            augment_video(self.color_augmenter, image=rgbs)
            rgbs = [rgb for rgb in rgbs]

            if np.random.rand() < self.reverse_prob:
                rgbs = np.stack(rgbs, 0)
                rgbs = np.flip(rgbs, axis=0)
                trajs = np.flip(trajs, axis=0)
                visibs = np.flip(visibs, axis=0)
                rgbs = [rgb for rgb in rgbs]

        if self.use_augs and (np.random.rand() < self.spatial_aug_prob):
            rgbs, trajs = self.add_spatial_augs(rgbs, trajs, visibs)
        else:
            # either crop or resize
            if np.random.rand() < 0.5:
                rgbs, trajs = self.just_crop(rgbs, trajs)
            else:
                rgbs, trajs = self.just_resize(rgbs, trajs)

        H,W,C = rgbs[0].shape
        assert(C==3)
        
        # update visibility annotations
        for si in range(S):
            # avoid 1px edge
            oob_inds = np.logical_or(
                np.logical_or(trajs[si,:,0] < 1, trajs[si,:,0] > W-2),
                np.logical_or(trajs[si,:,1] < 1, trajs[si,:,1] > H-2))
            visibs[si,oob_inds] = 0

            # when a point moves far oob, don't supervise with it
            very_oob_inds = np.logical_or(
                np.logical_or(trajs[si,:,0] < -64, trajs[si,:,0] > W+64),
                np.logical_or(trajs[si,:,1] < -64, trajs[si,:,1] > H+64))
            valids[si,very_oob_inds] = 0

        # ensure that the point is good in frame0
        vis_and_val = valids * visibs
        vis0 = vis_and_val[0] > 0
        trajs = trajs[:,vis0]
        visibs = visibs[:,vis0]
        valids = valids[:,vis0]

        # ensure that the point is good in frame1
        vis_and_val = valids * visibs
        vis1 = vis_and_val[1] > 0
        trajs = trajs[:,vis1]
        visibs = visibs[:,vis1]
        valids = valids[:,vis1]

        # ensure that the point is good in at least sqrt(S) frames
        val_ok = np.sum(valids, axis=0) >= max(np.sqrt(S),2)
        trajs = trajs[:,val_ok]
        visibs = visibs[:,val_ok]
        valids = valids[:,val_ok]

        # ensure that the per-frame motion isn't too crazy
        mot = np.max(np.linalg.norm(trajs[1:] - trajs[:-1], axis=-1), axis=0) # N
        mot_ok = mot < 128
        # if np.sum(~mot_ok):
        #     print('sum(~mot_ok)', np.sum(~mot_ok))
        trajs = trajs[:,mot_ok]
        visibs = visibs[:,mot_ok]
        valids = valids[:,mot_ok]
        
        N = trajs.shape[1]
        
        if N < min_N:
            if self.verbose:
                print('returning after cropping: N=%d; need at least N=%d' % (N, min_N))
            return None, False

        if N < self.N:
            if self.verbose:
                print('N=%d; ideally we want N=%d, but we will pad' % (N, self.N))

        if N > self.N*4:
            # fps based on position and motion
            xym = np.concatenate([np.mean(trajs, axis=0), np.mean(trajs[1:] - trajs[:-1], axis=0)], axis=-1)
            inds = utils.misc.farthest_point_sample_py(xym, self.N*4)
            trajs = trajs[:,inds]
            visibs = visibs[:,inds]
            valids = valids[:,inds]

        # we won't supervise with the extremes, but let's clamp anyway just to be safe
        trajs = np.minimum(np.maximum(trajs, np.array([-64,-64])), np.array([W+64, H+64])) # S,N,2
        
        N = trajs.shape[1]
        N_ = min(N, self.N)
        inds = np.random.choice(N, N_, replace=False)

        # prep for batching, by fixing N
        trajs_full = np.zeros((self.S, self.N, 2)).astype(np.float32)
        visibs_full = np.zeros((self.S, self.N)).astype(np.float32)
        valids_full = np.zeros((self.S, self.N)).astype(np.float32)
        trajs_full[:,:N_] = trajs[:,inds]
        visibs_full[:,:N_] = visibs[:,inds]
        valids_full[:,:N_] = valids[:,inds]

        rgbs = torch.from_numpy(np.stack(rgbs, 0)).permute(0,3,1,2).byte() # S,C,H,W
        trajs = torch.from_numpy(trajs_full).float() # S,N,2
        visibs = torch.from_numpy(visibs_full).float() # S,N
        valids = torch.from_numpy(valids_full).float() # S,N

        sample = {
            'rgbs': rgbs,
            'trajs': trajs,
            'visibs': visibs,
            'valids': valids,
        }
        return sample, True
    
    def __getitem__(self, index):
        gotit = False
        sample, gotit = self.getitem_helper(index)
        if not gotit:
            # return a fake sample, so we can still collate
            sample = {
                'rgbs': torch.zeros((self.S, 3, self.crop_size[0], self.crop_size[1]), dtype=torch.uint8),
                'trajs': torch.zeros((self.S, self.N, 2), dtype=torch.float32),
                'visibs': torch.zeros((self.S, self.N), dtype=torch.float32),
                'valids': torch.zeros((self.S, self.N), dtype=torch.float32),
            }
        return sample, gotit

    def add_spatial_augs(self, rgbs, trajs, visibs):
        T, N, _ = trajs.shape
        
        S = len(rgbs)
        H, W = rgbs[0].shape[:2]
        assert(S==T)

        rgbs = [rgb.astype(np.float32) for rgb in rgbs]
        
        # padding
        pad_x0 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
        pad_x1 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
        pad_y0 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
        pad_y1 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])

        rgbs = [np.pad(rgb, ((pad_y0, pad_y1), (pad_x0, pad_x1), (0, 0))) for rgb in rgbs]
        trajs[:,:,0] += pad_x0
        trajs[:,:,1] += pad_y0
        H, W = rgbs[0].shape[:2]

        # scaling + stretching
        scale = np.random.uniform(self.resize_lim[0], self.resize_lim[1])
        scale_x = scale
        scale_y = scale
        H_new = H
        W_new = W

        scale_delta_x = 0.0
        scale_delta_y = 0.0

        rgbs_scaled = []
        trajs_scaled = []
        
        scales_x = []
        scales_y = []
        for si in range(S):
            if si==1:
                scale_delta_x = np.random.uniform(-self.resize_delta, self.resize_delta)*0.1
                scale_delta_y = np.random.uniform(-self.resize_delta, self.resize_delta)*0.1
            elif si > 1:
                scale_delta_x = scale_delta_x*0.9 + np.random.uniform(-self.resize_delta, self.resize_delta)*0.1
                scale_delta_y = scale_delta_y*0.9 + np.random.uniform(-self.resize_delta, self.resize_delta)*0.1
            scale_x = scale_x + scale_delta_x
            scale_y = scale_y + scale_delta_y

            # bring h/w closer
            scale_xy = (scale_x + scale_y)*0.5
            scale_x = scale_x*0.5 + scale_xy*0.5
            scale_y = scale_y*0.5 + scale_xy*0.5
            
            # don't get too crazy
            scale_x = np.clip(scale_x, 0.2, 2.0)
            scale_y = np.clip(scale_y, 0.2, 2.0)
            
            H_new = int(H * scale_y)
            W_new = int(W * scale_x)

            # make it at least slightly bigger than the crop area,
            # so that the random cropping can add diversity
            H_new = np.clip(H_new, self.crop_size[0]+16, None)
            W_new = np.clip(W_new, self.crop_size[1]+16, None)
            # recompute scale in case we clipped
            scale_x = W_new/float(W)
            scale_y = H_new/float(H)

            rgbs_scaled.append(cv2.resize(rgbs[si], (W_new, H_new), interpolation=cv2.INTER_LINEAR))
            trajs[si,:,0] *= scale_x
            trajs[si,:,1] *= scale_y
        rgbs = rgbs_scaled
        
        ok_inds = visibs[0,:] > 0
        vis_trajs = trajs[:,ok_inds] # S,?,2
            
        if vis_trajs.shape[1] > 0:
            mid_x = np.mean(vis_trajs[0,:,0])
            mid_y = np.mean(vis_trajs[0,:,1])
        else:
            mid_y = self.crop_size[0]
            mid_x = self.crop_size[1]
            
        x0 = int(mid_x - self.crop_size[1]//2)
        y0 = int(mid_y - self.crop_size[0]//2)
        
        offset_x = 0
        offset_y = 0
        
        for si in range(S):
            # on each frame, shift a bit more 
            if si==1:
                offset_x = np.random.randint(-self.max_crop_offset, self.max_crop_offset)
                offset_y = np.random.randint(-self.max_crop_offset, self.max_crop_offset)
            elif si > 1:
                offset_x = int(offset_x*0.8 + np.random.randint(-self.max_crop_offset, self.max_crop_offset+1)*0.2)
                offset_y = int(offset_y*0.8 + np.random.randint(-self.max_crop_offset, self.max_crop_offset+1)*0.2)
            x0 = x0 + offset_x
            y0 = y0 + offset_y

            H_new, W_new = rgbs[si].shape[:2]
            if H_new==self.crop_size[0]:
                y0 = 0
            else:
                y0 = min(max(0, y0), H_new - self.crop_size[0] - 1)
                
            if W_new==self.crop_size[1]:
                x0 = 0
            else:
                x0 = min(max(0, x0), W_new - self.crop_size[1] - 1)
            rgbs[si] = rgbs[si][y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
            trajs[si,:,0] -= x0
            trajs[si,:,1] -= y0

            
        H_new = self.crop_size[0]
        W_new = self.crop_size[1]

        # flip
        h_flipped = False
        v_flipped = False
        if self.do_flip:
            if np.random.rand() < self.h_flip_prob:
                h_flipped = True
                rgbs = [rgb[:,::-1] for rgb in rgbs]
            if np.random.rand() < self.v_flip_prob:
                v_flipped = True
                rgbs = [rgb[::-1] for rgb in rgbs]
        if h_flipped:
            trajs[:,:,0] = W_new - trajs[:,:,0]
        if v_flipped:
            trajs[:,:,1] = H_new - trajs[:,:,1]
            
        return rgbs, trajs

    def just_crop(self, rgbs, trajs):
        T, N, _ = trajs.shape
        
        S = len(rgbs)
        H, W = rgbs[0].shape[:2]
        assert(S==T)

        H_new, W_new = self.crop_size[0], self.crop_size[1]

        y0 = np.random.randint(0, H-H_new)
        x0 = np.random.randint(0, W-W_new)
        rgbs = [rgb[y0:y0+H_new, x0:x0+W_new] for rgb in rgbs]
        trajs[:,:,0] -= x0
        trajs[:,:,1] -= y0

        return rgbs, trajs

    def just_resize(self, rgbs, trajs):
        T, N, _ = trajs.shape
        
        S = len(rgbs)
        H, W = rgbs[0].shape[:2]
        assert(S==T)

        H_new, W_new = self.crop_size[0], self.crop_size[1]

        sx_ = W_new / W
        sy_ = H_new / H
        rgbs = [cv2.resize(rgb, (W_new, H_new), interpolation=cv2.INTER_LINEAR) for rgb in rgbs]
        sc_py = np.array([sx_, sy_]).reshape([1,1,2])
        trajs = trajs * sc_py
        
        return rgbs, trajs

    def __len__(self):
        return len(self.rgb_paths)

datasets/pointodysseydataset.py
aharley/pips2
from numpy import random
import torch
import numpy as np
import os
import torchvision.transforms as transforms
import torch.nn.functional as F
from PIL import Image
import random
from torch._C import dtype, set_flush_denormal
import utils.basic
import utils.improc
import glob
import cv2
from torchvision.transforms import ColorJitter, GaussianBlur
import albumentations as A
from functools import partial
import sys

np.random.seed(125)
torch.multiprocessing.set_sharing_strategy('file_system')

class PointOdysseyDataset(torch.utils.data.Dataset):
    def __init__(self,
                 dataset_location='/orion/group/point_odyssey',
                 dset='train',
                 use_augs=False,
                 S=8,
                 N=32,
                 strides=[1],
                 quick=False,
                 verbose=False,
    ):
        print('loading pointodyssey dataset...')

        self.S = S
        self.N = N

        self.use_augs = use_augs
        self.dset = dset

        self.rgb_paths = []
        self.depth_paths = []
        self.normal_paths = []
        self.traj_paths = []
        self.annotation_paths = []
        self.full_idxs = []

        self.subdirs = []
        self.sequences = []
        self.subdirs.append(os.path.join(dataset_location, dset))

        for subdir in self.subdirs:
            for seq in glob.glob(os.path.join(subdir, "*/")):
                seq_name = seq.split('/')[-1]
                self.sequences.append(seq)

        self.sequences = sorted(self.sequences)
        if verbose:
            print(self.sequences)
        print('found %d unique videos in %s (dset=%s)' % (len(self.sequences), dataset_location, dset))
        
        ## load trajectories
        print('loading trajectories...')

        if quick:
           self.sequences = self.sequences[:1] 
        
        for seq in self.sequences:

            # if 'character1' not in seq:
            
            rgb_path = os.path.join(seq, 'rgbs')

            annotations_path = os.path.join(seq, 'annotations.npz')
            if os.path.isfile(annotations_path):

                if verbose: 
                    print('seq', seq)
                    
                for stride in strides:
                    for ii in range(0,len(os.listdir(rgb_path))-self.S*stride+1, 4):
                        full_idx = ii + np.arange(self.S)*stride
                        self.rgb_paths.append([os.path.join(seq, 'rgbs', 'rgb_%05d.jpg' % idx) for idx in full_idx])
                        self.depth_paths.append([os.path.join(seq, 'depths', 'depth_%05d.png' % idx) for idx in full_idx])
                        self.normal_paths.append([os.path.join(seq, 'normals', 'normal_%05d.jpg' % idx) for idx in full_idx])
                        self.annotation_paths.append(os.path.join(seq, 'annotations.npz'))
                        self.full_idxs.append(full_idx)
                        if verbose:
                            sys.stdout.write('.')
                            sys.stdout.flush()
                        else:
                            if verbose:
                                sys.stdout.write('v')
                                sys.stdout.flush()

        print('collected %d clips of length %d in %s (dset=%s)' % (
            len(self.rgb_paths), self.S, dataset_location, dset))

    def getitem_helper(self, index):
        sample = None
        gotit = False

        rgb_paths = self.rgb_paths[index]
        depth_paths = self.depth_paths[index]
        normal_paths = self.normal_paths[index]
        full_idx = self.full_idxs[index]
        annotations_path = self.annotation_paths[index]
        annotations = np.load(annotations_path, allow_pickle=True)
        trajs_2d = annotations['trajs_2d'][full_idx].astype(np.float32)
        visibs = annotations['visibilities'][full_idx].astype(np.float32)
        valids = (visibs<2).astype(np.float32)
        visibs = (visibs==1).astype(np.float32)
        trajs_world = annotations['trajs_3d'][full_idx].astype(np.float32)
        pix_T_cams = annotations['intrinsics'][full_idx].astype(np.float32)
        cams_T_world = annotations['extrinsics'][full_idx].astype(np.float32)

        # ensure that the point is good in frame0
        vis_and_val = valids * visibs
        vis0 = vis_and_val[0] > 0
        trajs_2d = trajs_2d[:,vis0]
        trajs_world = trajs_world[:,vis0]
        visibs = visibs[:,vis0]
        valids = valids[:,vis0]

        S,N,D = trajs_2d.shape
        assert(D==2)
        assert(S==self.S)
        
        if N < self.N//2:
            print('returning before cropping: N=%d; need at least N=%d' % (N, self.N//2))
            return None, False
        
        trajs_cam = utils.geom.apply_4x4_py(cams_T_world, trajs_world)
        trajs_pix = utils.geom.apply_pix_T_cam_py(pix_T_cams, trajs_cam)

        # get rid of infs and nans
        valids_xy = np.ones_like(trajs_2d)
        inf_idx = np.where(np.isinf(trajs_2d))
        trajs_2d[inf_idx] = 0
        valids_xy[inf_idx] = 0
        nan_idx = np.where(np.isnan(trajs_2d))
        trajs_2d[nan_idx] = 0
        valids_xy[nan_idx] = 0
        inv_idx = np.where(np.sum(valids_xy, axis=2)<2) # S,N
        visibs[inv_idx] = 0
        valids[inv_idx] = 0

        rgbs = []
        for rgb_path in rgb_paths:
            with Image.open(rgb_path) as im:
                rgbs.append(np.array(im)[:, :, :3])

        depths = []
        for depth_path in depth_paths:
            depth16 = cv2.imread(depth_path, cv2.IMREAD_ANYDEPTH)
            depth = depth16.astype(np.float32) / 65535.0 * 1000.0
            depths.append(depth)

        normals = []
        for normal_path in normal_paths:
            with Image.open(normal_path) as im:
                normals.append(np.array(im)[:, :, :3])
            

        H,W,C = rgbs[0].shape
        assert(C==3)
        
        # update visibility annotations
        for si in range(S):
            # avoid 1px edge
            oob_inds = np.logical_or(
                np.logical_or(trajs_2d[si,:,0] < 1, trajs_2d[si,:,0] > W-2),
                np.logical_or(trajs_2d[si,:,1] < 1, trajs_2d[si,:,1] > H-2))
            visibs[si,oob_inds] = 0

            # when a point moves far oob, don't supervise with it
            very_oob_inds = np.logical_or(
                np.logical_or(trajs_2d[si,:,0] < -64, trajs_2d[si,:,0] > W+64),
                np.logical_or(trajs_2d[si,:,1] < -64, trajs_2d[si,:,1] > H+64))
            valids[si,very_oob_inds] = 0

        # ensure that the point is good in frame0
        vis_and_val = valids * visibs
        vis0 = vis_and_val[0] > 0
        trajs_2d = trajs_2d[:,vis0]
        trajs_cam = trajs_cam[:,vis0]
        trajs_pix = trajs_pix[:,vis0]
        visibs = visibs[:,vis0]
        valids = valids[:,vis0]

        # ensure that the point is good in frame1
        vis_and_val = valids * visibs
        vis1 = vis_and_val[1] > 0
        trajs_2d = trajs_2d[:,vis1]
        trajs_cam = trajs_cam[:,vis1]
        trajs_pix = trajs_pix[:,vis1]
        visibs = visibs[:,vis1]
        valids = valids[:,vis1]

        # ensure that the point is visible at frame0
        vis0 = visibs[0] > 0
        trajs_2d = trajs_2d[:,vis0]
        trajs_cam = trajs_cam[:,vis0]
        trajs_pix = trajs_pix[:,vis0]
        visibs = visibs[:,vis0]
        valids = valids[:,vis0]

        # ensure that the point is good in at least sqrt(S) frames
        val_ok = np.sum(valids, axis=0) >= max(np.sqrt(S),2)
        trajs_2d = trajs_2d[:,val_ok]
        trajs_cam = trajs_cam[:,val_ok]
        trajs_pix = trajs_pix[:,val_ok]
        visibs = visibs[:,val_ok]
        valids = valids[:,val_ok]
        
        N = trajs_2d.shape[1]
        
        if N < self.N//2:
            # print('N=%d' % (N))
            return None, False
        
        if N < self.N:
            print('N=%d; ideally we want N=%d, but we will pad' % (N, self.N))

        if N > self.N*4:
            # fps based on position and motion
            xym = np.concatenate([np.mean(trajs_2d, axis=0), np.mean(trajs_2d[1:] - trajs_2d[:-1], axis=0)], axis=-1)
            inds = utils.misc.farthest_point_sample_py(xym, self.N*4)
            trajs_2d = trajs_2d[:,inds]
            trajs_cam = trajs_cam[:,inds]
            trajs_pix = trajs_pix[:,inds]
            visibs = visibs[:,inds]
            valids = valids[:,inds]

        # clamp so that the trajectories don't get too crazy
        trajs_2d = np.minimum(np.maximum(trajs_2d, np.array([-64,-64])), np.array([W+64, H+64])) # S,N,2
        trajs_pix = np.minimum(np.maximum(trajs_pix, np.array([-64,-64])), np.array([W+64, H+64])) # S,N,2
            
        N = trajs_2d.shape[1]
        N_ = min(N, self.N)
        inds = np.random.choice(N, N_, replace=False)

        # prep for batching, by fixing N
        trajs_2d_full = np.zeros((self.S, self.N, 2)).astype(np.float32)
        trajs_cam_full = np.zeros((self.S, self.N, 3)).astype(np.float32)
        trajs_pix_full = np.zeros((self.S, self.N, 2)).astype(np.float32)
        visibs_full = np.zeros((self.S, self.N)).astype(np.float32)
        valids_full = np.zeros((self.S, self.N)).astype(np.float32)
        trajs_2d_full[:,:N_] = trajs_2d[:,inds]
        trajs_cam_full[:,:N_] = trajs_cam[:,inds]
        trajs_pix_full[:,:N_] = trajs_pix[:,inds]
        visibs_full[:,:N_] = visibs[:,inds]
        valids_full[:,:N_] = valids[:,inds]

        rgbs = torch.from_numpy(np.stack(rgbs, 0)).permute(0,3,1,2) # S,3,H,W
        depths = torch.from_numpy(np.stack(depths, 0)).unsqueeze(1) # S,1,H,W
        normals = torch.from_numpy(np.stack(normals, 0)).permute(0,3,1,2) # S,3,H,W
        trajs_2d = torch.from_numpy(trajs_2d_full) # S,N,2
        trajs_cam = torch.from_numpy(trajs_cam_full) # S,N,3
        trajs_pix = torch.from_numpy(trajs_pix_full) # S,N,2
        visibs = torch.from_numpy(visibs_full) # S,N
        valids = torch.from_numpy(valids_full) # S,N

        sample = {
            'rgbs': rgbs,
            'depths': depths,
            'normals': normals,
            'trajs_2d': trajs_2d,
            'trajs_cam': trajs_cam,
            'trajs_pix': trajs_pix,
            'pix_T_cams': pix_T_cams,
            'cams_T_world': cams_T_world,
            'visibs': visibs,
            'valids': valids,
        }
        
        return sample, True

    
    def __getitem__(self, index):
        gotit = False
        sample, gotit = self.getitem_helper(index)
        if not gotit:
            print('warning: sampling failed')
            # return a fake sample, so we can still collate
            sample = {
                'rgbs': torch.zeros((self.S, 3, 540, 960)),
                'depths': torch.zeros((self.S, 1, 540, 960)),
                'trajs_2d': torch.zeros((self.S, self.N, 2)),
                'trajs_cam': torch.zeros((self.S, self.N, 3)),
                'trajs_pix': torch.zeros((self.S, self.N, 2)),
                'visibs': torch.zeros((self.S, self.N)),
                'valids': torch.zeros((self.S, self.N)),
            }
        return sample, gotit

    def __len__(self):
        return len(self.rgb_paths)

datasets/pointodysseydataset_3d.py
aharley/pips2
from numpy import random
import torch
import numpy as np
import os
import torchvision.transforms as transforms
import torch.nn.functional as F
from PIL import Image
import random
from torch._C import dtype, set_flush_denormal
import utils.basic
import utils.misc
import utils.improc
import glob
import cv2
from torchvision.transforms import ColorJitter, GaussianBlur
import albumentations as A
from functools import partial
import sys

class PointOdysseyDataset(torch.utils.data.Dataset):
    def __init__(self,
                 dataset_location='/orion/group/point_odyssey',
                 dset='test',
                 N=32,
                 verbose=False,
    ):
        print('loading pointodyssey fullseq dataset...')

        # in this version we load the full video,
        # rather than chopping it into subseqs of length S
        # and since this does not fit in memory,
        # we return paths instead of rgbs

        self.dset = dset
        self.N = N

        self.seq_paths = []
        self.rgb_paths = []
        self.traj_paths = []
        self.annotation_paths = []
        self.full_idxs = []

        self.subdirs = []
        self.sequences = []

        assert(dset in ['train', 'val', 'test'])
        self.subdirs.append(os.path.join(dataset_location, dset))

        for subdir in self.subdirs:
            for seq in glob.glob(os.path.join(subdir, "*/")):
                seq_name = seq.split('/')[-1]
                self.sequences.append(seq)

        self.sequences = sorted(self.sequences)
        if verbose:
            print(self.sequences)
        print('found %d unique videos in %s (dset=%s)' % (len(self.sequences), dataset_location, dset))
        
        print('loading trajectories...')
        for seq in self.sequences:
            rgb_path = os.path.join(seq, 'rgbs')
            annotations_path = os.path.join(seq, 'annotations.npz')
            assert(os.path.isfile(annotations_path))

            full_idx = np.arange(len(os.listdir(rgb_path)))
            self.rgb_paths.append([os.path.join(seq, 'rgbs', 'rgb_%05d.jpg' % (idx)) for idx in full_idx])
            self.annotation_paths.append(os.path.join(seq, 'annotations.npz'))
            self.full_idxs.append(full_idx)

            if verbose:
                sys.stdout.write('.')
                sys.stdout.flush()
            else:
                if verbose:
                    sys.stdout.write('v')
                    sys.stdout.flush()

        print('collected %d videos in %s (dset=%s)' % (
            len(self.rgb_paths), dataset_location, dset))

    def __getitem__(self, index):
        seq = self.sequences[index]
        rgb_paths = self.rgb_paths[index]
        full_idx = self.full_idxs[index]
        
        annotations_path = self.annotation_paths[index]
        annotations = np.load(annotations_path, allow_pickle=True)
        trajs = annotations['trajs_2d'][full_idx].astype(np.float32)
        visibs = annotations['visibilities'][full_idx].astype(np.float32)
        valids = (visibs<2).astype(np.float32)
        visibs = (visibs==1).astype(np.float32)

        S,N,D = trajs.shape
        assert(D==2)

        valids_xy = np.ones_like(trajs)

        # get rid of infs and nans
        inf_idx = np.where(np.isinf(trajs))
        trajs[inf_idx] = 0
        valids_xy[inf_idx] = 0
        nan_idx = np.where(np.isnan(trajs))
        trajs[nan_idx] = 0
        valids_xy[nan_idx] = 0

        inv_idx = np.where(np.sum(valids_xy, axis=2)<2) # S,N
        visibs[inv_idx] = 0
        valids[inv_idx] = 0
        
        H, W = 540, 960
        
        # update visibility annotations
        for si in range(S):
            # avoid 1px edge
            oob_inds = np.logical_or(
                np.logical_or(trajs[si,:,0] < 1, trajs[si,:,0] > W-2),
                np.logical_or(trajs[si,:,1] < 1, trajs[si,:,1] > H-2))
            visibs[si,oob_inds] = 0
            # exclude oob from eval
            valids[si,oob_inds] = 0

        # ensure that the point is good at frame0
        vis_and_val = valids * visibs
        vis0 = vis_and_val[0] > 0
        trajs = trajs[:,vis0]
        visibs = visibs[:,vis0]
        valids = valids[:,vis0]

        # ensure that the point is good in at least K frames total
        vis_and_val = valids * visibs
        val_ok = np.sum(vis_and_val, axis=0) >= 8
        trajs = trajs[:,val_ok]
        visibs = visibs[:,val_ok]
        valids = valids[:,val_ok]
        
        N = trajs.shape[1]
        
        if N < self.N:
            print('N=%d; ideally we want N=%d, but we will pad' % (N, self.N))

        # we won't supervise with the extremes, but let's clamp anyway just to be safe
        trajs = np.minimum(np.maximum(trajs, np.array([-32,-32])), np.array([W+32, H+32])) # S,N,2
        
        N = trajs.shape[1]
        N_ = min(N, self.N)
        # inds = np.random.choice(N, N_, replace=False)
        inds = np.linspace(0, N-1, N_).astype(np.int32)

        # prep for batching, by fixing N
        trajs_full = np.zeros((S, self.N, 2)).astype(np.float32)
        visibs_full = np.zeros((S, self.N)).astype(np.float32)
        valids_full = np.zeros((S, self.N)).astype(np.float32)

        trajs_full[:,:N_] = trajs[:,inds]
        visibs_full[:,:N_] = visibs[:,inds]
        valids_full[:,:N_] = valids[:,inds]

        # rgbs = torch.from_numpy(np.stack(rgbs, 0)).permute(0,3,1,2)  # S, C, H, W
        trajs = torch.from_numpy(trajs_full)  # S, N, 2
        visibs = torch.from_numpy(visibs_full)  # S, N
        valids = torch.from_numpy(valids_full)  # S, N

        sample = {
            'seq': seq,
            'rgb_paths': rgb_paths,
            'trajs': trajs,
            'visibs': visibs,
            'valids': valids,
        }
        return sample

    def __len__(self):
        return len(self.rgb_paths)

datasets/pointodysseydataset_fullseq.py
aharley/pips2
from numpy import random
import torch
import numpy as np
import os
import torchvision.transforms as transforms
import torch.nn.functional as F
from PIL import Image
import random
from torch._C import dtype, set_flush_denormal
import utils.basic
import utils.misc
import utils.improc
import glob
import cv2
import albumentations as A
from functools import partial
import sys

def augment_video(augmenter, **kwargs):
    assert isinstance(augmenter, A.ReplayCompose)
    keys = kwargs.keys()
    for i in range(len(next(iter(kwargs.values())))):
        data = augmenter(**{
            key: kwargs[key][i] if key not in ['bboxes', 'keypoints'] else [kwargs[key][i]] for key in keys
        })
        if i == 0:
            augmenter = partial(A.ReplayCompose.replay, data['replay'])
        for key in keys:
            if key == 'bboxes':
                kwargs[key][i] = np.array(data[key]).reshape(4)
            elif key == 'keypoints':
                kwargs[key][i] = np.array(data[key]).reshape(2)
            else:
                kwargs[key][i] = data[key]
                
class PointOdysseyDataset(torch.utils.data.Dataset):
    def __init__(self,
                 dataset_location='/orion/group/point_odyssey',
                 dset='train',
                 use_augs=False,
                 S=8,
                 N=32,
                 strides=[1,2,3,4],
                 crop_size=(368, 496),
                 req_full=False,
                 quick=False,
                 verbose=False,
    ):
        print('loading pointodyssey dataset...')

        self.S = S
        self.N = N
        self.req_full = False
        self.verbose = verbose

        self.use_augs = use_augs
        self.dset = dset

        self.rgb_paths = []
        self.mask_paths = []
        self.traj_paths = []
        self.annotation_paths = []
        self.full_idxs = []

        self.subdirs = []
        self.sequences = []
        
        assert(dset in ['train', 'val', 'test'])
        self.subdirs.append(os.path.join(dataset_location, dset))

        for subdir in self.subdirs:
            for seq in glob.glob(os.path.join(subdir, "*")):
                if os.path.isdir(seq):
                    seq_name = seq.split('/')[-1]
                    self.sequences.append(seq)

        self.sequences = sorted(self.sequences)
        if self.verbose:
            print(self.sequences)
        print('found %d unique videos in %s (dset=%s)' % (len(self.sequences), dataset_location, dset))
        
        print('loading trajectories...')
        if quick:
           self.sequences = self.sequences[:1] 
        
        for seq in self.sequences:
            
            rgb_path = os.path.join(seq, 'rgbs')

            annotations_path = os.path.join(seq, 'annotations.npz')
            if os.path.isfile(annotations_path):

                if self.verbose: 
                    print('seq', seq)
                    
                for stride in strides:
                    for ii in range(0,len(os.listdir(rgb_path))-self.S*stride+1, 4):
                        full_idx = ii + np.arange(self.S)*stride
                        self.rgb_paths.append([os.path.join(seq, 'rgbs', 'rgb_%05d.jpg' % idx) for idx in full_idx])
                        self.mask_paths.append([os.path.join(seq, 'masks', 'mask_%05d.png' % idx) for idx in full_idx])
                        self.annotation_paths.append(os.path.join(seq, 'annotations.npz'))
                        self.full_idxs.append(full_idx)
                        if self.verbose:
                            sys.stdout.write('.')
                            sys.stdout.flush()
                        else:
                            if self.verbose:
                                sys.stdout.write('v')
                                sys.stdout.flush()

        print('collected %d clips of length %d in %s (dset=%s)' % (
            len(self.rgb_paths), self.S, dataset_location, dset))

        self.spatial_aug_prob = 0.7
        self.reverse_prob = 0.5

        # occlusion augmentation
        self.eraser_aug_prob = 0.2
        self.eraser_bounds = [20, 300]

        # spatial augmentations
        self.pad_bounds = [0, 64]
        self.crop_size = crop_size
        self.resize_lim = [0.25, 1.5] # sample resizes from here
        self.resize_delta = 0.1
        self.max_crop_offset = 20
        
        self.do_flip = True
        self.h_flip_prob = 0.5
        self.v_flip_prob = 0.5

        self.color_augmenter = A.ReplayCompose([
            A.GaussNoise(p=0.2),
            A.OneOf([
                A.MotionBlur(p=0.2),
                A.MedianBlur(blur_limit=3, p=0.1),
                A.Blur(blur_limit=3, p=0.1),
            ], p=0.2),
            A.OneOf([
                A.CLAHE(clip_limit=2),
                A.Sharpen(),
                A.Emboss(),
            ], p=0.2),
            A.RGBShift(p=0.5),
            A.RandomBrightnessContrast(p=0.5),
            A.RandomGamma(p=0.5),
            A.HueSaturationValue(p=0.3),
            A.ImageCompression(quality_lower=50, quality_upper=100, p=0.3),
        ], p=0.8)
        

    def getitem_helper(self, index):
        sample = None
        gotit = False

        rgb_paths = self.rgb_paths[index]
        mask_paths = self.mask_paths[index]
        full_idx = self.full_idxs[index]
        annotations_path = self.annotation_paths[index]
        annotations = np.load(annotations_path, allow_pickle=True)
        trajs = annotations['trajs_2d'][full_idx].astype(np.float32)
        visibs = annotations['visibilities'][full_idx].astype(np.float32)
        valids = (visibs<2).astype(np.float32) # S,N
        visibs = (visibs==1).astype(np.float32) # S,N

        # ensure that the point is good in frame0
        vis_and_val = valids * visibs
        vis0 = vis_and_val[0] > 0
        trajs = trajs[:,vis0]
        visibs = visibs[:,vis0]
        valids = valids[:,vis0]

        S,N,D = trajs.shape
        assert(D==2)
        assert(S==self.S)

        if self.req_full:
            min_N = self.N//2
        else:
            min_N = self.N

        if N < min_N:
            if self.verbose:
                print('returning before cropping: N=%d; need at least N=%d' % (N, min_N))
            return None, False

        # get rid of infs and nans
        valids_xy = np.ones_like(trajs)
        inf_idx = np.where(np.isinf(trajs))
        trajs[inf_idx] = 0
        valids_xy[inf_idx] = 0
        nan_idx = np.where(np.isnan(trajs))
        trajs[nan_idx] = 0
        valids_xy[nan_idx] = 0
        inv_idx = np.where(np.sum(valids_xy, axis=2)<2) # S,N
        visibs[inv_idx] = 0
        valids[inv_idx] = 0

        rgbs = []
        for rgb_path in rgb_paths:
            with Image.open(rgb_path) as im:
                rgbs.append(np.array(im)[:, :, :3]) # H,W,3

        masks = []
        for mask_path in mask_paths:
            with Image.open(mask_path) as im:
                masks.append(np.array(im)) # H,W,1
        
        if self.use_augs:
            rgbs = np.stack(rgbs, 0)
            augment_video(self.color_augmenter, image=rgbs)
            rgbs = [rgb for rgb in rgbs]

            if np.random.rand() < self.reverse_prob:
                rgbs = np.stack(rgbs, 0)
                masks = np.stack(masks, 0)
                rgbs = np.flip(rgbs, axis=0)
                masks = np.flip(masks, axis=0)
                trajs = np.flip(trajs, axis=0)
                visibs = np.flip(visibs, axis=0)
                rgbs = [rgb for rgb in rgbs]
                masks = [mask for mask in masks]

        if self.use_augs and (np.random.rand() < self.spatial_aug_prob):
            rgbs, masks, trajs = self.add_spatial_augs(rgbs, masks, trajs, visibs)
        else:
            # either crop or resize
            if np.random.rand() < 0.5:
                rgbs, masks, trajs = self.just_crop(rgbs, masks, trajs)
            else:
                rgbs, masks, trajs = self.just_resize(rgbs, masks, trajs)

        H,W,C = rgbs[0].shape
        assert(C==3)
        
        # update visibility annotations
        for si in range(S):
            # avoid 1px edge
            oob_inds = np.logical_or(
                np.logical_or(trajs[si,:,0] < 1, trajs[si,:,0] > W-2),
                np.logical_or(trajs[si,:,1] < 1, trajs[si,:,1] > H-2))
            visibs[si,oob_inds] = 0

            # when a point moves far oob, don't supervise with it
            very_oob_inds = np.logical_or(
                np.logical_or(trajs[si,:,0] < -64, trajs[si,:,0] > W+64),
                np.logical_or(trajs[si,:,1] < -64, trajs[si,:,1] > H+64))
            valids[si,very_oob_inds] = 0

        # ensure that the point is good in frame0
        vis_and_val = valids * visibs
        vis0 = vis_and_val[0] > 0
        trajs = trajs[:,vis0]
        visibs = visibs[:,vis0]
        valids = valids[:,vis0]

        # ensure that the point is good in frame1
        vis_and_val = valids * visibs
        vis1 = vis_and_val[1] > 0
        trajs = trajs[:,vis1]
        visibs = visibs[:,vis1]
        valids = valids[:,vis1]

        # ensure that the point is good in at least sqrt(S) frames
        val_ok = np.sum(valids, axis=0) >= max(np.sqrt(S),2)
        trajs = trajs[:,val_ok]
        visibs = visibs[:,val_ok]
        valids = valids[:,val_ok]

        # ensure that the per-frame motion isn't too crazy
        mot = np.max(np.linalg.norm(trajs[1:] - trajs[:-1], axis=-1), axis=0) # N
        mot_ok = mot < 128
        # if np.sum(~mot_ok):
        #     print('sum(~mot_ok)', np.sum(~mot_ok))
        trajs = trajs[:,mot_ok]
        visibs = visibs[:,mot_ok]
        valids = valids[:,mot_ok]
        
        N = trajs.shape[1]
        
        if N < min_N:
            if self.verbose:
                print('returning after cropping: N=%d; need at least N=%d' % (N, min_N))
            return None, False

        if N < self.N:
            if self.verbose:
                print('N=%d; ideally we want N=%d, but we will pad' % (N, self.N))

        if N > self.N*4:
            # fps based on position and motion
            xym = np.concatenate([np.mean(trajs, axis=0), np.mean(trajs[1:] - trajs[:-1], axis=0)], axis=-1)
            inds = utils.misc.farthest_point_sample_py(xym, self.N*4)
            trajs = trajs[:,inds]
            visibs = visibs[:,inds]
            valids = valids[:,inds]

        # we won't supervise with the extremes, but let's clamp anyway just to be safe
        trajs = np.minimum(np.maximum(trajs, np.array([-64,-64])), np.array([W+64, H+64])) # S,N,2
        
        N = trajs.shape[1]
        N_ = min(N, self.N)
        inds = np.random.choice(N, N_, replace=False)

        # prep for batching, by fixing N
        trajs_full = np.zeros((self.S, self.N, 2)).astype(np.float32)
        visibs_full = np.zeros((self.S, self.N)).astype(np.float32)
        valids_full = np.zeros((self.S, self.N)).astype(np.float32)
        trajs_full[:,:N_] = trajs[:,inds]
        visibs_full[:,:N_] = visibs[:,inds]
        valids_full[:,:N_] = valids[:,inds]

        rgbs = torch.from_numpy(np.stack(rgbs, 0)).permute(0,3,1,2).byte() # S,C,H,W
        masks = torch.from_numpy(np.stack(masks, 0)).unsqueeze(1).byte() # S,C,H,W
        trajs = torch.from_numpy(trajs_full).float() # S,N,2
        visibs = torch.from_numpy(visibs_full).float() # S,N
        valids = torch.from_numpy(valids_full).float() # S,N

        sample = {
            'rgbs': rgbs,
            'masks': masks,
            'trajs': trajs,
            'visibs': visibs,
            'valids': valids,
        }
        return sample, True
    
    def __getitem__(self, index):
        gotit = False
        sample, gotit = self.getitem_helper(index)
        if not gotit:
            # return a fake sample, so we can still collate
            sample = {
                'rgbs': torch.zeros((self.S, 3, self.crop_size[0], self.crop_size[1]), dtype=torch.uint8),
                'trajs': torch.zeros((self.S, self.N, 2), dtype=torch.float32),
                'visibs': torch.zeros((self.S, self.N), dtype=torch.float32),
                'valids': torch.zeros((self.S, self.N), dtype=torch.float32),
            }
        return sample, gotit

    def add_spatial_augs(self, rgbs, masks, trajs, visibs):
        T, N, _ = trajs.shape
        
        S = len(rgbs)
        H, W = rgbs[0].shape[:2]
        assert(S==T)

        rgbs = [rgb.astype(np.float32) for rgb in rgbs]
        masks = [mask.astype(np.float32) for mask in masks]
        
        # padding
        pad_x0 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
        pad_x1 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
        pad_y0 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])
        pad_y1 = np.random.randint(self.pad_bounds[0], self.pad_bounds[1])

        rgbs = [np.pad(rgb, ((pad_y0, pad_y1), (pad_x0, pad_x1), (0, 0))) for rgb in rgbs]
        masks = [np.pad(mask, ((pad_y0, pad_y1), (pad_x0, pad_x1), (0, 0))) for mask in masks]
        trajs[:,:,0] += pad_x0
        trajs[:,:,1] += pad_y0
        H, W = rgbs[0].shape[:2]

        # scaling + stretching
        scale = np.random.uniform(self.resize_lim[0], self.resize_lim[1])
        scale_x = scale
        scale_y = scale
        H_new = H
        W_new = W

        scale_delta_x = 0.0
        scale_delta_y = 0.0

        rgbs_scaled = []
        trajs_scaled = []
        
        scales_x = []
        scales_y = []
        for si in range(S):
            if si==1:
                scale_delta_x = np.random.uniform(-self.resize_delta, self.resize_delta)*0.1
                scale_delta_y = np.random.uniform(-self.resize_delta, self.resize_delta)*0.1
            elif si > 1:
                scale_delta_x = scale_delta_x*0.9 + np.random.uniform(-self.resize_delta, self.resize_delta)*0.1
                scale_delta_y = scale_delta_y*0.9 + np.random.uniform(-self.resize_delta, self.resize_delta)*0.1
            scale_x = scale_x + scale_delta_x
            scale_y = scale_y + scale_delta_y

            # bring h/w closer
            scale_xy = (scale_x + scale_y)*0.5
            scale_x = scale_x*0.5 + scale_xy*0.5
            scale_y = scale_y*0.5 + scale_xy*0.5
            
            # don't get too crazy
            scale_x = np.clip(scale_x, 0.2, 2.0)
            scale_y = np.clip(scale_y, 0.2, 2.0)
            
            H_new = int(H * scale_y)
            W_new = int(W * scale_x)

            # make it at least slightly bigger than the crop area,
            # so that the random cropping can add diversity
            H_new = np.clip(H_new, self.crop_size[0]+16, None)
            W_new = np.clip(W_new, self.crop_size[1]+16, None)
            # recompute scale in case we clipped
            scale_x = W_new/float(W)
            scale_y = H_new/float(H)

            rgbs_scaled.append(cv2.resize(rgbs[si], (W_new, H_new), interpolation=cv2.INTER_LINEAR))
            masks_scaled.append(cv2.resize(masks[si], (W_new, H_new), interpolation=cv2.INTER_NEAREST))
            trajs[si,:,0] *= scale_x
            trajs[si,:,1] *= scale_y
        rgbs = rgbs_scaled
        masks = masks_scaled
        
        ok_inds = visibs[0,:] > 0
        vis_trajs = trajs[:,ok_inds] # S,?,2
            
        if vis_trajs.shape[1] > 0:
            mid_x = np.mean(vis_trajs[0,:,0])
            mid_y = np.mean(vis_trajs[0,:,1])
        else:
            mid_y = self.crop_size[0]
            mid_x = self.crop_size[1]
            
        x0 = int(mid_x - self.crop_size[1]//2)
        y0 = int(mid_y - self.crop_size[0]//2)
        
        offset_x = 0
        offset_y = 0
        
        for si in range(S):
            # on each frame, shift a bit more 
            if si==1:
                offset_x = np.random.randint(-self.max_crop_offset, self.max_crop_offset)
                offset_y = np.random.randint(-self.max_crop_offset, self.max_crop_offset)
            elif si > 1:
                offset_x = int(offset_x*0.8 + np.random.randint(-self.max_crop_offset, self.max_crop_offset+1)*0.2)
                offset_y = int(offset_y*0.8 + np.random.randint(-self.max_crop_offset, self.max_crop_offset+1)*0.2)
            x0 = x0 + offset_x
            y0 = y0 + offset_y

            H_new, W_new = rgbs[si].shape[:2]
            if H_new==self.crop_size[0]:
                y0 = 0
            else:
                y0 = min(max(0, y0), H_new - self.crop_size[0] - 1)
                
            if W_new==self.crop_size[1]:
                x0 = 0
            else:
                x0 = min(max(0, x0), W_new - self.crop_size[1] - 1)
            rgbs[si] = rgbs[si][y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
            masks[si] = masks[si][y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
            trajs[si,:,0] -= x0
            trajs[si,:,1] -= y0

            
        H_new = self.crop_size[0]
        W_new = self.crop_size[1]

        # flip
        h_flipped = False
        v_flipped = False
        if self.do_flip:
            if np.random.rand() < self.h_flip_prob:
                h_flipped = True
                rgbs = [rgb[:,::-1] for rgb in rgbs]
                masks = [mask[:,::-1] for mask in masks]
            if np.random.rand() < self.v_flip_prob:
                v_flipped = True
                rgbs = [rgb[::-1] for rgb in rgbs]
                masks = [mask[::-1] for mask in masks]
        if h_flipped:
            trajs[:,:,0] = W_new - trajs[:,:,0]
        if v_flipped:
            trajs[:,:,1] = H_new - trajs[:,:,1]
            
        return rgbs, masks, trajs

    def just_crop(self, rgbs, masks, trajs):
        T, N, _ = trajs.shape
        
        S = len(rgbs)
        H, W = rgbs[0].shape[:2]
        assert(S==T)

        H_new, W_new = self.crop_size[0], self.crop_size[1]

        y0 = np.random.randint(0, H-H_new)
        x0 = np.random.randint(0, W-W_new)
        rgbs = [rgb[y0:y0+H_new, x0:x0+W_new] for rgb in rgbs]
        masks = [mask[y0:y0+H_new, x0:x0+W_new] for mask in masks]
        trajs[:,:,0] -= x0
        trajs[:,:,1] -= y0

        return rgbs, masks, trajs

    def just_resize(self, rgbs, masks, trajs):
        T, N, _ = trajs.shape
        
        S = len(rgbs)
        H, W = rgbs[0].shape[:2]
        assert(S==T)

        H_new, W_new = self.crop_size[0], self.crop_size[1]

        sx_ = W_new / W
        sy_ = H_new / H
        rgbs = [cv2.resize(rgb, (W_new, H_new), interpolation=cv2.INTER_LINEAR) for rgb in rgbs]
        masks = [cv2.resize(mask, (W_new, H_new), interpolation=cv2.INTER_NEAREST) for mask in masks]
        sc_py = np.array([sx_, sy_]).reshape([1,1,2])
        trajs = trajs * sc_py
        
        return rgbs, masks, trajs

    def __len__(self):
        return len(self.rgb_paths)

datasets/pointodysseydataset_seg.py
aharley/pips2
from numpy import random
from numpy.core.numeric import full
import torch
import numpy as np
import pickle

class TapVidDavis(torch.utils.data.Dataset):
    def __init__(self, dataset_location='../datasets/tapvid_davis'):

        print('loading TAPVID-DAVIS dataset...')

        input_path = '%s/tapvid_davis.pkl' % dataset_location
        with open(input_path, 'rb') as f:
            data = pickle.load(f)
            if isinstance(data, dict):
                data = list(data.values())
        self.data = data
        print('found %d videos in %s' % (len(self.data), dataset_location))
        
    def __getitem__(self, index):
        dat = self.data[index]
        rgbs = dat['video'] # list of H,W,C uint8 images
        trajs = dat['points'] # N,S,2 array
        valids = 1-dat['occluded'] # N,S array
        # note the annotations are only valid when not occluded
        
        trajs = trajs.transpose(1,0,2) # S,N,2
        valids = valids.transpose(1,0) # S,N

        vis_ok = valids[0] > 0
        trajs = trajs[:,vis_ok]
        valids = valids[:,vis_ok]

        # 1.0,1.0 should lie at the bottom-right corner pixel
        H, W, C = rgbs[0].shape
        trajs[:,:,0] *= W-1
        trajs[:,:,1] *= H-1

        rgbs = torch.from_numpy(np.stack(rgbs,0)).permute(0,3,1,2) # S,C,H,W
        trajs = torch.from_numpy(trajs) # S,N,2
        valids = torch.from_numpy(valids) # S,N

        sample = {
            'rgbs': rgbs,
            'trajs': trajs,
            'valids': valids,
            'visibs': valids,
        }
        return sample

    def __len__(self):
        return len(self.data)



datasets/tapviddataset_fullseq.py
aharley/pips2
import time
import numpy as np
import saverloader
from nets.pips2 import Pips
import utils.improc
from utils.basic import print_, print_stats
import torch
from tensorboardX import SummaryWriter
import torch.nn.functional as F
from fire import Fire
import sys
import cv2
from pathlib import Path

def read_mp4(fn):
    vidcap = cv2.VideoCapture(fn)
    frames = []
    while(vidcap.isOpened()):
        ret, frame = vidcap.read()
        if ret == False:
            break
        frames.append(frame)
    vidcap.release()
    return frames

def run_model(model, rgbs, S_max=128, N=64, iters=16, sw=None):
    rgbs = rgbs.cuda().float() # B, S, C, H, W

    B, S, C, H, W = rgbs.shape
    assert(B==1)

    # pick N points to track; we'll use a uniform grid
    N_ = np.sqrt(N).round().astype(np.int32)
    grid_y, grid_x = utils.basic.meshgrid2d(B, N_, N_, stack=False, norm=False, device='cuda')
    grid_y = 8 + grid_y.reshape(B, -1)/float(N_-1) * (H-16)
    grid_x = 8 + grid_x.reshape(B, -1)/float(N_-1) * (W-16)
    xy0 = torch.stack([grid_x, grid_y], dim=-1) # B, N_*N_, 2
    _, S, C, H, W = rgbs.shape

    # zero-vel init
    trajs_e = xy0.unsqueeze(1).repeat(1,S,1,1)

    iter_start_time = time.time()
    
    preds, preds_anim, _, _ = model(trajs_e, rgbs, iters=iters, feat_init=None, delta_mult=0.5)
    trajs_e = preds[-1]

    iter_time = time.time()-iter_start_time
    print('inference time: %.2f seconds (%.1f fps)' % (iter_time, S/iter_time))

    if sw is not None and sw.save_this:
        rgbs_prep = utils.improc.preprocess_color(rgbs)
        sw.summ_traj2ds_on_rgbs('outputs/trajs_on_rgbs', trajs_e[0:1], utils.improc.preprocess_color(rgbs[0:1]), cmap='hot', linewidth=1, show_dots=False)
    return trajs_e


def main(
        filename='./stock_videos/camel.mp4',
        S=48, # seqlen
        N=1024, # number of points per clip
        stride=8, # spatial stride of the model
        timestride=1, # temporal stride of the model
        iters=16, # inference steps of the model
        image_size=(512,896), # input resolution
        max_iters=4, # number of clips to run
        shuffle=False, # dataset shuffling
        log_freq=1, # how often to make image summaries
        log_dir='./logs_demo',
        init_dir='./reference_model',
        device_ids=[0],
):

    # the idea in this file is to run the model on a demo video,
    # and return some visualizations
    
    exp_name = 'de00' # copy from dev repo

    print('filename', filename)
    name = Path(filename).stem
    print('name', name)
    
    rgbs = read_mp4(filename)
    rgbs = np.stack(rgbs, axis=0) # S,H,W,3
    rgbs = rgbs[:,:,:,::-1].copy() # BGR->RGB
    rgbs = rgbs[::timestride]
    S_here,H,W,C = rgbs.shape
    print('rgbs', rgbs.shape)

    # autogen a name
    model_name = "%s_%d_%d_%s" % (name, S, N, exp_name)
    import datetime
    model_date = datetime.datetime.now().strftime('%H:%M:%S')
    model_name = model_name + '_' + model_date
    print('model_name', model_name)
    
    log_dir = 'logs_demo'
    writer_t = SummaryWriter(log_dir + '/' + model_name + '/t', max_queue=10, flush_secs=60)

    global_step = 0

    model = Pips(stride=8).cuda()
    parameters = list(model.parameters())
    if init_dir:
        _ = saverloader.load(init_dir, model)
    global_step = 0
    model.eval()

    idx = list(range(0, max(S_here-S,1), S))
    if max_iters:
        idx = idx[:max_iters]
    
    for si in idx:
        global_step += 1
        
        iter_start_time = time.time()

        sw_t = utils.improc.Summ_writer(
            writer=writer_t,
            global_step=global_step,
            log_freq=log_freq,
            fps=16,
            scalar_freq=int(log_freq/2),
            just_gif=True)

        rgb_seq = rgbs[si:si+S]
        rgb_seq = torch.from_numpy(rgb_seq).permute(0,3,1,2).to(torch.float32) # S,3,H,W
        rgb_seq = F.interpolate(rgb_seq, image_size, mode='bilinear').unsqueeze(0) # 1,S,3,H,W
        
        with torch.no_grad():
            trajs_e = run_model(model, rgb_seq, S_max=S, N=N, iters=iters, sw=sw_t)

        iter_time = time.time()-iter_start_time
        
        print('%s; step %06d/%d; itime %.2f' % (
            model_name, global_step, max_iters, iter_time))
        
            
    writer_t.close()

if __name__ == '__main__':
    Fire(main)

demo.py
aharley/pips2
import time
import numpy as np
import timeit
import utils.improc
import utils.geom
import random
from utils.basic import print_, print_stats
import torch
import torch.nn as nn
from tensorboardX import SummaryWriter
from datasets.pointodysseydataset import PointOdysseyDataset
import torch.nn.functional as F
from fire import Fire
import sys
import os
from PIL import Image
from torch.utils.data import Dataset, DataLoader

# random.seed(125)
# np.random.seed(125)
# torch.manual_seed(125)

def run_model(d, exp_out_dir, device, mname, export=True, sw=None):
    metrics = {}

    rgbs = d['rgbs'].to(device).float().unsqueeze(0) # B,S,C,H,W
    trajs = d['trajs'].to(device).float().unsqueeze(0) # B,S,N,2
    visibs = d['visibs'].to(device).float().unsqueeze(0) # B,S,N
    valids = d['valids'].to(device).float().unsqueeze(0) # B,S,N

    B,S,C,H,W = rgbs.shape
    assert(C==3)
    
    B,S,N,D = trajs.shape
    assert(D==2)
    
    if torch.sum(valids)<B*S*N//2:
        sys.stdout.write('x')
        sys.stdout.flush()
        return False

    track_g = torch.cat([trajs, visibs.unsqueeze(-1), valids.unsqueeze(-1)], dim=3)

    temp_dir = 'temp_%s' % mname
    if not os.path.exists(temp_dir):
        os.makedirs(temp_dir)
        
    npz_out_f = '%s/track.npz' % (exp_out_dir)
    rgb_out_f = '%s/rgb.mp4' % (exp_out_dir)

    if (export and
        os.path.isfile(npz_out_f) and
        os.path.isfile(rgb_out_f)):
        sys.stdout.write(':')
        sys.stdout.flush()
        return False
    
    if sw is not None and sw.save_this:
        prep_rgbs = utils.improc.preprocess_color(rgbs)
        prep_grays = prep_rgbs.mean(dim=2, keepdim=True).repeat(1,1,3,1,1)
        sw.summ_traj2ds_on_rgb('0_inputs/trajs_on_rgbs', trajs[0:1], utils.improc.preprocess_color(rgbs[0:1].mean(dim=1)), cmap='winter', linewidth=2)
        sw.summ_traj2ds_on_rgbs2('0_inputs/trajs_on_rgbs2', trajs[0:1,::4], visibs[0:1,::4], prep_rgbs[0:1,::4], valids=valids[0:1,::4], frame_ids=list(range(0,S,4)))
        # sw.summ_traj2ds_on_rgbs('0_inputs/trajs_on_rgbs', trajs[0:1], utils.improc.preprocess_color(rgbs[0:1].mean(dim=1)), cmap='winter', linewidth=2)

    rgbs = rgbs[0].byte().cpu().numpy() # S,3,H,W
    track_g = track_g[0].cpu().numpy() # S,N,2

    rgbs = rgbs.transpose(0,2,3,1) # channels last

    if export:
        if not os.path.exists(exp_out_dir):
            os.makedirs(exp_out_dir)
        
        for si in range(S):
            temp_out_f = '%s/%03d.jpeg' % (temp_dir, si)
            im = Image.fromarray(rgbs[si])
            im.save(temp_out_f, "JPEG")
        os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate 24 -pattern_type glob -i "./%s/*.jpeg" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (temp_dir, rgb_out_f))

        # save npz only if we made all the way through
        if os.path.isfile(rgb_out_f):
            np.savez_compressed(
                npz_out_f,
                track_g=track_g,
            )
            sys.stdout.write('.')
            sys.stdout.flush()
            return True
        else:
            sys.stdout.write('f')
            sys.stdout.flush()
            return False
    return False

def main(
        dset='train', 
        S=36, # seqlen
        N=128, # number of particles to export per clip
        crop_size=(384,512), 
        use_augs=True, # resizing/jittering/color/blur augs
        shuffle=False, # dataset shuffling
        log_dir='./logs_export_mp4_dataset',
        dataset_location='/orion/group/point_odyssey',
        max_iters=0,
        log_freq=100,
        device_ids=[0],
):
    device = 'cpu:%d' % device_ids[0]

    # the idea in this file is:
    # walk through the dataset,
    # and export rgb mp4s for all valid samples,
    # so that dataloading is not a bottleneck in training
    
    exp_name = 'em00' # copy from dev repo
    
    mod = 'aa' # copy from dev repo; crop_size=(256x384), S=36
    mod = 'ab' # allow more OOB, by updating threshs to 64; export at 384,512; output 256; export as long as we have N//2
    mod = 'ac' # N=128
    mod = 'ad' # put more info into name; also print rtime
    mod = 'ae' # allow trajs to go behind camera during S

    assert(crop_size[0] % 64 == 0)
    assert(crop_size[1] % 64 == 0)
    
    # autogen a descriptive name
    model_name = "%d" % (S)
    if use_augs:
        model_name += "_A"
    model_name += "_%s" % exp_name
    model_name += "_%s" % mod
    import datetime
    model_date = datetime.datetime.now().strftime('%H%M%S')
    model_name = model_name + '_' + model_date
    print('model_name', model_name)
    
    writer_t = SummaryWriter(log_dir + '/' + model_name + '/t', max_queue=10, flush_secs=60)

    def worker_init_fn(worker_id):
        np.random.seed(np.random.get_state()[1][0] + worker_id)
    dataset_t = PointOdysseyDataset(
        dataset_location=dataset_location,
        dset=dset,
        S=S,
        N=N,
        use_augs=use_augs,
        crop_size=crop_size,
        verbose=True,
    )
    dataloader_t = DataLoader(
        dataset_t,
        batch_size=1,
        shuffle=shuffle,
        num_workers=4,
        worker_init_fn=worker_init_fn,
        drop_last=True)
    iterloader_t = iter(dataloader_t)

    if max_iters==0:
        max_iters = len(dataloader_t) # number of samples to export

    perm = np.random.permutation(max_iters) # write in random order, for parallel
    # perm = np.arange(max_iters) # write in sequential order, for debug

    global_step = 0
    while global_step < max_iters:
        global_step += 1
        this_step = perm[global_step-1]
        iter_start_time = time.time()
        sw_t = utils.improc.Summ_writer(
            writer=writer_t,
            global_step=this_step,
            log_freq=log_freq,
            fps=4,
            scalar_freq=int(log_freq/4),
            just_gif=True)

        H, W = crop_size
        out_dir = './pod_export/%s_%d_%d_%dx%d' % (mod, S, N, H, W)
        exp_out_dir = '%s/%06d' % (out_dir, this_step)

        npz_out_f = '%s/track.npz' % (exp_out_dir)
        rgb_out_f = '%s/rgb.mp4' % (exp_out_dir)

        if (os.path.isfile(npz_out_f) and
            os.path.isfile(rgb_out_f)):
            sys.stdout.write(':')
            sys.stdout.flush()
        else:
            sample, gotit = dataset_t.__getitem__(this_step % len(dataloader_t))
            if gotit:
                iter_rtime = time.time()-iter_start_time
                out = run_model(sample, exp_out_dir, device, model_name, sw=sw_t)
                iter_itime = time.time()-iter_start_time
                if out:
                    print('%s; step %06d/%d; this_step %06d; rtime %.2f; itime %.2f' % (
                        model_name, global_step, max_iters, this_step, iter_rtime, iter_itime))
    writer_t.close()
            

if __name__ == '__main__':
    Fire(main)

export_mp4_dataset.py
aharley/pips2
import time
import numpy as np
import timeit
import saverloader
from nets.pips2 import Pips
import utils.improc
import utils.geom
import utils.misc
import random
from utils.basic import print_, print_stats
from datasets.pointodysseydataset import PointOdysseyDataset
import torch
import torch.nn as nn
from tensorboardX import SummaryWriter
import torch.nn.functional as F
from fire import Fire
import sys
from torch import nn, einsum
from einops import rearrange, repeat
from einops.layers.torch import Rearrange, Reduce
from torch.utils.data import Dataset, DataLoader

def create_pools(n_pool=1000):
    pools = {}
    pool_names = [
        'l1',
        'd_1',
        'd_2',
        'd_4',
        'd_8',
        'd_16',
        'd_avg',
        'l1_vis',
        'ate_all',
        'ate_vis',
        'ate_occ',
        'median_l2',
        'survival',
        'total_loss',
    ]
    for pool_name in pool_names:
        pools[pool_name] = utils.misc.SimplePool(n_pool, version='np')
    return pools

def requires_grad(parameters, flag=True):
    for p in parameters:
        p.requires_grad = flag

def fetch_optimizer(lr, wdecay, epsilon, num_steps, params):
    optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=wdecay, eps=epsilon)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, lr, num_steps+100, pct_start=0.1, cycle_momentum=False, anneal_strategy='linear')
    return optimizer, scheduler

def val_model(model, d, device, iters=8, sw=None, is_train=False):
    metrics = {}

    rgbs = d['rgbs'].float().to(device) # B,S,C,H,W
    trajs_g = d['trajs'].float().to(device) # B,S,N,8
    vis_g = d['visibs'].float().to(device) # B,S,N
    valids = d['valids'].float().to(device) # B,S,N

    B, S, C, H, W = rgbs.shape
    B, S, N, D = trajs_g.shape
    assert(D==2)

    # zero-vel init
    trajs_e0 = trajs_g[:,0:1].repeat(1,S,1,1)

    preds, preds_anim, _, _ = model(trajs_e0, rgbs, iters=iters)
    trajs_e = preds[-1]

    l1_dists = torch.abs(trajs_e - trajs_g).sum(dim=-1) # B,S,N
    l1_loss = utils.basic.reduce_masked_mean(l1_dists, valids)
    l1_vis = utils.basic.reduce_masked_mean(l1_dists, valids*vis_g)
    
    ate = torch.norm(trajs_e - trajs_g, dim=-1) # B,S,N
    ate_all = utils.basic.reduce_masked_mean(ate, valids, dim=[1,2])
    ate_vis = utils.basic.reduce_masked_mean(ate, valids*vis_g)
    ate_occ = utils.basic.reduce_masked_mean(ate, valids*(1.0-vis_g))

    metrics['l1'] = l1_loss.mean().item()
    metrics['l1_vis'] = l1_vis.mean().item()
    metrics['ate_all'] = ate_all.mean().item()
    metrics['ate_vis'] = ate_vis.item()
    metrics['ate_occ'] = ate_occ.item()
    
    if sw is not None and sw.save_this:
        prep_rgbs = utils.improc.preprocess_color(rgbs)
        prep_grays = torch.mean(prep_rgbs, dim=2, keepdim=True).repeat(1, 1, 3, 1, 1)
        gt_rgb = utils.improc.preprocess_color(sw.summ_traj2ds_on_rgb('', trajs_g[0:1], prep_grays[0:1].mean(dim=1), valids=valids[0:1], cmap='winter', only_return=True))
        rgb_vis = []
        for tre in preds_anim:
            ate = torch.norm(tre - trajs_g, dim=-1) # B,S,N
            ate_all = utils.basic.reduce_masked_mean(ate, valids, dim=[1,2]) # B
            rgb_vis.append(sw.summ_traj2ds_on_rgb('', tre[0:1], gt_rgb, valids=valids[0:1], only_return=True, cmap='spring', frame_id=ate_all[0]))
        sw.summ_rgbs('3_test/animated_trajs_on_rgb', rgb_vis)
        
    d_sum = 0.0
    thrs = [1,2,4,8,16]
    sx_ = W / 256.0
    sy_ = H / 256.0
    sc_py = np.array([sx_, sy_]).reshape([1,1,2])
    sc_pt = torch.from_numpy(sc_py).float().cuda()
    for thr in thrs:
        # note we exclude timestep0 from this eval
        d_ = (torch.norm(trajs_e[:,1:]/sc_pt - trajs_g[:,1:]/sc_pt, dim=-1) < thr).float() # B,S-1,N
        d_ = utils.basic.reduce_masked_mean(d_, valids[:,1:]).item()*100.0
        d_sum += d_
        metrics['d_%d' % thr] = d_
    d_avg = d_sum / len(thrs)
    metrics['d_avg'] = d_avg

    sur_thr = 16
    dists = torch.norm(trajs_e/sc_pt - trajs_g/sc_pt, dim=-1) # B,S,N
    dist_ok = 1 - (dists > sur_thr).float() * valids # B,S,N
    survival = torch.cumprod(dist_ok, dim=1) # B,S,N
    metrics['survival'] = torch.mean(survival).item()*100.0
    
    # get the median l2 error for each trajectory
    dists_ = dists.permute(0,2,1).reshape(B*N,S)
    valids_ = valids.permute(0,2,1).reshape(B*N,S)
    median_l2 = utils.basic.reduce_masked_median(dists_, valids_, keep_batch=True) # B*N
    metrics['median_l2'] = median_l2.mean().item()

    return metrics

def run_model(model, d, device, iters=8, sw=None, is_train=True, use_augs=True):
    total_loss = torch.tensor(0.0, requires_grad=True).to(device)

    metrics = {}

    rgbs = d['rgbs'].float().to(device) # B,S,C,H,W
    trajs_g = d['trajs'].float().to(device) # B,S,N,8
    vis_g = d['visibs'].float().to(device) # B,S,N
    valids = d['valids'].float().to(device) # B,S,N

    if use_augs and np.random.rand() < 0.5: # rot90 aug
        rgbs = rgbs.permute(0,1,2,4,3) # swap xy
        trajs_g = trajs_g.flip([3]) # swap xy
        
    B, S, C, H, W = rgbs.shape
    assert(C==3)
    B, S, N, D = trajs_g.shape
    assert(D==2)

    # full random
    x = torch.from_numpy(np.random.uniform(0, W-1, (B,S,N))).float().to(trajs_g.device)
    y = torch.from_numpy(np.random.uniform(0, H-1, (B,S,N))).float().to(trajs_g.device)
    trajs_e0 = torch.stack([x,y], dim=-1) # B,S,N,2
    # mix with gt a random amount
    coeff = torch.from_numpy(np.random.uniform(0, 1, (B,1,N,1))).float().to(trajs_g.device)
    trajs_e0 = trajs_e0*coeff + trajs_g*(1-coeff)
    # use zero-velocity init for some
    trajs_z = trajs_g[:,0:1].repeat(1,S,1,1)
    mask = (torch.from_numpy(np.random.uniform(0, 1, (B,1,N,1))).float().to(trajs_g.device)>0.5).float()
    trajs_e0 = trajs_e0*mask + trajs_z*(1-mask)
    # reset zeroth on all
    trajs_e0[:,0:1] = trajs_g[:,0:1]

    # # zero-vel init
    # trajs_e0 = trajs_g[:,0:1].repeat(1,S,1,1)
    

    # measure our initial distance, so we can check our improvement
    ate0 = torch.norm(trajs_e0 - trajs_g, dim=-1) # B,S,N
    ate0_all = utils.basic.reduce_masked_mean(ate0, valids, dim=[1,2])
    
    preds, preds_anim, _, loss = model(trajs_e0, rgbs, iters=iters, trajs_g=trajs_g, vis_g=vis_g, valids=valids, is_train=is_train)
    trajs_e = preds[-1]

    total_loss += loss

    # collect stats
    l1_dists = torch.abs(trajs_e - trajs_g).sum(dim=-1) # B,S,N
    l1_loss = utils.basic.reduce_masked_mean(l1_dists, valids)
    l1_vis = utils.basic.reduce_masked_mean(l1_dists, valids*vis_g)
    ate = torch.norm(trajs_e - trajs_g, dim=-1) # B,S,N
    ate_all = utils.basic.reduce_masked_mean(ate, valids, dim=[1,2])
    ate_vis = utils.basic.reduce_masked_mean(ate, valids*vis_g)
    ate_occ = utils.basic.reduce_masked_mean(ate, valids*(1.0-vis_g))
    metrics['l1'] = l1_loss.mean().item()
    metrics['l1_vis'] = l1_vis.mean().item()
    metrics['ate_all'] = ate_all.mean().item()
    metrics['ate_vis'] = ate_vis.item()
    metrics['ate_occ'] = ate_occ.item()
    metrics['total_loss'] = total_loss.item()

    d_sum = 0.0
    thrs = [1,2,4,8,16]
    sx_ = W / 256.0
    sy_ = H / 256.0
    sc_py = np.array([sx_, sy_]).reshape([1,1,2])
    sc_pt = torch.from_numpy(sc_py).float().to(device)
    for thr in thrs:
        # note we exclude timestep0 from this eval
        d_ = (torch.norm(trajs_e[:,1:]/sc_pt - trajs_g[:,1:]/sc_pt, dim=-1) < thr).float() # B,S-1,N
        d_ = utils.basic.reduce_masked_mean(d_, valids[:,1:]).item()*100.0
        d_sum += d_
        metrics['d_%d' % thr] = d_
    d_avg = d_sum / len(thrs)
    metrics['d_avg'] = d_avg

    sur_thr = 16
    dists = torch.norm(trajs_e/sc_pt - trajs_g/sc_pt, dim=-1) # B,S,N
    dist_ok = 1 - (dists > sur_thr).float() * valids # B,S,N
    survival = torch.cumprod(dist_ok, dim=1) # B,S,N
    metrics['survival'] = torch.mean(survival).item()*100.0
    
    # get the median l2 error for each trajectory
    dists_ = dists.permute(0,2,1).reshape(B*N,S)
    valids_ = valids.permute(0,2,1).reshape(B*N,S)
    val_ok = valids_[:,0] > 0 # get rid of the ones we padded in
    dists_ = dists_[val_ok]
    valids_ = valids_[val_ok]
    median_l2 = utils.basic.reduce_masked_median(dists_, valids_, keep_batch=True) # B*N
    metrics['median_l2'] = median_l2.mean().item()

    if sw is not None and sw.save_this:
        prep_rgbs = utils.improc.preprocess_color(rgbs)
        prep_grays = torch.mean(prep_rgbs, dim=2, keepdim=True).repeat(1, 1, 3, 1, 1)

        rgb0 = sw.summ_traj2ds_on_rgb('', trajs_g[0:1], prep_rgbs[0:1,0], valids=valids[0:1], cmap='winter', linewidth=2, only_return=True)
        sw.summ_traj2ds_on_rgb('0_inputs/trajs_e0_on_rgb0', trajs_e0[0:1], utils.improc.preprocess_color(rgb0), valids=valids[0:1], cmap='spring', linewidth=2, frame_id=ate0_all[0].mean().item())
        sw.summ_traj2ds_on_rgb('2_outputs/trajs_e_on_rgb0', trajs_e[0:1], utils.improc.preprocess_color(rgb0), valids=valids[0:1], cmap='spring', linewidth=2, frame_id=ate_all[0].mean().item())
        sw.summ_traj2ds_on_rgbs2('0_inputs/trajs_g_on_rgbs2', trajs_g[0:1,::4], vis_g[0:1,::4], prep_rgbs[0:1,::4], valids=valids[0:1,::4], frame_ids=list(range(0,S,4)))
        
        # in the kp vis, clamp so that we can see everything
        trajs_g_clamp = trajs_g.clone()
        trajs_g_clamp[:,:,:,0] = trajs_g_clamp[:,:,:,0].clip(0,W-1)
        trajs_g_clamp[:,:,:,1] = trajs_g_clamp[:,:,:,1].clip(0,H-1)
        trajs_e_clamp = trajs_e.clone()
        trajs_e_clamp[:,:,:,0] = trajs_e_clamp[:,:,:,0].clip(0,W-1)
        trajs_e_clamp[:,:,:,1] = trajs_e_clamp[:,:,:,1].clip(0,H-1)

        gt_rgb = utils.improc.preprocess_color(sw.summ_traj2ds_on_rgb('', trajs_g[0:1], prep_grays[0:1].mean(dim=1), valids=valids[0:1], cmap='winter', only_return=True))
        rgb_vis = []
        for tre in preds_anim:
            ate = torch.norm(tre - trajs_g, dim=-1) # B,S,N
            ate_all = utils.basic.reduce_masked_mean(ate, valids, dim=[1,2]) # B
            rgb_vis.append(sw.summ_traj2ds_on_rgb('', tre[0:1], gt_rgb, valids=valids[0:1], only_return=True, cmap='spring', frame_id=ate_all[0]))
        sw.summ_rgbs('3_test/animated_trajs_on_rgb', rgb_vis)

        outs = sw.summ_pts_on_rgbs(
            '',
            trajs_g_clamp[0:1,::4],
            prep_grays[0:1,::4],
            valids=valids[0:1,::4],
            cmap='winter', linewidth=3, only_return=True)
        sw.summ_pts_on_rgbs(
            '0_inputs/kps_gv_on_rgbs',
            trajs_g_clamp[0:1,::4],
            utils.improc.preprocess_color(outs),
            valids=valids[0:1,::4]*vis_g[0:1,::4],
            cmap='spring', linewidth=2)

        outs = sw.summ_pts_on_rgbs(
            '',
            trajs_g_clamp[0:1,::4],
            prep_grays[0:1,::4],
            valids=valids[0:1,::4],
            cmap='winter', linewidth=3, only_return=True)
        sw.summ_pts_on_rgbs(
            '2_outputs/kps_eg_on_rgbs',
            trajs_e_clamp[0:1,::4],
            utils.improc.preprocess_color(outs),
            valids=valids[0:1,::4],
            cmap='spring', linewidth=2)
        
    return total_loss, metrics
    

def main(
        B=2, # batchsize 
        S=38, # seqlen
        N=64, # number of points per clip
        stride=8, # spatial stride of the model 
        iters=4, # inference steps of the model
        crop_size=(512,896), # raw flt data is 540,960
        use_augs=True, # resizing/jittering/color/blur augs
        shuffle=True, # dataset shuffling
        cache_len=0, # how many samples to cache into ram (for overfitting/debug)
        cache_freq=0, # how often to add a new sample to cache
        dataset_location='/orion/group/point_odyssey', 
        n_pool=1000, # size of running avg for stats
        quick=False, # debug
        # optimization
        lr=5e-4,
        grad_acc=1, 
        use_scheduler=True,
        max_iters=200000,
        # summaries
        log_dir='./logs_finetune_highres',
        log_freq=1000,
        val_freq=0,
        # saving/loading
        ckpt_dir='./checkpoints',
        save_freq=1000,
        keep_latest=2,
        init_dir='',
        load_optimizer=True,
        load_step=True,
        ignore_load=None,
        device_ids=[0],
):
    device = 'cuda:%d' % device_ids[0]

    # the idea in this file is:
    # finetune on pointodyssey, using the raw high-res jpgs

    exp_name = 'ac00' # copy from dev repo

    if quick: # (debug)
        B = 1
        log_freq = 100
        max_iters = 1000
        shuffle = False
        val_freq = 10
        n_pool = 10
        use_augs = False
        cache_len = 3 # overfit on this many
        cache_freq = 0
        save_freq = 99999999
        
    if init_dir:
        init_dir = '%s/%s' % (ckpt_dir, init_dir)
        
    assert(crop_size[0] % 32 == 0)
    assert(crop_size[1] % 32 == 0)
    
    # autogen a descriptive name
    model_name = "%d_%d_%d" % (B,S,N)
    model_name += "_i%d" % (iters)
    if grad_acc > 1:
        model_name += "x%d" % grad_acc
    lrn = "%.1e" % lr # e.g., 5.0e-04
    lrn = lrn[0] + lrn[3:5] + lrn[-1] # e.g., 5e-4
    model_name += "_%s" % lrn
    if use_scheduler:
        model_name += "s"
    if cache_len:
        model_name += "_c%d_f%d" % (cache_len, cache_freq)
    if use_augs:
        model_name += "_A"
    model_name += "_%s" % exp_name
    import datetime
    model_date = datetime.datetime.now().strftime('%H%M%S')
    model_name = model_name + '_' + model_date
    print('model_name', model_name)
    
    ckpt_path = '%s/%s' % (ckpt_dir, model_name)
    writer_t = SummaryWriter(log_dir + '/' + model_name + '/t', max_queue=10, flush_secs=60)
    if val_freq:
        writer_v = SummaryWriter(log_dir + '/' + model_name + '/v', max_queue=10, flush_secs=60)

    def worker_init_fn(worker_id):
        np.random.seed(np.random.get_state()[1][0] + worker_id)
    dataset_t = PointOdysseyDataset(
        dataset_location=dataset_location,
        dset='train',
        S=S,
        N=N,
        strides=[1,2],
        use_augs=use_augs,
        crop_size=crop_size,
        quick=quick,
        req_full=True,
        verbose=False,
    )
    dataloader_t = DataLoader(
        dataset_t,
        batch_size=B,
        shuffle=shuffle,
        num_workers=6,
        worker_init_fn=worker_init_fn,
        drop_last=True)
    iterloader_t = iter(dataloader_t)

    if cache_len:
        sample_pool_t = utils.misc.SimplePool(cache_len, version='np')
    
    model = Pips(stride=stride).to(device)
    model = torch.nn.DataParallel(model, device_ids=device_ids)

    parameters = list(model.parameters())
    weight_decay = 1e-6
    if use_scheduler:
        optimizer, scheduler = fetch_optimizer(lr, weight_decay, 1e-8, max_iters, model.parameters())
    else:
        optimizer = torch.optim.AdamW(parameters, lr=lr, weight_decay=weight_decay) 
        scheduler = None

    utils.misc.count_parameters(model)

    global_step = 0
    if init_dir:
        if load_step and load_optimizer:
            global_step = saverloader.load(init_dir, model.module, optimizer=optimizer, scheduler=scheduler, ignore_load=ignore_load)
        elif load_step:
            global_step = saverloader.load(init_dir, model.module, ignore_load=ignore_load)
        else:
            _ = saverloader.load(init_dir, model.module, ignore_load=ignore_load)
            global_step = 0
    
    requires_grad(parameters, True)
    model.train()

    pools_t = create_pools(n_pool)
    if val_freq:
        pools_v = create_pools(n_pool)
    
    while global_step < max_iters:
        global_step += 1
        iter_start_time = time.time()
        iter_rtime = 0.0
        
        for internal_step in range(grad_acc):
            read_start_time = time.time()

            if internal_step==grad_acc-1:
                sw_t = utils.improc.Summ_writer(
                    writer=writer_t,
                    global_step=global_step,
                    log_freq=log_freq,
                    fps=min(S,8),
                    scalar_freq=log_freq//4,
                    just_gif=True)
            else:
                sw_t = None

            read_new = True # read something from the dataloder
            if cache_len: 
                read_new = False
                if len(sample_pool_t) < cache_len:
                    read_new = True
                if cache_freq > 0 and global_step % cache_freq == 0:
                    read_new = True
            if read_new:
                gotit = (False,False)
                while not all(gotit):
                    try:
                        sample, gotit = next(iterloader_t)
                    except StopIteration:
                        iterloader_t = iter(dataloader_t)
                        sample, gotit = next(iterloader_t)
                if cache_len:
                    sample_pool_t.update([sample])
                    print('cached a new sample into sample_pool (len %d)' % (len(sample_pool_t)))
            if cache_len:
                sample = sample_pool_t.sample()
                    
            iter_rtime += time.time()-read_start_time

            total_loss, metrics = run_model(
                model, sample, device,
                iters=iters,
                sw=sw_t,
                is_train=True,
                use_augs=use_augs)

            if torch.isnan(total_loss):
                print('nan in loss; quitting')
                return False

            total_loss /= grad_acc
            total_loss.backward()

        sw_t.summ_scalar('total_loss', metrics['total_loss'])

        for key in list(pools_t.keys()):
            if key in metrics:
                pools_t[key].update([metrics[key]])
            sw_t.summ_scalar('_/%s' % (key), pools_t[key].mean())

        current_lr = optimizer.param_groups[0]['lr']
        sw_t.summ_scalar('_/current_lr', current_lr)
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        if use_scheduler:
            scheduler.step()
        optimizer.zero_grad()

        if np.mod(global_step, save_freq)==0:
            saverloader.save(ckpt_path, optimizer, model.module, global_step, scheduler=scheduler, keep_latest=keep_latest)

        if val_freq and (global_step) % val_freq == 0:
            model.eval()
            del sample
            with torch.no_grad():
                torch.cuda.empty_cache()
            sw_v = utils.improc.Summ_writer(
                writer=writer_v,
                global_step=global_step,
                log_freq=log_freq,
                fps=min(S,8),
                scalar_freq=log_freq//4,
                just_gif=True)
            if cache_len:
                sample = sample_pool_t.sample()
            else:
                gotit = (False,False)
                while not all(gotit):
                    try:
                        sample, gotit = next(iterloader_t)
                    except StopIteration:
                        iterloader_t = iter(dataloader_t)
                        sample, gotit = next(iterloader_t)
            with torch.no_grad():
                metrics = val_model(
                    model, sample, device,
                    iters=iters*2,
                    sw=sw_v,
                    is_train=False)
            for key in list(pools_v.keys()):
                if key in metrics:
                    pools_v[key].update([metrics[key]])
                sw_v.summ_scalar('_/%s' % (key), pools_v[key].mean())
            model.train()
                
                    
        iter_itime = time.time()-iter_start_time

        if val_freq:
            print('%s; step %06d/%d; rtime %.2f; itime %.2f; loss %.3f; loss_t %.2f; d_t %.1f; d_v %.1f' % (
                model_name, global_step, max_iters, iter_rtime, iter_itime,
                total_loss.item(), pools_t['total_loss'].mean(), pools_t['d_avg'].mean(), pools_v['d_avg'].mean()))
        else:
            print('%s; step %06d/%d; rtime %.2f; itime %.2f; loss %.3f; loss_t %.2f; d_t %.1f' % (
                model_name, global_step, max_iters, iter_rtime, iter_itime,
                total_loss.item(), pools_t['total_loss'].mean(), pools_t['d_avg'].mean()))
            
    writer_t.close()
    if val_freq:
        writer_v.close()

if __name__ == '__main__':
    Fire(main)

finetune_highres.py
aharley/pips2
import time
import numpy as np
import timeit
import saverloader
import utils.improc
import utils.geom
import utils.misc
import random
from utils.basic import print_, print_stats
from datasets.pointodysseydataset_seg import PointOdysseyDataset
import torch
import torch.nn as nn
from tensorboardX import SummaryWriter
import torch.nn.functional as F
from fire import Fire
import sys
from torch import nn, einsum
from einops import rearrange, repeat
from einops.layers.torch import Rearrange, Reduce
from torch.utils.data import Dataset, DataLoader

random.seed(125)
np.random.seed(125)
torch.manual_seed(125)
    
def run_model(d, device, sw=None):
    rgbs = d['rgbs'].to(device).float() # B,S,C,H,W
    masks = d['masks'].to(device).float() # B,S,1,H,W
    trajs_g = d['trajs'].to(device).float() # B,S,N,2
    vis_g = d['visibs'].to(device).float() # B,S,N
    valids = d['valids'].to(device).float() # B,S,N

    print_stats('rgbs', rgbs)
    print_stats('masks', masks)
    print_stats('trajs_g', trajs_g)
    print_stats('vis_g', vis_g)
    print_stats('valids', valids)

    B, S, C, H, W = rgbs.shape
    assert(C==3)
    B, S, N, D = trajs_g.shape
    assert(D==2)
    
    
    if sw is not None and sw.save_this:

        prep_rgbs = utils.improc.preprocess_color(rgbs)
        prep_grays = torch.mean(prep_rgbs, dim=2, keepdim=True).repeat(1, 1, 3, 1, 1)

        sw.summ_traj2ds_on_rgb('0_inputs/trajs_g_on_rgb', trajs_g[0:1], prep_rgbs.mean(dim=1), valids=valids[0:1], cmap='winter')
        # sw.summ_traj2ds_on_rgbs('0_inputs/trajs_g_on_rgbs', trajs_g[0:1], prep_rgbs, valids=valids[0:1])
        sw.summ_traj2ds_on_rgbs2('0_inputs/trajs_g_on_rgbs2', trajs_g[0:1], vis_g[0:1], utils.improc.preprocess_color(rgbs[0:1]), valids=valids[0:1])

        mask_max = torch.max(masks)
        masks_vis = masks/(1e-4 + mask_max)
        sw.summ_oneds('0_inputs/masks', masks_vis.unbind(1), norm=False)

        label_colors = utils.improc.get_n_colors(int(mask_max.item())+1, sequential=False)

        mask_vis = []
        for si in range(S):
            mask_vis.append(sw.summ_seg('', masks[:,si,0], only_return=True, label_colors=label_colors))
        sw.summ_rgbs('0_inputs/masks_colored', mask_vis)
        
        # for the kp vis, we will clamp so that we can see everything
        trajs_g_clamp = trajs_g.clone()
        trajs_g_clamp[:,:,:,0] = trajs_g_clamp[:,:,:,0].clip(0,W-1)
        trajs_g_clamp[:,:,:,1] = trajs_g_clamp[:,:,:,1].clip(0,H-1)
        
        outs = sw.summ_pts_on_rgbs(
            '',
            trajs_g_clamp[0:1],
            prep_grays[0:1],
            valids=valids[0:1],
            cmap='winter', linewidth=3, only_return=True)
        sw.summ_pts_on_rgbs(
            '0_inputs/kps_gv_on_rgbs',
            trajs_g_clamp[0:1],
            utils.improc.preprocess_color(outs),
            valids=valids[0:1]*vis_g[0:1],
            cmap='spring', linewidth=2)

    return None 
    

def main(
        exp_name='debug',
        dset='train',
        B=1, # batchsize 
        S=32, # seqlen
        N=128, # number of points per clip
        crop_size=(256,448), 
        use_augs=False, # resizing/jittering/color/blur augs
        shuffle=False, # dataset shuffling
        log_dir='./logs_just_vis_2d',
        dataset_location='/orion/group/point_odyssey', 
        log_freq=1,
        max_iters=10,
        quick=False,
        dname=None,
):
    device = 'cpu:0'

    # the idea in this file is:
    # load pointodyssey data and visualize it
    
    exp_name = 'jv00' # copy from dev repo
    exp_name = 'jv01' # fix color bug...
    
    import socket
    host = socket.gethostname()

    assert(crop_size[0] % 32 == 0)
    assert(crop_size[1] % 32 == 0)
    
    # autogen a descriptive name
    model_name = "%d_%d_%d" % (B, S, N)
    if use_augs:
        model_name += "_A"
    model_name += "_%s" % exp_name
    import datetime
    model_date = datetime.datetime.now().strftime('%H:%M:%S')
    model_name = model_name + '_' + model_date
    print('model_name', model_name)
    
    writer_t = SummaryWriter(log_dir + '/' + model_name + '/t', max_queue=10, flush_secs=60)

    # get dataset
    def worker_init_fn(worker_id):
        np.random.seed(np.random.get_state()[1][0] + worker_id)
    dataset_t = PointOdysseyDataset(
        dataset_location=dataset_location,
        dset=dset,
        S=S,
        N=N,
        use_augs=use_augs,
        crop_size=crop_size,
        quick=quick,
        verbose=True,
    )
    dataloader_t = DataLoader(
        dataset_t,
        batch_size=B,
        shuffle=shuffle,
        num_workers=0,
        worker_init_fn=worker_init_fn,
        drop_last=True)
    iterloader_t = iter(dataloader_t)

    
    global_step = 0
    while global_step < max_iters:
        global_step += 1
        iter_start_time = time.time()
        sw_t = utils.improc.Summ_writer(
            writer=writer_t,
            global_step=global_step,
            log_freq=log_freq,
            fps=min(S,8),
            scalar_freq=log_freq//5,
            just_gif=True)

        gotit = (False,False)
        while not all(gotit):
            try:
                sample, gotit = next(iterloader_t)
            except StopIteration:
                iterloader_t = iter(dataloader_t)
                sample, gotit = next(iterloader_t)

        iter_rtime = time.time()-iter_start_time
        _ = run_model(sample, device, sw=sw_t)
        iter_itime = time.time()-iter_start_time
        
        print('%s; step %06d/%d; rtime %.2f; itime %.2f' % (
            model_name, global_step, max_iters, iter_rtime, iter_itime))
            
    writer_t.close()
            

if __name__ == '__main__':
    Fire(main)

just_vis_2d.py
aharley/pips2
import time
import numpy as np
import timeit
import saverloader
import utils.improc
import utils.geom
import utils.misc
import utils.vox
import random
from utils.basic import print_, print_stats
from datasets.pointodysseydataset_3d import PointOdysseyDataset
import torch
import torch.nn as nn
from tensorboardX import SummaryWriter
import torch.nn.functional as F
from fire import Fire
import sys
from torch import nn, einsum
from einops import rearrange, repeat
from einops.layers.torch import Rearrange, Reduce
from torch.utils.data import Dataset, DataLoader

# define the scene centroid for 3d vox
scene_centroid_x = 0.0 # right
scene_centroid_y = 0.0 # down
scene_centroid_z = 4.0 # forward
scene_centroid_py = np.array([scene_centroid_x,
                              scene_centroid_y,
                              scene_centroid_z]).reshape([1, 3])
scene_centroid = torch.from_numpy(scene_centroid_py).float()

# define a volume around the centroid
XMIN, XMAX = -4, 4
ZMIN, ZMAX = -4, 4
YMIN, YMAX = -4, 4
bounds = (XMIN, XMAX, YMIN, YMAX, ZMIN, ZMAX)

Z, Y, X = 500, 50, 500

    
def run_model(d, device, sw=None):
    rgbs = d['rgbs'].to(device).float() # B,S,3,H,W
    depths = d['depths'].to(device).float() # B,S,1,H,W
    normals = d['normals'].to(device).float() # B,S,3,H,W
    trajs_x = d['trajs_pix'].to(device).float() # B,S,N,2
    trajs_g = d['trajs_2d'].to(device).float() # B,S,N,2
    vis_g = d['visibs'].to(device).float() # B,S,N
    valids = d['valids'].to(device).float() # B,S,N

    pix_T_cams = d['pix_T_cams'].to(device).float() # B,S,4,4
    cams_T_world = d['cams_T_world'].to(device).float() # B,S,4,4

    print('rgbs', rgbs.shape)
    print('trajs_x', trajs_x.shape)
    print('trajs_g', trajs_g.shape)
    print('vis_g', vis_g.shape)
    print('valids', valids.shape, torch.sum(valids[:,0]))
    print('pix_T_cams', pix_T_cams.shape)
    print('cams_T_world', cams_T_world.shape)

    B, S, C, H, W = rgbs.shape
    assert(C==3)
    B, S, N, D = trajs_g.shape
    assert(D==2)

    assert(B==1)


    __p = lambda x: utils.basic.pack_seqdim(x, B)
    __u = lambda x: utils.basic.unpack_seqdim(x, B)
    
    # if torch.median(depth) > 100:

    normals = utils.improc.preprocess_color(normals)
    
    print_stats('normals', normals)

    depths_valid = (depths < 1000.0).float() * (depths > 0.0).float()

    # med = utils.basic.reduce_masked_median(depths[:,0], depths_valid[:,0])
    med = utils.basic.reduce_masked_median(depths, depths_valid)
    # med = torch.median(depths[:,0])
    print_('masked median depth', med)

    if med > 10:
        xyz_cam0 = utils.geom.depth2pointcloud(depths[:,0], pix_T_cams[:,0])
        # print('xyz_cam0[0,:10] bef', xyz_cam0[0,:10])
        
        depths = depths / 12.0
        xyz_cam0 = utils.geom.depth2pointcloud(depths[:,0], pix_T_cams[:,0])
        # print('xyz_cam0[0,:10] aft', xyz_cam0[0,:10])
        # input()

        # so,
        # it seems all of the values are divided by 12
        # which makes perfect sense via the math


        # what i need to do is:
        # undo the scale temporarily, apply the extrinsics, and redo it
        
        # sc = utils.geom.eye_3x3(B*S, device=device)
        sc = utils.geom.eye_4x4(B*S, device=device)
        sc[:,:3,:3] /= 12.0
        
        # print_('sc[0]', sc[0])
        # print_('inverse(sc[0])', sc.inverse()[0])
        # print_stats('cams_T_world before', cams_T_world)
        
        # cams_T_world = __u(utils.geom.matmul2(sc, __p(cams_T_world)))
        cams_T_world = __u(utils.geom.matmul3(sc, __p(cams_T_world), sc.inverse()))

        # print_stats('cams_T_world after', cams_T_world)
        
        # print_('sc[0]', sc[0])
        # # cams_T_world = __u(utils.geom.matmul2(sc, __p(cams_T_world)))
        # pix_T_cams = __u(utils.geom.matmul2(sc, __p(pix_T_cams)))
        # pix_T_cams = __u(utils.geom.matmul2(__p(pix_T_cams), sc))

        
        # cams_T_world = __u(utils.geom.matmul2(sc, __p(cams_T_world)))

        # med = utils.basic.reduce_masked_median(depths, depths_valid)
        # # med = torch.median(depths[:,0])
        # print_('new masked median depth', med)

    vox_util = utils.vox.Vox_util(
        Z, Y, X,
        scene_centroid=scene_centroid.to(device),
        bounds=bounds,
        assert_cube=False)

    # world_T_cams = __u(utils.geom.safe_inverse(__p(cams_T_world)))
    # utils.geom.get_camM_T_camXs(origin_T_camXs, ind=0)
    # cam0_T_camXs = utils.geom.get_camM_T_camXs(velo_T_cams, ind=0)
    # camXs_T_cam0 = __u(utils.geom.safe_inverse(__p(cam0_T_camXs)))

    cam0_T_world = cams_T_world[:,0]

    occ0_vis = []
    occI_vis = []
    for si in range(S):

        # cam0_T_camXs = utils.geom.get_camM_T_camXs(velo_T_cams, ind=0)

        # camI_T_world = cams_T_world[:,si]
        world_T_camI = utils.geom.safe_inverse(cams_T_world[:,si])
        cam0_T_camI = utils.geom.matmul2(cam0_T_world, world_T_camI)

        xyz_camI = utils.geom.depth2pointcloud(depths[:,si], pix_T_cams[:,si])
        depth_valid_ = depths_valid[:,si].reshape(-1)
        xyz_camI = xyz_camI[:,depth_valid_>0]
        xyz_cam0 = utils.geom.apply_4x4(cam0_T_camI, xyz_camI)

        if si==0:
            print_stats('xyz_camI', xyz_camI)
            print_stats('xyz_cam0', xyz_cam0)

        occ_memI = vox_util.voxelize_xyz(xyz_camI, Z, Y, X, assert_cube=False)
        occI_vis.append(sw.summ_occ('', occ_memI, only_return=True))

        occ_mem0 = vox_util.voxelize_xyz(xyz_cam0, Z, Y, X, assert_cube=False)
        occ0_vis.append(sw.summ_occ('', occ_mem0, only_return=True))


        # print_stats('depths', depths)
        # max_depth = 16
        # depths_valid = (depths < max_depth).float() * (depths > 0.01).float()
        # depths = depths * depths_valid
        # print_stats('depths', depths)

        if sw is not None and sw.save_this:
            sw.summ_rgbs('0_inputs/occI_vis', occI_vis)
            sw.summ_rgbs('0_inputs/occ0_vis', occ0_vis)
            # sw.summ_rgb('0_inputs/normal0', utils.basic.normalize(normals[:,0])-0.5)
            # sw.summ_rgb('0_inputs/normal0', normals[:,0])
            sw.summ_rgbs('0_inputs/normals', normals[:,:4].unbind(1))
            # prep_rgbs = utils.improc.preprocess_color(rgbs)
            # prep_grays = torch.mean(prep_rgbs, dim=2, keepdim=True).repeat(1, 1, 3, 1, 1)

            # sw.summ_traj2ds_on_rgb('0_inputs/trajs_x_on_rgb', trajs_x[0:1], prep_rgbs.mean(dim=1), valids=valids[0:1], cmap='winter')
            # sw.summ_traj2ds_on_rgb('0_inputs/trajs_g_on_rgb', trajs_g[0:1], prep_rgbs.mean(dim=1), valids=valids[0:1], cmap='winter')

            # sw.summ_oned('0_inputs/depth0', depths[:,0], norm=True, frame_id=med.item())
            sw.summ_rgb('0_inputs/rgb0', rgbs[:,0].byte(), frame_id=med.item())
            # sw.summ_oned('0_inputs/depth0', depths[:,0] * depths_valid[:,0], norm=True, frame_id=med.item())
            # sw.summ_oned('0_inputs/depth0', depths[:,0] * depths_valid[:,0], norm=False, max_val=8.0, frame_id=med.item())

            print_stats('depths[:,0]', depths[:,0])

            sw.summ_oned('0_inputs/depth0', depths[:,0] * depths_valid[:,0], max_val=scene_centroid_z+ZMAX, frame_id=med.item())
        
        
    return None 
    

def main(
        exp_name='debug',
        dset='train',
        B=1, # batchsize 
        S=8, # seqlen
        N=256, # number of points per clip
        use_augs=False, # resizing/jittering/color/blur augs
        shuffle=False, # dataset shuffling
        log_dir='./logs_just_vis_3d',
        dataset_location='/orion/group/point_odyssey',
        log_freq=1,
        max_iters=10,
        quick=False,
        dname=None,
):
    device = 'cpu:0'

    # the idea in this file is:
    # load the 3d pointodyssey data and visualize it
    
    exp_name = 'jw00' # copy from dev repo
    exp_name = 'jw01' # clean up
    exp_name = 'jw02' # go
    exp_name = 'jw03' # collect vis*valid first
    exp_name = 'jw04' # show me depth
    exp_name = 'jw05' # print median
    exp_name = 'jw06' # scale extrinsics

    # autogen a descriptive name
    model_name = "%d_%d_%d" % (B, S, N)
    if use_augs:
        model_name += "_A"
    model_name += "_%s" % exp_name
    import datetime
    model_date = datetime.datetime.now().strftime('%H:%M:%S')
    model_name = model_name + '_' + model_date
    print('model_name', model_name)
    
    writer_t = SummaryWriter(log_dir + '/' + model_name + '/t', max_queue=10, flush_secs=60)

    # get dataset
    def worker_init_fn(worker_id):
        np.random.seed(np.random.get_state()[1][0] + worker_id)
    dataset_t = PointOdysseyDataset(
        dataset_location=dataset_location,
        dset=dset,
        S=S,
        N=N,
        quick=quick,
        verbose=True,
    )
    dataloader_t = DataLoader(
        dataset_t,
        batch_size=B,
        shuffle=shuffle,
        num_workers=0,
        worker_init_fn=worker_init_fn,
        drop_last=True)
    iterloader_t = iter(dataloader_t)

    global_step = 0
    while global_step < max_iters:
        global_step += 1
        iter_start_time = time.time()
        sw_t = utils.improc.Summ_writer(
            writer=writer_t,
            global_step=global_step,
            log_freq=log_freq,
            fps=min(S,8),
            scalar_freq=log_freq//5,
            just_gif=True)

        gotit = (False,False)
        while not all(gotit):
            try:
                sample, gotit = next(iterloader_t)
            except StopIteration:
                iterloader_t = iter(dataloader_t)
                sample, gotit = next(iterloader_t)

        iter_rtime = time.time()-iter_start_time
        _ = run_model(sample, device, sw=sw_t)
        iter_itime = time.time()-iter_start_time
        
        print('%s; step %06d/%d; rtime %.2f; itime %.2f' % (
            model_name, global_step, max_iters, iter_rtime, iter_itime))
            
    writer_t.close()
            

if __name__ == '__main__':
    Fire(main)

just_vis_3d.py
aharley/pips2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
import utils.basic
from utils.basic import print_stats
import utils.samp
import utils.misc
from torch import nn, einsum
from einops import rearrange, repeat
from einops.layers.torch import Rearrange, Reduce

class Conv1dPad(nn.Module):
    """
    nn.Conv1d with auto-computed padding ("same" padding)
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, groups=1):
        super(Conv1dPad, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.groups = groups
        self.conv = torch.nn.Conv1d(
            in_channels=self.in_channels, 
            out_channels=self.out_channels, 
            kernel_size=self.kernel_size, 
            stride=self.stride, 
            groups=self.groups)

    def forward(self, x):
        net = x
        in_dim = net.shape[-1]
        out_dim = (in_dim + self.stride - 1) // self.stride
        p = max(0, (out_dim - 1) * self.stride + self.kernel_size - in_dim)
        pad_left = p // 2
        pad_right = p - pad_left
        net = F.pad(net, (pad_left, pad_right), "constant", 0)
        net = self.conv(net)
        return net
    
class ResidualBlock1d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, groups, use_norm, use_do, is_first_block=False):
        super(ResidualBlock1d, self).__init__()
        
        self.in_channels = in_channels
        self.kernel_size = kernel_size
        self.out_channels = out_channels
        self.stride = stride
        self.groups = groups
        self.stride = 1
        self.is_first_block = is_first_block
        self.use_norm = use_norm
        self.use_do = use_do

        self.norm1 = nn.InstanceNorm1d(in_channels)
        self.relu1 = nn.ReLU()
        self.do1 = nn.Dropout(p=0.5)
        self.conv1 = Conv1dPad(
            in_channels=in_channels, 
            out_channels=out_channels, 
            kernel_size=kernel_size, 
            stride=self.stride,
            groups=self.groups)

        self.norm2 = nn.InstanceNorm1d(out_channels)
        self.relu2 = nn.ReLU()
        self.do2 = nn.Dropout(p=0.5)
        self.conv2 = Conv1dPad(
            in_channels=out_channels, 
            out_channels=out_channels, 
            kernel_size=kernel_size, 
            stride=1,
            groups=self.groups)

    def forward(self, x):
        
        identity = x
        
        out = x
        if not self.is_first_block:
            if self.use_norm:
                out = self.norm1(out)
            out = self.relu1(out)
            if self.use_do:
                out = self.do1(out)
        out = self.conv1(out)
        
        if self.use_norm:
            out = self.norm2(out)
        out = self.relu2(out)
        if self.use_do:
            out = self.do2(out)
        out = self.conv2(out)
            
        if self.out_channels != self.in_channels:
            identity = identity.transpose(-1,-2)
            ch1 = (self.out_channels-self.in_channels)//2
            ch2 = self.out_channels-self.in_channels-ch1
            identity = F.pad(identity, (ch1, ch2), "constant", 0)
            identity = identity.transpose(-1,-2)
        
        out += identity
        return out

def sequence_loss(flow_preds, flow_gt, vis, valids, gamma=0.8):
    """ Loss function defined over sequence of flow predictions """
    B, S, N, D = flow_gt.shape
    assert(D==2)
    B, S1, N = vis.shape
    B, S2, N = valids.shape
    assert(S==S1)
    assert(S==S2)
    n_predictions = len(flow_preds)    
    flow_loss = 0.0
    for i in range(n_predictions):
        i_weight = gamma**(n_predictions - i - 1)
        flow_pred = flow_preds[i]
        i_loss = (flow_pred - flow_gt).abs() # B,S,N,2
        i_loss = torch.mean(i_loss, dim=3) # B,S,N
        flow_loss += i_weight * utils.basic.reduce_masked_mean(i_loss, valids)
    flow_loss = flow_loss/n_predictions
    return flow_loss

class ResidualBlock2d(nn.Module):
    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
        super(ResidualBlock2d, 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.InstanceNorm2d(planes)
            self.norm2 = nn.InstanceNorm2d(planes)
            if not stride == 1:
                self.norm3 = nn.InstanceNorm2d(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=8, norm_fn='batch', dropout=0.0):
        super(BasicEncoder, self).__init__()
        self.stride = stride
        self.norm_fn = norm_fn

        self.in_planes = 64
        
        if self.norm_fn == 'group':
            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=self.in_planes)
            self.norm2 = nn.GroupNorm(num_groups=8, num_channels=output_dim*2)
            
        elif self.norm_fn == 'batch':
            self.norm1 = nn.InstanceNorm2d(self.in_planes)
            self.norm2 = nn.InstanceNorm2d(output_dim*2)

        elif self.norm_fn == 'instance':
            self.norm1 = nn.InstanceNorm2d(self.in_planes)
            self.norm2 = nn.InstanceNorm2d(output_dim*2)

        elif self.norm_fn == 'none':
            self.norm1 = nn.Sequential()
            
        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(64,  stride=1)
        self.layer2 = self._make_layer(96, stride=2)
        self.layer3 = self._make_layer(128, stride=2)
        self.layer4 = self._make_layer(128, stride=2)

        self.conv2 = nn.Conv2d(128+128+96+64, 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)
        
        self.dropout = None
        if dropout > 0:
            self.dropout = nn.Dropout2d(p=dropout)

        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, nn.InstanceNorm2d, nn.GroupNorm)):
                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 = ResidualBlock2d(self.in_planes, dim, self.norm_fn, stride=stride)
        layer2 = ResidualBlock2d(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
        
        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)
        a = F.interpolate(a, (H//self.stride, W//self.stride), mode='bilinear', align_corners=True)
        b = F.interpolate(b, (H//self.stride, W//self.stride), mode='bilinear', align_corners=True)
        c = F.interpolate(c, (H//self.stride, W//self.stride), mode='bilinear', align_corners=True)
        d = F.interpolate(d, (H//self.stride, W//self.stride), mode='bilinear', align_corners=True)
        x = self.conv2(torch.cat([a,b,c,d], dim=1))
        x = self.norm2(x)
        x = self.relu2(x)
        x = self.conv3(x)

        if self.training and self.dropout is not None:
            x = self.dropout(x)
        return x

class DeltaBlock(nn.Module):
    def __init__(self, latent_dim=128, hidden_dim=128, corr_levels=4, corr_radius=3):
        super(DeltaBlock, self).__init__()
        
        kitchen_dim = (corr_levels * (2*corr_radius + 1)**2)*3 + latent_dim + 2

        self.latent_dim = latent_dim
        self.hidden_dim = hidden_dim
        
        in_channels = kitchen_dim
        base_filters = 128
        self.n_block = 8
        self.kernel_size = 3
        self.groups = 1
        self.use_norm = True
        self.use_do = False

        self.increasefilter_gap = 2 

        self.first_block_conv = Conv1dPad(in_channels=in_channels, out_channels=base_filters, kernel_size=self.kernel_size, stride=1)
        self.first_block_norm = nn.InstanceNorm1d(base_filters)
        self.first_block_relu = nn.ReLU()
        out_channels = base_filters
                
        self.basicblock_list = nn.ModuleList()
        for i_block in range(self.n_block):

            if i_block == 0:
                is_first_block = True
            else:
                is_first_block = False

            if is_first_block:
                in_channels = base_filters
                out_channels = in_channels
            else:
                in_channels = int(base_filters*2**((i_block-1)//self.increasefilter_gap))
                if (i_block % self.increasefilter_gap == 0) and (i_block != 0):
                    out_channels = in_channels * 2
                else:
                    out_channels = in_channels
            
            tmp_block = ResidualBlock1d(
                in_channels=in_channels, 
                out_channels=out_channels, 
                kernel_size=self.kernel_size, 
                stride=1, 
                groups=self.groups, 
                use_norm=self.use_norm, 
                use_do=self.use_do, 
                is_first_block=is_first_block)
            self.basicblock_list.append(tmp_block)

        self.final_norm = nn.InstanceNorm1d(out_channels)
        self.final_relu = nn.ReLU(inplace=True)
        self.dense = nn.Linear(out_channels, 2)
        
            
    def forward(self, fcorr, flow):
        B, S, D = flow.shape
        assert(D==2)
        flow_sincos = utils.misc.posemb_sincos_2d_xy(flow, self.latent_dim, cat_coords=True)
        x = torch.cat([fcorr, flow_sincos], dim=2) # B,S,-1
        
        # conv1d wants channels in the middle
        out = x.permute(0,2,1)
        out = self.first_block_conv(out)
        out = self.first_block_relu(out)
        for i_block in range(self.n_block):
            net = self.basicblock_list[i_block]
            out = net(out)
        out = self.final_relu(out)
        out = out.permute(0,2,1)
        
        delta = self.dense(out)
        return delta

def bilinear_sampler(img, coords, mode='bilinear', mask=False):
    """ Wrapper for grid_sample, uses pixel coordinates """
    H, W = img.shape[-2:]
    xgrid, ygrid = coords.split([1,1], dim=-1)
    # go to 0,1 then 0,2 then -1,1
    xgrid = 2*xgrid/(W-1) - 1
    ygrid = 2*ygrid/(H-1) - 1

    grid = torch.cat([xgrid, ygrid], dim=-1)
    img = F.grid_sample(img, grid, align_corners=True)

    if mask:
        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
        return img, mask.float()
    return img

def coords_grid(batch, ht, wd):
    coords = torch.meshgrid(torch.arange(ht), torch.arange(wd), indexing='ij')
    coords = torch.stack(coords[::-1], dim=0).float()
    return coords[None].repeat(batch, 1, 1, 1)

class CorrBlock:
    def __init__(self, fmaps, num_levels=4, radius=4):
        B, S, C, H, W = fmaps.shape
        self.S, self.C, self.H, self.W = S, C, H, W
        self.num_levels = num_levels
        self.radius = radius
        self.fmaps_pyramid = []
        self.fmaps_pyramid.append(fmaps)
        for i in range(self.num_levels-1):
            fmaps_ = fmaps.reshape(B*S, C, H, W)
            fmaps_ = F.avg_pool2d(fmaps_, 2, stride=2)
            _, _, H, W = fmaps_.shape
            fmaps = fmaps_.reshape(B, S, C, H, W)
            self.fmaps_pyramid.append(fmaps)

    def sample(self, coords):
        r = self.radius
        B, S, N, D = coords.shape
        assert(D==2)

        x0 = coords[:,0,:,0].round().clamp(0, self.W-1).long()
        y0 = coords[:,0,:,1].round().clamp(0, self.H-1).long()

        H, W = self.H, self.W
        out_pyramid = []
        for i in range(self.num_levels):
            corrs = self.corrs_pyramid[i] # B,S,N,H,W
            _, _, _, H, W = corrs.shape
            
            dx = torch.linspace(-r, r, 2*r+1)
            dy = torch.linspace(-r, r, 2*r+1)
            delta = torch.stack(torch.meshgrid(dy, dx, indexing='ij'), axis=-1).to(coords.device) 

            centroid_lvl = coords.reshape(B*S*N, 1, 1, 2) / 2**i
            delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
            coords_lvl = centroid_lvl + delta_lvl

            corrs = bilinear_sampler(corrs.reshape(B*S*N, 1, H, W), coords_lvl)
            corrs = corrs.view(B, S, N, -1)
            out_pyramid.append(corrs)

        out = torch.cat(out_pyramid, dim=-1) # B,S,N,LRR*2
        return out.contiguous().float()

    def corr(self, targets):
        B, S, N, C = targets.shape
        assert(C==self.C)
        self.corrs_pyramid = []
        for fmaps in self.fmaps_pyramid:
            _, _, _, H, W = fmaps.shape
            fmap2s = fmaps.view(B, S, C, H*W)
            corrs = torch.matmul(targets, fmap2s)
            corrs = corrs.view(B, S, N, H, W) 
            corrs = corrs / torch.sqrt(torch.tensor(C).float())
            self.corrs_pyramid.append(corrs)

class Pips(nn.Module):
    def __init__(self, stride=8):
        super(Pips, self).__init__()

        self.stride = stride

        self.hidden_dim = hdim = 256
        self.latent_dim = latent_dim = 128
        self.corr_levels = 4
        self.corr_radius = 3
        
        self.fnet = BasicEncoder(output_dim=self.latent_dim, norm_fn='instance', dropout=0, stride=stride)
        self.delta_block = DeltaBlock(hidden_dim=self.hidden_dim, corr_levels=self.corr_levels, corr_radius=self.corr_radius)
        self.norm = nn.GroupNorm(1, self.latent_dim)

    def forward(self, trajs_e0, rgbs, iters=3, trajs_g=None, vis_g=None, valids=None, sw=None, feat_init=None, is_train=False, beautify=False):
        total_loss = torch.tensor(0.0).cuda()

        B,S,N,D = trajs_e0.shape
        assert(D==2)

        B,S,C,H,W = rgbs.shape
        rgbs = 2 * (rgbs / 255.0) - 1.0
        
        H8 = H//self.stride
        W8 = W//self.stride

        device = rgbs.device

        rgbs_ = rgbs.reshape(B*S, C, H, W)
        fmaps_ = self.fnet(rgbs_)
        fmaps = fmaps_.reshape(B, S, self.latent_dim, H8, W8)

        if sw is not None and sw.save_this:
            sw.summ_feats('1_model/0_fmaps', fmaps.unbind(1))

        coords = trajs_e0.clone()/float(self.stride)

        hdim = self.hidden_dim

        fcorr_fn1 = CorrBlock(fmaps, num_levels=self.corr_levels, radius=self.corr_radius)
        fcorr_fn2 = CorrBlock(fmaps, num_levels=self.corr_levels, radius=self.corr_radius)
        fcorr_fn4 = CorrBlock(fmaps, num_levels=self.corr_levels, radius=self.corr_radius)
        if feat_init is not None:
            feats1, feats2, feats4 = feat_init
        else:
            feat1 = utils.samp.bilinear_sample2d(fmaps[:,0], coords[:,0,:,0], coords[:,0,:,1]).permute(0, 2, 1) # B,N,C
            feats1 = feat1.unsqueeze(1).repeat(1, S, 1, 1) # B,S,N,C
            feats2 = feat1.unsqueeze(1).repeat(1, S, 1, 1) # B,S,N,C
            feats4 = feat1.unsqueeze(1).repeat(1, S, 1, 1) # B,S,N,C
        
        coords_bak = coords.clone()

        coords[:,0] = coords_bak[:,0] # lock coord0 for target
        
        coord_predictions1 = [] # for loss
        coord_predictions2 = [] # for vis

        fcorr_fn1.corr(feats1) # we only need to run this corr once

        coord_predictions2.append(coords.detach() * self.stride)
        
        for itr in range(iters):
            coords = coords.detach()

            if itr >= 1:
                # timestep indices
                inds2 = (torch.arange(S)-2).clip(min=0)
                inds4 = (torch.arange(S)-4).clip(min=0)
                # coordinates at these timesteps
                coords2_ = coords[:,inds2].reshape(B*S,N,2)
                coords4_ = coords[:,inds4].reshape(B*S,N,2)
                # featuremaps at these timesteps
                fmaps2_ = fmaps[:,inds2].reshape(B*S,self.latent_dim,H8,W8)
                fmaps4_ = fmaps[:,inds4].reshape(B*S,self.latent_dim,H8,W8)
                # features at these coords/times
                feats2_ = utils.samp.bilinear_sample2d(fmaps2_, coords2_[:,:,0], coords2_[:,:,1]).permute(0, 2, 1) # B*S, N, C
                feats2 = feats2_.reshape(B,S,N,self.latent_dim)
                feats4_ = utils.samp.bilinear_sample2d(fmaps4_, coords4_[:,:,0], coords4_[:,:,1]).permute(0, 2, 1) # B*S, N, C
                feats4 = feats4_.reshape(B,S,N,self.latent_dim)

            fcorr_fn2.corr(feats2)
            fcorr_fn4.corr(feats4)

            # now we want costs at the current locations
            fcorrs1 = fcorr_fn1.sample(coords) # B,S,N,LRR
            fcorrs2 = fcorr_fn2.sample(coords) # B,S,N,LRR
            fcorrs4 = fcorr_fn4.sample(coords) # B,S,N,LRR
            LRR = fcorrs1.shape[3]

            # we want everything in the format B*N, S, C
            fcorrs1_ = fcorrs1.permute(0, 2, 1, 3).reshape(B*N, S, LRR)
            fcorrs2_ = fcorrs2.permute(0, 2, 1, 3).reshape(B*N, S, LRR)
            fcorrs4_ = fcorrs4.permute(0, 2, 1, 3).reshape(B*N, S, LRR)
            fcorrs_ = torch.cat([fcorrs1_, fcorrs2_, fcorrs4_], dim=2)
            flows_ = (coords[:,1:] - coords[:,:-1]).permute(0,2,1,3).reshape(B*N, S-1, 2)
            flows_ = torch.cat([flows_, flows_[:,-1:]], dim=1) # B*N,S,2

            delta_coords_ = self.delta_block(fcorrs_, flows_) # B*N,S,2

            if beautify and itr > 3*iters//4:
                # this smooths the results a bit, but does not really help perf
                delta_coords_ = delta_coords_ * 0.5

            coords = coords + delta_coords_.reshape(B, N, S, 2).permute(0,2,1,3)

            coord_predictions1.append(coords * self.stride)

            coords[:,0] = coords_bak[:,0] # lock coord0 for target
            coord_predictions2.append(coords * self.stride)
            
        # pause at the end, to make the summs more interpretable
        coord_predictions2.append(coords * self.stride)

        if trajs_g is not None:
            loss = sequence_loss(coord_predictions1, trajs_g, vis_g, valids, 0.8)
        else:
            loss = None
            
        coord_predictions1.append(coords * self.stride)
        feats = (feats1, feats2, feats4)
        return coord_predictions1, coord_predictions2, feats, loss

nets/pips2.py
aharley/pips2
import torch
import os, pathlib
import numpy as np

def save(ckpt_dir, optimizer, model, global_step, scheduler=None, keep_latest=5, model_name='model'):
    if not os.path.exists(ckpt_dir):
        os.makedirs(ckpt_dir)

    prev_ckpts = list(pathlib.Path(ckpt_dir).glob('%s-*' % 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()
    model_path = '%s/%s-%09d.pth' % (ckpt_dir, model_name, global_step)
    
    ckpt = {'optimizer_state_dict': optimizer.state_dict()}
    ckpt['model_state_dict'] = model.state_dict()
    if scheduler is not None:
        ckpt['scheduler_state_dict'] = scheduler.state_dict()
    torch.save(ckpt, model_path)
    print("saved a checkpoint: %s" % (model_path))

def load(ckpt_dir, model, optimizer=None, scheduler=None, step=0, model_name='model', ignore_load=None):
    print('reading ckpt from %s' % ckpt_dir)
    if not os.path.exists(ckpt_dir):
        print('...there is no full checkpoint here!')
        print('-- note this function no longer appends "saved_checkpoints/" before the ckpt_dir --')
    else:
        ckpt_names = os.listdir(ckpt_dir)
        steps = [int((i.split('-')[1]).split('.')[0]) for i in ckpt_names]
        if len(ckpt_names) > 0:
            if step==0:
                step = max(steps)
            model_name = '%s-%09d.pth' % (model_name, step)
            path = os.path.join(ckpt_dir, model_name)
            print('...found checkpoint %s'%(path))

            if ignore_load is not None:
                
                print('ignoring', ignore_load)

                checkpoint = torch.load(path)['model_state_dict']

                model_dict = model.state_dict()

                # 1. filter out ignored keys
                pretrained_dict = {k: v for k, v in checkpoint.items()}
                for ign in ignore_load:
                    pretrained_dict = {k: v for k, v in pretrained_dict.items() if not ign in k}
                    
                # 2. overwrite entries in the existing state dict
                model_dict.update(pretrained_dict)
                # 3. load the new state dict
                model.load_state_dict(model_dict, strict=False)
            else:
                checkpoint = torch.load(path)
                model.load_state_dict(checkpoint['model_state_dict'], strict=False)
                
            if optimizer is not None:
                optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            if scheduler is not None:
                scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        else:
            print('...there is no full checkpoint here!')
    return step

saverloader.py
aharley/pips2
import time
import numpy as np
import saverloader
from nets.pips2 import Pips
import utils.improc
import utils.misc
import random, os, cv2
from utils.basic import print_, print_stats
import torch
from tensorboardX import SummaryWriter
import torch.nn.functional as F
from fire import Fire
from torch.utils.data import Dataset, DataLoader
from datasets.crohddataset import CrohdDataset

def create_pools(n_pool=1000):
    pools = {}
    pool_names = [
        'd_1',
        'd_2',
        'd_4',
        'd_8',
        'd_16',
        'd_avg',
        'median_l2',
        'survival',
        'ate_all',
        'ate_vis',
        'ate_occ',
        'total_loss',
    ]
    for pool_name in pool_names:
        pools[pool_name] = utils.misc.SimplePool(n_pool, version='np')
    return pools

def test_on_fullseq(model, d, sw, iters=8, S_max=8, image_size=(384,512)):
    metrics = {}

    folder = str(d['folder'][0])
    print('folder', folder)
    trajs_g = d['trajs'].cuda().float() # B,S,N,2
    visibs = d['visibs'].cuda().float() # B,S,N
    valids = d['valids'].cuda().float() # B,S,N
    start_frame = int(d['start_frame'][0])

    B, S, N, D = trajs_g.shape
    assert(D==2)
    assert(B==1)
    print('this video is %d frames long' % S)

    # collect the seq of rgb paths
    rgb_paths = []
    for ii in range(start_frame, start_frame+S):
        rgb_path = os.path.join(folder, 'img1', str(ii+1).zfill(6)+'.jpg')
        rgb_paths.append(rgb_path)

    # load one to check H,W
    rgb0_bak = cv2.imread(rgb_paths[0])
    H_bak, W_bak = rgb0_bak.shape[:2]
    H, W = image_size
    sy = H/H_bak
    sx = W/W_bak
    trajs_g[:,:,:,0] *= sx
    trajs_g[:,:,:,1] *= sy
    rgb0_bak = cv2.resize(rgb0_bak, (W, H), interpolation=cv2.INTER_LINEAR)
    rgb0_bak = torch.from_numpy(rgb0_bak[:,:,::-1].copy()).permute(2,0,1) # 3,H,W
    rgb0_bak = rgb0_bak.unsqueeze(0).to(trajs_g.device) # 1,3,H,W

    if sw is not None and sw.save_this:
        prep_rgb0 = utils.improc.preprocess_color(rgb0_bak)
        sw.summ_traj2ds_on_rgb('0_inputs/trajs_g_on_rgb0', trajs_g[0:1], prep_rgb0, valids=valids[0:1], cmap='winter', linewidth=1)
    
    # zero-vel init
    trajs_e = trajs_g[:,0].repeat(1,S,1,1)
    
    cur_frame = 0
    done = False
    feat_init = None
    while not done:
        end_frame = cur_frame + S_max

        if end_frame > S:
            diff = end_frame-S
            end_frame = end_frame-diff
            cur_frame = max(cur_frame-diff,0)
        print('working on subseq %d:%d' % (cur_frame, end_frame))

        traj_seq = trajs_e[:, cur_frame:end_frame]

        idx_seq = np.arange(cur_frame, end_frame)
        rgb_paths_seq = [rgb_paths[ii] for ii in idx_seq]
        rgbs = [cv2.imread(rgb_path) for rgb_path in rgb_paths_seq]
        rgbs = [rgb[:,:,::-1] for rgb in rgbs] # BGR->RGB
        H_load, W_load = rgbs[0].shape[:2]
        assert(H_load==H_bak and W_load==W_bak)
        rgbs = [cv2.resize(rgb, (W, H), interpolation=cv2.INTER_LINEAR) for rgb in rgbs]
        rgb_seq = torch.from_numpy(np.stack(rgbs, 0)).permute(0,3,1,2) # S,3,H,W
        rgb_seq = rgb_seq.unsqueeze(0).to(traj_seq.device) # 1,S,3,H,W
        S_local = rgb_seq.shape[1]

        if feat_init is not None:
            feat_init = [fi[:,:S_local] for fi in feat_init]
            
        preds, preds_anim, feat_init, _ = model(traj_seq, rgb_seq, iters=iters, feat_init=feat_init)

        trajs_e[:, cur_frame:end_frame] = preds[-1][:, :S_local]
        trajs_e[:, end_frame:] = trajs_e[:, end_frame-1:end_frame] # update the future with new zero-vel
        
        # if sw is not None and sw.save_this:
        #     traj_seq_e = preds[-1]
        #     traj_seq_g = trajs_g[:,cur_frame:end_frame]
        #     valid_seq = valids[:,cur_frame:end_frame]
        #     prep_rgbs = utils.improc.preprocess_color(rgb_seq)
        #     gray_rgbs = torch.mean(prep_rgbs, dim=2, keepdim=True).repeat(1, 1, 3, 1, 1)
        #     gt_rgb = utils.improc.preprocess_color(sw.summ_traj2ds_on_rgb('', traj_seq_g, gray_rgbs[0:1].mean(dim=1), valids=valid_seq, cmap='winter', only_return=True))
        #     rgb_vis = []
        #     for tre in preds_anim:
        #         ate = torch.norm(tre - traj_seq_g, dim=-1) # B,S,N
        #         ate_all = utils.basic.reduce_masked_mean(ate, valid_seq, dim=[1,2]) # B
        #         rgb_vis.append(sw.summ_traj2ds_on_rgb('', tre[0:1], gt_rgb, valids=valid_seq, only_return=True, cmap='spring', frame_id=ate_all[0]))
        #     sw.summ_rgbs('3_test/animated_trajs_on_rgb_cur%02d' % cur_frame, rgb_vis)

        if end_frame >= S:
            done = True
        else:
            cur_frame = cur_frame + S_max - 1

    d_sum = 0.0
    thrs = [1,2,4,8,16]
    sx_ = W / 256.0
    sy_ = H / 256.0
    sc_py = np.array([sx_, sy_]).reshape([1,1,2])
    sc_pt = torch.from_numpy(sc_py).float().cuda()
    for thr in thrs:
        # note we exclude timestep0 from this eval
        d_ = (torch.norm(trajs_e[:,1:]/sc_pt - trajs_g[:,1:]/sc_pt, dim=-1) < thr).float() # B,S-1,N
        d_ = utils.basic.reduce_masked_mean(d_, valids[:,1:]).item()*100.0
        d_sum += d_
        metrics['d_%d' % thr] = d_
    d_avg = d_sum / len(thrs)
    metrics['d_avg'] = d_avg

    sur_thr = 16
    dists = torch.norm(trajs_e/sc_pt - trajs_g/sc_pt, dim=-1) # B,S,N
    dist_ok = 1 - (dists > sur_thr).float() * valids # B,S,N
    survival = torch.cumprod(dist_ok, dim=1) # B,S,N
    metrics['survival'] = torch.mean(survival).item()*100.0
    
    # get the median l2 error for each trajectory
    dists_ = dists.permute(0,2,1).reshape(B*N,S)
    valids_ = valids.permute(0,2,1).reshape(B*N,S)
    median_l2 = utils.basic.reduce_masked_median(dists_, valids_, keep_batch=True)
    metrics['median_l2'] = median_l2.mean().item()

    if sw is not None and sw.save_this:
        rgb0_g = sw.summ_traj2ds_on_rgb('', trajs_g[0:1], prep_rgb0, valids=valids[0:1], cmap='winter', linewidth=2, only_return=True)
        sw.summ_traj2ds_on_rgb('2_outputs/trajs_e_on_rgb0', trajs_e[0:1], utils.improc.preprocess_color(rgb0_g), valids=valids[0:1], cmap='spring', linewidth=2, frame_id=d_avg)
        
    return metrics

    
def main(
        B=1, # batchsize 
        S=128, # seqlen
        N=256, # number of points per clip
        stride=8, # spatial stride of the model
        iters=16, # inference steps of the model
        image_size=(512,896), # input resolution
        shuffle=False, # dataset shuffling
        log_freq=99, # how often to make image summaries
        max_iters=99, # how many samples to test
        log_dir='./logs_test_on_cro',
        dataset_location='/orion/u/aharley/datasets/head_tracking',
        init_dir='./reference_model',
        device_ids=[0],
        n_pool=1000, # how long the running averages should be
):
    device = 'cuda:%d' % device_ids[0]

    # the idea in this file is:
    # load a ckpt, and test it in crohd,
    # tracking points from frame0 to the end.

    exp_name = 'cro00' # copy from dev repo
    exp_name = 'cro01' # clean up
    exp_name = 'cro02' # survival thr 16
    

    assert(B==1) # B>1 not implemented here
    assert(image_size[0] % 32 == 0)
    assert(image_size[1] % 32 == 0)

    # autogen a descriptive name
    model_name = "%d_%d" % (B,S)
    model_name += "_i%d" % (iters)
    model_name += "_%s" % exp_name
    import datetime
    model_date = datetime.datetime.now().strftime('%H%M%S')
    model_name = model_name + '_' + model_date
    print('model_name', model_name)

    writer_x = SummaryWriter(log_dir + '/' + model_name + '/x', max_queue=10, flush_secs=60)
        
    dataset_x = CrohdDataset(
        dataset_location=dataset_location,
        S=1000, # create clips of this length
    )
    dataloader_x = DataLoader(
        dataset_x,
        batch_size=B,
        shuffle=shuffle,
        num_workers=0,
        drop_last=True)
    iterloader_x = iter(dataloader_x)

    model = Pips(stride=stride).to(device)
    model = torch.nn.DataParallel(model, device_ids=device_ids)

    utils.misc.count_parameters(model)

    _ = saverloader.load(init_dir, model.module)
    model.eval()

    pools_x = create_pools(n_pool)
    
    global_step = 0
    max_iters = min(max_iters, len(dataset_x))
    while global_step < max_iters:
        global_step += 1
        iter_start_time = time.time()
        with torch.no_grad():
            torch.cuda.empty_cache()
        sw_x = utils.improc.Summ_writer(
            writer=writer_x,
            global_step=global_step,
            log_freq=log_freq,
            fps=min(S,8),
            scalar_freq=1,
            just_gif=True)
        try:
            sample = next(iterloader_x)
        except StopIteration:
            iterloader_x = iter(dataloader_x)
            sample = next(iterloader_x)
        iter_rtime = time.time()-iter_start_time
        with torch.no_grad():
            metrics = test_on_fullseq(model, sample, sw_x, iters=iters, S_max=S, image_size=image_size)
        for key in list(pools_x.keys()):
            if key in metrics:
                pools_x[key].update([metrics[key]])
                sw_x.summ_scalar('_/%s' % (key), pools_x[key].mean())
        iter_itime = time.time()-iter_start_time

        print('%s; step %06d/%d; rtime %.2f; itime %.2f; d_x %.1f; sur_x %.1f; med_x %.1f' % (
            model_name, global_step, max_iters, iter_rtime, iter_itime,
            pools_x['d_avg'].mean(), pools_x['survival'].mean(), pools_x['median_l2'].mean()))
            
    writer_x.close()
            

if __name__ == '__main__':
    Fire(main)

test_on_cro.py
aharley/pips2
import time
import numpy as np
import saverloader
from nets.pips2 import Pips
import utils.improc
import utils.misc
import random, os, cv2
from utils.basic import print_, print_stats
import torch
from tensorboardX import SummaryWriter
import torch.nn.functional as F
from fire import Fire
from torch.utils.data import Dataset, DataLoader
from datasets.pointodysseydataset_fullseq import PointOdysseyDataset

def create_pools(n_pool=1000):
    pools = {}
    pool_names = [
        'd_1',
        'd_2',
        'd_4',
        'd_8',
        'd_16',
        'd_avg',
        'median_l2',
        'survival',
        'ate_all',
        'ate_vis',
        'ate_occ',
        'total_loss',
    ]
    for pool_name in pool_names:
        pools[pool_name] = utils.misc.SimplePool(n_pool, version='np')
    return pools

def test_on_fullseq(model, d, sw, iters=8, S_max=8, image_size=(384,512)):
    metrics = {}

    seq = str(d['seq'][0])
    print('seq', seq)
    trajs_g = d['trajs'].cuda().float() # B,S,N,2
    visibs = d['visibs'].cuda().float() # B,S,N
    valids = d['valids'].cuda().float() # B,S,N

    B, S, N, D = trajs_g.shape
    assert(D==2)
    assert(B==1)
    print('this video is %d frames long' % S)

    # load one to check H,W
    rgb_path0 = os.path.join(seq, 'rgbs', 'rgb_%05d.jpg' % (0))
    rgb0_bak = cv2.imread(rgb_path0)
    H_bak, W_bak = rgb0_bak.shape[:2]
    H, W = image_size
    sy = H/H_bak
    sx = W/W_bak
    trajs_g[:,:,:,0] *= sx
    trajs_g[:,:,:,1] *= sy
    rgb0_bak = cv2.resize(rgb0_bak, (W, H), interpolation=cv2.INTER_LINEAR)
    rgb0_bak = torch.from_numpy(rgb0_bak[:,:,::-1].copy()).permute(2,0,1) # 3,H,W
    rgb0_bak = rgb0_bak.unsqueeze(0).to(trajs_g.device) # 1,3,H,W

    if sw is not None and sw.save_this:
        prep_rgb0 = utils.improc.preprocess_color(rgb0_bak)
        sw.summ_traj2ds_on_rgb('0_inputs/trajs_g_on_rgb0', trajs_g[0:1], prep_rgb0, valids=valids[0:1], cmap='winter', linewidth=1)
        
    # zero-vel init
    trajs_e = trajs_g[:,0].repeat(1,S,1,1)
    
    cur_frame = 0
    done = False
    feat_init = None
    while not done:
        end_frame = cur_frame + S_max

        if end_frame > S:
            diff = end_frame-S
            end_frame = end_frame-diff
            cur_frame = max(cur_frame-diff,0)
        print('working on subseq %d:%d' % (cur_frame, end_frame))

        traj_seq = trajs_e[:, cur_frame:end_frame]

        idx_seq = np.arange(cur_frame, end_frame)
        rgb_paths_seq = [os.path.join(seq, 'rgbs', 'rgb_%05d.jpg' % (idx)) for idx in idx_seq]
        rgbs = [cv2.imread(rgb_path) for rgb_path in rgb_paths_seq]
        rgbs = [rgb[:,:,::-1] for rgb in rgbs] # BGR->RGB
        H_load, W_load = rgbs[0].shape[:2]
        assert(H_load==H_bak and W_load==W_bak)
        rgbs = [cv2.resize(rgb, (W, H), interpolation=cv2.INTER_LINEAR) for rgb in rgbs]
        rgb_seq = torch.from_numpy(np.stack(rgbs, 0)).permute(0,3,1,2) # S,3,H,W
        rgb_seq = rgb_seq.unsqueeze(0).to(traj_seq.device) # 1,S,3,H,W
        S_local = rgb_seq.shape[1]

        if feat_init is not None:
            feat_init = [fi[:,:S_local] for fi in feat_init]
        
        preds, preds_anim, feat_init, _ = model(traj_seq, rgb_seq, iters=iters, feat_init=feat_init)

        trajs_e[:, cur_frame:end_frame] = preds[-1][:, :S_local]
        trajs_e[:, end_frame:] = trajs_e[:, end_frame-1:end_frame] # update the future with new zero-vel

        # if sw is not None and sw.save_this:
        #     traj_seq_e = preds[-1]
        #     traj_seq_g = trajs_g[:,cur_frame:end_frame]
        #     valid_seq = valids[:,cur_frame:end_frame]
        #     prep_rgbs = utils.improc.preprocess_color(rgb_seq)
        #     gray_rgbs = torch.mean(prep_rgbs, dim=2, keepdim=True).repeat(1, 1, 3, 1, 1)
        #     gt_rgb = utils.improc.preprocess_color(sw.summ_traj2ds_on_rgb('', traj_seq_g, gray_rgbs[0:1].mean(dim=1), valids=valid_seq, cmap='winter', only_return=True))
        #     rgb_vis = []
        #     for tre in preds_anim:
        #         ate = torch.norm(tre - traj_seq_g, dim=-1) # B,S,N
        #         ate_all = utils.basic.reduce_masked_mean(ate, valid_seq, dim=[1,2]) # B
        #         rgb_vis.append(sw.summ_traj2ds_on_rgb('', tre[0:1], gt_rgb, valids=valid_seq, only_return=True, cmap='spring', frame_id=ate_all[0]))
        #     sw.summ_rgbs('3_test/animated_trajs_on_rgb_cur%02d' % cur_frame, rgb_vis)
        
        if end_frame >= S:
            done = True
        else:
            cur_frame = cur_frame + S_max - 1

    d_sum = 0.0
    thrs = [1,2,4,8,16]
    sx_ = W / 256.0
    sy_ = H / 256.0
    sc_py = np.array([sx_, sy_]).reshape([1,1,2])
    sc_pt = torch.from_numpy(sc_py).float().cuda()
    for thr in thrs:
        # note we exclude timestep0 from this eval
        d_ = (torch.norm(trajs_e[:,1:]/sc_pt - trajs_g[:,1:]/sc_pt, dim=-1) < thr).float() # B,S-1,N
        d_ = utils.basic.reduce_masked_mean(d_, valids[:,1:]).item()*100.0
        d_sum += d_
        metrics['d_%d' % thr] = d_
    d_avg = d_sum / len(thrs)
    metrics['d_avg'] = d_avg

    sur_thr = 16
    dists = torch.norm(trajs_e/sc_pt - trajs_g/sc_pt, dim=-1) # B,S,N
    dist_ok = 1 - (dists > sur_thr).float() * valids # B,S,N
    survival = torch.cumprod(dist_ok, dim=1) # B,S,N
    metrics['survival'] = torch.mean(survival).item()*100.0
    
    # get the median l2 error for each trajectory
    dists_ = dists.permute(0,2,1).reshape(B*N,S)
    valids_ = valids.permute(0,2,1).reshape(B*N,S)
    median_l2 = utils.basic.reduce_masked_median(dists_, valids_, keep_batch=True)
    metrics['median_l2'] = median_l2.mean().item()

    if sw is not None and sw.save_this:
        rgb0 = sw.summ_traj2ds_on_rgb('', trajs_g[0:1], prep_rgb0, valids=valids[0:1], cmap='winter', linewidth=2, only_return=True)
        sw.summ_traj2ds_on_rgb('2_outputs/trajs_e_on_rgb0', trajs_e[0:1], utils.improc.preprocess_color(rgb0), valids=valids[0:1], cmap='spring', linewidth=2, frame_id=d_avg)
        
    return metrics
    
def main(
        dset='test', 
        B=1, # batchsize 
        S=128, # seqlen
        N=256, # number of points per clip
        stride=8, # spatial stride of the model
        iters=16, # inference steps of the model
        image_size=(512,896), # input resolution
        shuffle=False, # dataset shuffling
        log_freq=99, # how often to make image summaries
        max_iters=99, # how many samples to test
        log_dir='./logs_test_on_pod',
        dataset_location='/orion/group/point_odyssey',
        init_dir='./reference_model',
        device_ids=[0],
        n_pool=1000, # how long the running averages should be
):
    device = 'cuda:%d' % device_ids[0]

    # the idea in this file is:
    # load a ckpt, and test it in pointodyssey,
    # tracking points from frame0 to the end.
    
    exp_name = 'pod00' # copy from dev repo
    exp_name = 'pod01' # pod test
    exp_name = 'pod02' # clean up
    exp_name = 'pod03' # survival thr 16
    exp_name = 'pod04' # refmodel repeat
    exp_name = 'pod05' # fix bug with d_ stats

    assert(B==1) # B>1 not implemented here
    assert(image_size[0] % 32 == 0)
    assert(image_size[1] % 32 == 0)
    
    # autogen a descriptive name
    model_name = "%d_%d" % (B,S)
    model_name += "_i%d" % (iters)
    model_name += "_%s" % exp_name
    import datetime
    model_date = datetime.datetime.now().strftime('%H%M%S')
    model_name = model_name + '_' + model_date
    print('model_name', model_name)
    
    writer_x = SummaryWriter(log_dir + '/' + model_name + '/x', max_queue=10, flush_secs=60)

    dataset_x = PointOdysseyDataset(
        dataset_location=dataset_location,
        dset=dset,
        N=N,
        verbose=True,
    )
    dataloader_x = DataLoader(
        dataset_x,
        batch_size=B,
        shuffle=shuffle,
        num_workers=0,
        drop_last=True)
    iterloader_x = iter(dataloader_x)

    model = Pips(stride=stride).to(device)
    model = torch.nn.DataParallel(model, device_ids=device_ids)

    utils.misc.count_parameters(model)

    _ = saverloader.load(init_dir, model.module)
    model.eval()

    pools_x = create_pools(n_pool)
    
    global_step = 0
    max_iters = min(max_iters, len(dataset_x))
    while global_step < max_iters:
        global_step += 1
        iter_start_time = time.time()
        with torch.no_grad():
            torch.cuda.empty_cache()
        sw_x = utils.improc.Summ_writer(
            writer=writer_x,
            global_step=global_step,
            log_freq=log_freq,
            fps=min(S,8),
            scalar_freq=1,
            just_gif=True)
        try:
            sample = next(iterloader_x)
        except StopIteration:
            iterloader_x = iter(dataloader_x)
            sample = next(iterloader_x)
        iter_rtime = time.time()-iter_start_time
        with torch.no_grad():
            metrics = test_on_fullseq(model, sample, sw_x, iters=iters, S_max=S, image_size=image_size)
        for key in list(pools_x.keys()):
            if key in metrics:
                pools_x[key].update([metrics[key]])
                sw_x.summ_scalar('_/%s' % (key), pools_x[key].mean())
        iter_itime = time.time()-iter_start_time
        
        print('%s; step %06d/%d; rtime %.2f; itime %.2f; d_x %.1f; sur_x %.1f; med_x %.1f' % (
            model_name, global_step, max_iters, iter_rtime, iter_itime,
            pools_x['d_avg'].mean(), pools_x['survival'].mean(), pools_x['median_l2'].mean()))
        
    writer_x.close()

if __name__ == '__main__':
    Fire(main)

test_on_pod.py
aharley/pips2
import time
import numpy as np
import saverloader
from nets.pips2 import Pips
import utils.improc
import utils.misc
import random
from utils.basic import print_, print_stats
import torch
from tensorboardX import SummaryWriter
import torch.nn.functional as F
from fire import Fire
from torch.utils.data import Dataset, DataLoader
from datasets.tapviddataset_fullseq import TapVidDavis

def create_pools(n_pool=1000):
    pools = {}
    pool_names = [
        'd_1',
        'd_2',
        'd_4',
        'd_8',
        'd_16',
        'd_avg',
        'median_l2',
        'survival',
        'ate_all',
        'ate_vis',
        'ate_occ',
        'total_loss',
    ]
    for pool_name in pool_names:
        pools[pool_name] = utils.misc.SimplePool(n_pool, version='np')
    return pools

def test_on_fullseq(model, d, sw, iters=8, S_max=8, image_size=(384,512)):
    metrics = {}

    rgbs = d['rgbs'].cuda().float() # B,S,C,H,W
    trajs_g = d['trajs'].cuda().float() # B,S,N,2
    valids = d['valids'].cuda().float() # B,S,N

    B, S, C, H, W = rgbs.shape
    B, S, N, D = trajs_g.shape
    assert(D==2)
    assert(B==1)
    # print('this video is %d frames long' % S)

    rgbs_ = rgbs.reshape(B*S, C, H, W)
    H_, W_ = image_size
    sy = H_/H
    sx = W_/W
    rgbs_ = F.interpolate(rgbs_, (H_, W_), mode='bilinear')
    rgbs = rgbs_.reshape(B, S, 3, H_, W_)
    trajs_g[:,:,:,0] *= sx
    trajs_g[:,:,:,1] *= sy
    H, W = H_, W_
    
    # zero-vel init
    trajs_e = trajs_g[:,0].repeat(1,S,1,1)
    
    cur_frame = 0
    done = False
    feat_init = None
    while not done:
        end_frame = cur_frame + S_max

        if end_frame > S:
            diff = end_frame-S
            end_frame = end_frame-diff
            cur_frame = max(cur_frame-diff,0)
        # print('working on subseq %d:%d' % (cur_frame, end_frame))

        traj_seq = trajs_e[:, cur_frame:end_frame]
        rgb_seq = rgbs[:, cur_frame:end_frame]
        S_local = rgb_seq.shape[1]

        if feat_init is not None:
            feat_init = [fi[:,:S_local] for fi in feat_init]
            
        preds, preds_anim, feat_init, _ = model(traj_seq, rgb_seq, iters=iters, feat_init=feat_init)

        trajs_e[:, cur_frame:end_frame] = preds[-1][:, :S_local]
        trajs_e[:, end_frame:] = trajs_e[:, end_frame-1:end_frame] # update the future with new zero-vel

        if sw is not None and sw.save_this:
            traj_seq_e = preds[-1]
            traj_seq_g = trajs_g[:,cur_frame:end_frame]
            valid_seq = valids[:,cur_frame:end_frame]
            prep_rgbs = utils.improc.preprocess_color(rgb_seq)
            gray_rgbs = torch.mean(prep_rgbs, dim=2, keepdim=True).repeat(1, 1, 3, 1, 1)
            gt_rgb = utils.improc.preprocess_color(sw.summ_traj2ds_on_rgb('', traj_seq_g, gray_rgbs[0:1].mean(dim=1), valids=valid_seq, cmap='winter', only_return=True))
            rgb_vis = []
            for tre in preds_anim:
                ate = torch.norm(tre - traj_seq_g, dim=-1) # B,S,N
                ate_all = utils.basic.reduce_masked_mean(ate, valid_seq, dim=[1,2]) # B
                rgb_vis.append(sw.summ_traj2ds_on_rgb('', tre[0:1], gt_rgb, valids=valid_seq, only_return=True, cmap='spring', frame_id=ate_all[0]))
            sw.summ_rgbs('3_test/animated_trajs_on_rgb_cur%02d' % cur_frame, rgb_vis)
        
        if end_frame >= S:
            done = True
        else:
            cur_frame = cur_frame + S_max - 1

    d_sum = 0.0
    thrs = [1,2,4,8,16]
    sx_ = W / 256.0
    sy_ = H / 256.0
    sc_py = np.array([sx_, sy_]).reshape([1,1,2])
    sc_pt = torch.from_numpy(sc_py).float().cuda()
    for thr in thrs:
        # note we exclude timestep0 from this eval
        d_ = (torch.norm(trajs_e[:,1:]/sc_pt - trajs_g[:,1:]/sc_pt, dim=-1) < thr).float() # B,S-1,N
        d_ = utils.basic.reduce_masked_mean(d_, valids[:,1:]).item()*100.0
        d_sum += d_
        metrics['d_%d' % thr] = d_
    d_avg = d_sum / len(thrs)
    metrics['d_avg'] = d_avg
    
    sur_thr = 16
    dists = torch.norm(trajs_e/sc_pt - trajs_g/sc_pt, dim=-1) # B,S,N
    dist_ok = 1 - (dists > sur_thr).float() * valids # B,S,N
    survival = torch.cumprod(dist_ok, dim=1) # B,S,N
    metrics['survival'] = torch.mean(survival).item()*100.0

    # get the median l2 error for each trajectory
    dists_ = dists.permute(0,2,1).reshape(B*N,S)
    valids_ = valids.permute(0,2,1).reshape(B*N,S)
    median_l2 = utils.basic.reduce_masked_median(dists_, valids_, keep_batch=True)
    metrics['median_l2'] = median_l2.mean().item()

    if sw is not None and sw.save_this:
        prep_rgbs = utils.improc.preprocess_color(rgbs)
        rgb0 = sw.summ_traj2ds_on_rgb('', trajs_g[0:1], prep_rgbs[0:1,0], valids=valids[0:1], cmap='winter', linewidth=2, only_return=True)
        sw.summ_traj2ds_on_rgb('2_outputs/trajs_e_on_rgb0', trajs_e[0:1], utils.improc.preprocess_color(rgb0), valids=valids[0:1], cmap='spring', linewidth=2, frame_id=d_avg)
        st = 4
        sw.summ_traj2ds_on_rgbs2('2_outputs/trajs_e_on_rgbs2', trajs_e[0:1,::st], valids[0:1,::st], prep_rgbs[0:1,::st], valids=valids[0:1,::st], frame_ids=list(range(0,S,st)))

    return metrics

def main(
        B=1, # batchsize 
        S=128, # seqlen
        stride=8, # spatial stride of the model 
        iters=16, # inference steps of the model
        image_size=(512,896), # input resolution
        shuffle=False, # dataset shuffling
        log_freq=99, # how often to make image summaries
        max_iters=30, # how many samples to test
        log_dir='./logs_test_on_tap',
        dataset_location='/orion/u/aharley/datasets/tapvid_davis',
        init_dir='./reference_model',
        device_ids=[0], 
        n_pool=1000, # how long the running averages should be
):
    device = 'cuda:%d' % device_ids[0]

    # the idea in this file is:
    # load a ckpt, and test it in tapvid,
    # tracking points from frame0 to the end.
    
    exp_name = 'tap00' # copy from dev repo
    exp_name = 'tap01' # clean up

    assert(B==1) # B>1 not implemented here
    assert(image_size[0] % 32 == 0)
    assert(image_size[1] % 32 == 0)
    
    # autogen a descriptive name
    model_name = "%d_%d" % (B,S)
    model_name += "_i%d" % (iters)
    model_name += "_%s" % exp_name
    import datetime
    model_date = datetime.datetime.now().strftime('%H%M%S')
    model_name = model_name + '_' + model_date
    print('model_name', model_name)
    
    writer_x = SummaryWriter(log_dir + '/' + model_name + '/x', max_queue=10, flush_secs=60)

    dataset_x = TapVidDavis(
        dataset_location=dataset_location,
    )
    dataloader_x = DataLoader(
        dataset_x,
        batch_size=1,
        shuffle=shuffle,
        num_workers=1)
    iterloader_x = iter(dataloader_x)

    model = Pips(stride=stride).to(device)
    model = torch.nn.DataParallel(model, device_ids=device_ids)

    utils.misc.count_parameters(model)

    _ = saverloader.load(init_dir, model.module)
    model.eval()

    pools_x = create_pools(n_pool)

    global_step = 0
    max_iters = min(max_iters, len(dataset_x))
    while global_step < max_iters:
        global_step += 1
        iter_start_time = time.time()
        with torch.no_grad():
            torch.cuda.empty_cache()
        sw_x = utils.improc.Summ_writer(
            writer=writer_x,
            global_step=global_step,
            log_freq=log_freq,
            fps=min(S,8),
            scalar_freq=1,
            just_gif=True)
        try:
            sample = next(iterloader_x)
        except StopIteration:
            iterloader_x = iter(dataloader_x)
            sample = next(iterloader_x)
        iter_rtime = time.time()-iter_start_time
        with torch.no_grad():
            metrics = test_on_fullseq(model, sample, sw_x, iters=iters, S_max=S, image_size=image_size)
        for key in list(pools_x.keys()):
            if key in metrics:
                pools_x[key].update([metrics[key]])
                sw_x.summ_scalar('_/%s' % (key), pools_x[key].mean())
        iter_itime = time.time()-iter_start_time
        
        print('%s; step %06d/%d; rtime %.2f; itime %.2f; d_x %.1f; sur_x %.1f; med_x %.1f' % (
            model_name, global_step, max_iters, iter_rtime, iter_itime,
            pools_x['d_avg'].mean(), pools_x['survival'].mean(), pools_x['median_l2'].mean()))
            
    writer_x.close()
            
if __name__ == '__main__':
    Fire(main)

test_on_tap.py
aharley/pips2
import time
import numpy as np
import timeit
import saverloader
from nets.pips2 import Pips
import utils.improc
import utils.geom
import utils.misc
import random
from utils.basic import print_, print_stats
from datasets.exportdataset import ExportDataset
import torch
import torch.nn as nn
from tensorboardX import SummaryWriter
import torch.nn.functional as F
from fire import Fire
import sys
from torch import nn, einsum
from einops import rearrange, repeat
from einops.layers.torch import Rearrange, Reduce
from torch.utils.data import Dataset, DataLoader

def create_pools(n_pool=1000):
    pools = {}
    pool_names = [
        'l1',
        'd_1',
        'd_2',
        'd_4',
        'd_8',
        'd_16',
        'd_avg',
        'l1_vis',
        'ate_all',
        'ate_vis',
        'ate_occ',
        'median_l2',
        'survival',
        'total_loss',
    ]
    for pool_name in pool_names:
        pools[pool_name] = utils.misc.SimplePool(n_pool, version='np')
    return pools

def requires_grad(parameters, flag=True):
    for p in parameters:
        p.requires_grad = flag

def fetch_optimizer(lr, wdecay, epsilon, num_steps, params):
    optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=wdecay, eps=epsilon)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, lr, num_steps+100, pct_start=0.1, cycle_momentum=False, anneal_strategy='linear')
    return optimizer, scheduler

def val_model(model, d, device, iters=8, sw=None, is_train=False):
    metrics = {}

    rgbs = d['rgbs'].float().to(device) # B,S,C,H,W
    track_g = d['track_g'].float().to(device) # B,S,N,8
    trajs_g = track_g[:,:,:,:2]
    vis_g = track_g[:,:,:,2]
    valids = track_g[:,:,:,3]

    B, S, C, H, W = rgbs.shape
    B, S, N, D = trajs_g.shape
    assert(D==2)

    # zero-vel init
    trajs_e0 = trajs_g[:,0:1].repeat(1,S,1,1)

    preds, preds_anim, _, _ = model(trajs_e0, rgbs, iters=iters)
    trajs_e = preds[-1]

    l1_dists = torch.abs(trajs_e - trajs_g).sum(dim=-1) # B,S,N
    l1_loss = utils.basic.reduce_masked_mean(l1_dists, valids)
    l1_vis = utils.basic.reduce_masked_mean(l1_dists, valids*vis_g)
    
    ate = torch.norm(trajs_e - trajs_g, dim=-1) # B,S,N
    ate_all = utils.basic.reduce_masked_mean(ate, valids, dim=[1,2])
    ate_vis = utils.basic.reduce_masked_mean(ate, valids*vis_g)
    ate_occ = utils.basic.reduce_masked_mean(ate, valids*(1.0-vis_g))

    metrics['l1'] = l1_loss.mean().item()
    metrics['l1_vis'] = l1_vis.mean().item()
    metrics['ate_all'] = ate_all.mean().item()
    metrics['ate_vis'] = ate_vis.item()
    metrics['ate_occ'] = ate_occ.item()
    
    if sw is not None and sw.save_this:
        prep_rgbs = utils.improc.preprocess_color(rgbs)
        prep_grays = torch.mean(prep_rgbs, dim=2, keepdim=True).repeat(1, 1, 3, 1, 1)
        gt_rgb = utils.improc.preprocess_color(sw.summ_traj2ds_on_rgb('', trajs_g[0:1], prep_grays[0:1].mean(dim=1), valids=valids[0:1], cmap='winter', only_return=True))
        rgb_vis = []
        for tre in preds_anim:
            ate = torch.norm(tre - trajs_g, dim=-1) # B,S,N
            ate_all = utils.basic.reduce_masked_mean(ate, valids, dim=[1,2]) # B
            rgb_vis.append(sw.summ_traj2ds_on_rgb('', tre[0:1], gt_rgb, valids=valids[0:1], only_return=True, cmap='spring', frame_id=ate_all[0]))
        sw.summ_rgbs('3_test/animated_trajs_on_rgb', rgb_vis)
        
    d_sum = 0.0
    thrs = [1,2,4,8,16]
    sx_ = W / 256.0
    sy_ = H / 256.0
    sc_py = np.array([sx_, sy_]).reshape([1,1,2])
    sc_pt = torch.from_numpy(sc_py).float().cuda()
    for thr in thrs:
        # note we exclude timestep0 from this eval
        d_ = (torch.norm(trajs_e[:,1:]/sc_pt - trajs_g[:,1:]/sc_pt, dim=-1) < thr).float() # B,S-1,N
        d_ = utils.basic.reduce_masked_mean(d_, valids[:,1:]).item()*100.0
        d_sum += d_
        metrics['d_%d' % thr] = d_
    d_avg = d_sum / len(thrs)
    metrics['d_avg'] = d_avg

    sur_thr = 16
    dists = torch.norm(trajs_e/sc_pt - trajs_g/sc_pt, dim=-1) # B,S,N
    dist_ok = 1 - (dists > sur_thr).float() * valids # B,S,N
    survival = torch.cumprod(dist_ok, dim=1) # B,S,N
    metrics['survival'] = torch.mean(survival).item()*100.0
    
    # get the median l2 error for each trajectory
    dists_ = dists.permute(0,2,1).reshape(B*N,S)
    valids_ = valids.permute(0,2,1).reshape(B*N,S)
    median_l2 = utils.basic.reduce_masked_median(dists_, valids_, keep_batch=True) # B*N
    metrics['median_l2'] = median_l2.mean().item()

    return metrics

def run_model(model, d, device, iters=8, sw=None, is_train=True, use_augs=True):
    total_loss = torch.tensor(0.0, requires_grad=True).to(device)

    metrics = {}

    rgbs = d['rgbs'].float().to(device) # B,S,C,H,W
    track_g = d['track_g'].float().to(device) # B,S,N,8
    trajs_g = track_g[:,:,:,:2]
    vis_g = track_g[:,:,:,2]
    valids = track_g[:,:,:,3]

    if use_augs and np.random.rand() < 0.5: # rot90 aug
        rgbs = rgbs.permute(0,1,2,4,3) # swap xy
        trajs_g = trajs_g.flip([3]) # swap xy
        
    B, S, C, H, W = rgbs.shape
    assert(C==3)
    B, S, N, D = trajs_g.shape
    assert(D==2)

    # full random
    x = torch.from_numpy(np.random.uniform(0, W-1, (B,S,N))).float().to(trajs_g.device)
    y = torch.from_numpy(np.random.uniform(0, H-1, (B,S,N))).float().to(trajs_g.device)
    trajs_e0 = torch.stack([x,y], dim=-1) # B,S,N,2
    # mix with gt a random amount
    coeff = torch.from_numpy(np.random.uniform(0, 1, (B,1,N,1))).float().to(trajs_g.device)
    trajs_e0 = trajs_e0*coeff + trajs_g*(1-coeff)
    # use zero-velocity init for some
    trajs_z = trajs_g[:,0:1].repeat(1,S,1,1)
    mask = (torch.from_numpy(np.random.uniform(0, 1, (B,1,N,1))).float().to(trajs_g.device)>0.5).float()
    trajs_e0 = trajs_e0*mask + trajs_z*(1-mask)
    # reset zeroth on all
    trajs_e0[:,0:1] = trajs_g[:,0:1]

    # measure our initial distance, so we can check our improvement
    ate0 = torch.norm(trajs_e0 - trajs_g, dim=-1) # B,S,N
    ate0_all = utils.basic.reduce_masked_mean(ate0, valids, dim=[1,2])
    
    preds, preds_anim, _, loss = model(trajs_e0, rgbs, iters=iters, trajs_g=trajs_g, vis_g=vis_g, valids=valids, is_train=is_train)
    trajs_e = preds[-1]

    total_loss += loss

    # collect stats
    l1_dists = torch.abs(trajs_e - trajs_g).sum(dim=-1) # B,S,N
    l1_loss = utils.basic.reduce_masked_mean(l1_dists, valids)
    l1_vis = utils.basic.reduce_masked_mean(l1_dists, valids*vis_g)
    ate = torch.norm(trajs_e - trajs_g, dim=-1) # B,S,N
    ate_all = utils.basic.reduce_masked_mean(ate, valids, dim=[1,2])
    ate_vis = utils.basic.reduce_masked_mean(ate, valids*vis_g)
    ate_occ = utils.basic.reduce_masked_mean(ate, valids*(1.0-vis_g))
    metrics['l1'] = l1_loss.mean().item()
    metrics['l1_vis'] = l1_vis.mean().item()
    metrics['ate_all'] = ate_all.mean().item()
    metrics['ate_vis'] = ate_vis.item()
    metrics['ate_occ'] = ate_occ.item()
    metrics['total_loss'] = total_loss.item()

    d_sum = 0.0
    thrs = [1,2,4,8,16]
    sx_ = W / 256.0
    sy_ = H / 256.0
    sc_py = np.array([sx_, sy_]).reshape([1,1,2])
    sc_pt = torch.from_numpy(sc_py).float().to(device)
    for thr in thrs:
        # note we exclude timestep0 from this eval
        d_ = (torch.norm(trajs_e[:,1:]/sc_pt - trajs_g[:,1:]/sc_pt, dim=-1) < thr).float() # B,S-1,N
        d_ = utils.basic.reduce_masked_mean(d_, valids[:,1:]).item()*100.0
        d_sum += d_
        metrics['d_%d' % thr] = d_
    d_avg = d_sum / len(thrs)
    metrics['d_avg'] = d_avg

    sur_thr = 16
    dists = torch.norm(trajs_e/sc_pt - trajs_g/sc_pt, dim=-1) # B,S,N
    dist_ok = 1 - (dists > sur_thr).float() * valids # B,S,N
    survival = torch.cumprod(dist_ok, dim=1) # B,S,N
    metrics['survival'] = torch.mean(survival).item()*100.0
    
    # get the median l2 error for each trajectory
    dists_ = dists.permute(0,2,1).reshape(B*N,S)
    valids_ = valids.permute(0,2,1).reshape(B*N,S)
    val_ok = valids_[:,0] > 0 # get rid of the ones we padded in
    dists_ = dists_[val_ok]
    valids_ = valids_[val_ok]
    median_l2 = utils.basic.reduce_masked_median(dists_, valids_, keep_batch=True) # B*N
    metrics['median_l2'] = median_l2.mean().item()

    if sw is not None and sw.save_this:
        prep_rgbs = utils.improc.preprocess_color(rgbs)
        prep_grays = torch.mean(prep_rgbs, dim=2, keepdim=True).repeat(1, 1, 3, 1, 1)

        rgb0 = sw.summ_traj2ds_on_rgb('', trajs_g[0:1], prep_rgbs[0:1,0], valids=valids[0:1], cmap='winter', linewidth=2, only_return=True)
        sw.summ_traj2ds_on_rgb('0_inputs/trajs_e0_on_rgb0', trajs_e0[0:1], utils.improc.preprocess_color(rgb0), valids=valids[0:1], cmap='spring', linewidth=2, frame_id=ate0_all[0].mean().item())
        sw.summ_traj2ds_on_rgb('2_outputs/trajs_e_on_rgb0', trajs_e[0:1], utils.improc.preprocess_color(rgb0), valids=valids[0:1], cmap='spring', linewidth=2, frame_id=ate_all[0].mean().item())
        sw.summ_traj2ds_on_rgbs2('0_inputs/trajs_g_on_rgbs2', trajs_g[0:1,::4], vis_g[0:1,::4], prep_rgbs[0:1,::4], valids=valids[0:1,::4], frame_ids=list(range(0,S,4)))
        
        # in the kp vis, clamp so that we can see everything
        trajs_g_clamp = trajs_g.clone()
        trajs_g_clamp[:,:,:,0] = trajs_g_clamp[:,:,:,0].clip(0,W-1)
        trajs_g_clamp[:,:,:,1] = trajs_g_clamp[:,:,:,1].clip(0,H-1)
        trajs_e_clamp = trajs_e.clone()
        trajs_e_clamp[:,:,:,0] = trajs_e_clamp[:,:,:,0].clip(0,W-1)
        trajs_e_clamp[:,:,:,1] = trajs_e_clamp[:,:,:,1].clip(0,H-1)

        gt_rgb = utils.improc.preprocess_color(sw.summ_traj2ds_on_rgb('', trajs_g[0:1], prep_grays[0:1].mean(dim=1), valids=valids[0:1], cmap='winter', only_return=True))
        rgb_vis = []
        for tre in preds_anim:
            ate = torch.norm(tre - trajs_g, dim=-1) # B,S,N
            ate_all = utils.basic.reduce_masked_mean(ate, valids, dim=[1,2]) # B
            rgb_vis.append(sw.summ_traj2ds_on_rgb('', tre[0:1], gt_rgb, valids=valids[0:1], only_return=True, cmap='spring', frame_id=ate_all[0]))
        sw.summ_rgbs('3_test/animated_trajs_on_rgb', rgb_vis)

        outs = sw.summ_pts_on_rgbs(
            '',
            trajs_g_clamp[0:1,::4],
            prep_grays[0:1,::4],
            valids=valids[0:1,::4],
            cmap='winter', linewidth=3, only_return=True)
        sw.summ_pts_on_rgbs(
            '0_inputs/kps_gv_on_rgbs',
            trajs_g_clamp[0:1,::4],
            utils.improc.preprocess_color(outs),
            valids=valids[0:1,::4]*vis_g[0:1,::4],
            cmap='spring', linewidth=2)

        outs = sw.summ_pts_on_rgbs(
            '',
            trajs_g_clamp[0:1,::4],
            prep_grays[0:1,::4],
            valids=valids[0:1,::4],
            cmap='winter', linewidth=3, only_return=True)
        sw.summ_pts_on_rgbs(
            '2_outputs/kps_eg_on_rgbs',
            trajs_e_clamp[0:1,::4],
            utils.improc.preprocess_color(outs),
            valids=valids[0:1,::4],
            cmap='spring', linewidth=2)
        
    return total_loss, metrics
    

def main(
        B=2, # batchsize 
        S=36, # seqlen
        N=64, # number of points per clip
        stride=8, # spatial stride of the model 
        iters=6, # inference steps of the model
        crop_size=(384,512), # raw flt data is 540,960
        use_augs=True, # resizing/jittering/color/blur augs
        shuffle=True, # dataset shuffling
        cache_len=0, # how many samples to cache into ram (for overfitting/debug)
        cache_freq=0, # how often to add a new sample to cache
        dataset_location='./pod_export', # where we exported the data
        dataset_version='ae_36_128_384x512', # export version
        n_pool=1000, # size of running avg for stats
        quick=False, # debug
        # optimization
        lr=5e-4,
        grad_acc=1, 
        use_scheduler=True,
        max_iters=200000,
        # summaries
        log_dir='./logs_train',
        log_freq=1000,
        val_freq=100,
        # saving/loading
        ckpt_dir='./checkpoints',
        save_freq=1000,
        keep_latest=2,
        init_dir='',
        load_optimizer=True,
        load_step=True,
        ignore_load=None,
        device_ids=[0],
):
    device = 'cuda:%d' % device_ids[0]

    # the idea in this file is:
    # train from scratch on pointodyssey.
    # on val steps, unroll the inference.

    exp_name = 'aa00' # copy from dev repo
    exp_name = 'aa01' # clean up

    if quick: # (debug)
        B = 1
        log_freq = 100
        max_iters = 1000
        shuffle = False
        val_freq = 10
        n_pool = 100
        use_augs = False
        cache_len = 0 # overfit on this many
        cache_freq = 0
        save_freq = 99999999
        
    if init_dir:
        init_dir = '%s/%s' % (ckpt_dir, init_dir)
        
    assert(crop_size[0] % 32 == 0)
    assert(crop_size[1] % 32 == 0)
    
    # autogen a descriptive name
    model_name = "%d_%d_%d" % (B,S,N)
    model_name += "_i%d" % (iters)
    if grad_acc > 1:
        model_name += "x%d" % grad_acc
    lrn = "%.1e" % lr # e.g., 5.0e-04
    lrn = lrn[0] + lrn[3:5] + lrn[-1] # e.g., 5e-4
    model_name += "_%s" % lrn
    if use_scheduler:
        model_name += "s"
    if cache_len:
        model_name += "_c%d_f%d" % (cache_len, cache_freq)
    if use_augs:
        model_name += "_A"
    model_name += "_%s" % exp_name
    import datetime
    model_date = datetime.datetime.now().strftime('%H%M%S')
    model_name = model_name + '_' + model_date
    print('model_name', model_name)
    
    ckpt_path = '%s/%s' % (ckpt_dir, model_name)
    writer_t = SummaryWriter(log_dir + '/' + model_name + '/t', max_queue=10, flush_secs=60)
    if val_freq:
        writer_v = SummaryWriter(log_dir + '/' + model_name + '/v', max_queue=10, flush_secs=60)

    def worker_init_fn(worker_id):
        np.random.seed(np.random.get_state()[1][0] + worker_id)
    dataset_t = ExportDataset(
        dataset_location=dataset_location,
        dataset_version=dataset_version,
        S=S,
        crop_size=crop_size,
        use_augs=use_augs)
    dataloader_t = DataLoader(
        dataset_t,
        batch_size=B,
        shuffle=shuffle,
        num_workers=6,
        worker_init_fn=worker_init_fn,
        drop_last=True)
    iterloader_t = iter(dataloader_t)

    if cache_len:
        sample_pool_t = utils.misc.SimplePool(cache_len, version='np')
    
    model = Pips(stride=stride).to(device)
    model = torch.nn.DataParallel(model, device_ids=device_ids)

    parameters = list(model.parameters())
    weight_decay = 1e-6
    if use_scheduler:
        optimizer, scheduler = fetch_optimizer(lr, weight_decay, 1e-8, max_iters, model.parameters())
    else:
        optimizer = torch.optim.AdamW(parameters, lr=lr, weight_decay=weight_decay) 
        scheduler = None

    utils.misc.count_parameters(model)

    global_step = 0
    if init_dir:
        if load_step and load_optimizer:
            global_step = saverloader.load(init_dir, model.module, optimizer=optimizer, scheduler=scheduler, ignore_load=ignore_load)
        elif load_step:
            global_step = saverloader.load(init_dir, model.module, ignore_load=ignore_load)
        else:
            _ = saverloader.load(init_dir, model.module, ignore_load=ignore_load)
            global_step = 0
    
    requires_grad(parameters, True)
    model.train()

    pools_t = create_pools(n_pool)
    if val_freq:
        pools_v = create_pools(n_pool)
    
    while global_step < max_iters:
        global_step += 1
        iter_start_time = time.time()
        iter_rtime = 0.0
        
        for internal_step in range(grad_acc):
            read_start_time = time.time()

            if internal_step==grad_acc-1:
                sw_t = utils.improc.Summ_writer(
                    writer=writer_t,
                    global_step=global_step,
                    log_freq=log_freq,
                    fps=min(S,8),
                    scalar_freq=log_freq//4,
                    just_gif=True)
            else:
                sw_t = None

            read_new = True # read something from the dataloder
            if cache_len: 
                read_new = False
                if len(sample_pool_t) < cache_len:
                    read_new = True
                if cache_freq > 0 and global_step % cache_freq == 0:
                    read_new = True
            if read_new:
                try:
                    sample = next(iterloader_t)
                except StopIteration:
                    iterloader_t = iter(dataloader_t)
                    sample = next(iterloader_t)
                if cache_len:
                    sample_pool_t.update([sample])
                    print('cached a new sample into sample_pool (len %d)' % (len(sample_pool_t)))
            if cache_len:
                sample = sample_pool_t.sample()
                    
            iter_rtime += time.time()-read_start_time

            total_loss, metrics = run_model(
                model, sample, device,
                iters=iters,
                sw=sw_t,
                is_train=True,
                use_augs=use_augs)

            if torch.isnan(total_loss):
                print('nan in loss; quitting')
                return False

            total_loss /= grad_acc
            total_loss.backward()

        sw_t.summ_scalar('total_loss', metrics['total_loss'])

        for key in list(pools_t.keys()):
            if key in metrics:
                pools_t[key].update([metrics[key]])
            sw_t.summ_scalar('_/%s' % (key), pools_t[key].mean())

        current_lr = optimizer.param_groups[0]['lr']
        sw_t.summ_scalar('_/current_lr', current_lr)
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        optimizer.step()
        if use_scheduler:
            scheduler.step()
        optimizer.zero_grad()

        if np.mod(global_step, save_freq)==0:
            saverloader.save(ckpt_path, optimizer, model.module, global_step, scheduler=scheduler, keep_latest=keep_latest)

        if val_freq and (global_step) % val_freq == 0:
            model.eval()
            del sample
            with torch.no_grad():
                torch.cuda.empty_cache()
            sw_v = utils.improc.Summ_writer(
                writer=writer_v,
                global_step=global_step,
                log_freq=log_freq,
                fps=min(S,8),
                scalar_freq=log_freq//4,
                just_gif=True)
            if cache_len:
                sample = sample_pool_t.sample()
            else:
                try:
                    sample = next(iterloader_t)
                except StopIteration:
                    iterloader_t = iter(dataloader_t)
                    sample = next(iterloader_t)
            with torch.no_grad():
                metrics = val_model(
                    model, sample, device,
                    iters=iters*2,
                    sw=sw_v,
                    is_train=False)
            for key in list(pools_v.keys()):
                if key in metrics:
                    pools_v[key].update([metrics[key]])
                sw_v.summ_scalar('_/%s' % (key), pools_v[key].mean())
            model.train()
                
                    
        iter_itime = time.time()-iter_start_time

        if val_freq:
            print('%s; step %06d/%d; rtime %.2f; itime %.2f; loss %.3f; loss_t %.2f; d_t %.1f; d_v %.1f' % (
                model_name, global_step, max_iters, iter_rtime, iter_itime,
                total_loss.item(), pools_t['total_loss'].mean(), pools_t['d_avg'].mean(), pools_v['d_avg'].mean()))
        else:
            print('%s; step %06d/%d; rtime %.2f; itime %.2f; loss %.3f; loss_t %.2f; d_t %.1f' % (
                model_name, global_step, max_iters, iter_rtime, iter_itime,
                total_loss.item(), pools_t['total_loss'].mean(), pools_t['d_avg'].mean()))
            
    writer_t.close()
    if val_freq:
        writer_v.close()

if __name__ == '__main__':
    Fire(main)

train.py
aharley/pips2
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())
    if d.is_cuda:
        out = out.cuda()
    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):
    # 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
    for (a,b) in zip(x.size(), mask.size()):
        # if not b==1: 
        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 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 cuda:
    #     grid_z = grid_z.cuda()
    #     grid_y = grid_y.cuda()
    #     grid_x = grid_x.cuda()
        
    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):
    C = channels
    xy_grid = gridcloud2d(C, kernel_size, kernel_size) # 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)
    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

utils/basic.py
aharley/pips2
import torch
import utils.basic
import numpy as np
import torchvision.ops as ops
from utils.basic import print_

def matmul2(mat1, mat2):
    return torch.matmul(mat1, mat2)

def matmul3(mat1, mat2, mat3):
    return torch.matmul(mat1, torch.matmul(mat2, mat3))

def eye_3x3(B, device='cuda'):
    rt = torch.eye(3, device=torch.device(device)).view(1,3,3).repeat([B, 1, 1])
    return rt

def eye_4x4(B, device='cuda'):
    rt = torch.eye(4, device=torch.device(device)).view(1,4,4).repeat([B, 1, 1])
    return rt

def safe_inverse(a): #parallel version
    B, _, _ = list(a.shape)
    inv = a.clone()
    r_transpose = a[:, :3, :3].transpose(1,2) #inverse of rotation matrix

    inv[:, :3, :3] = r_transpose
    inv[:, :3, 3:4] = -torch.matmul(r_transpose, a[:, :3, 3:4])

    return inv

def safe_inverse_single(a):
    r, t = split_rt_single(a)
    t = t.view(3,1)
    r_transpose = r.t()
    inv = torch.cat([r_transpose, -torch.matmul(r_transpose, t)], 1)
    bottom_row = a[3:4, :] # this is [0, 0, 0, 1]
    # bottom_row = torch.tensor([0.,0.,0.,1.]).view(1,4)
    inv = torch.cat([inv, bottom_row], 0)
    return inv

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(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_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 get_camM_T_camXs(origin_T_camXs, ind=0):
    B, S = list(origin_T_camXs.shape)[0:2]
    camM_T_camXs = torch.zeros_like(origin_T_camXs)
    for b in list(range(B)):
        camM_T_origin = safe_inverse_single(origin_T_camXs[b,ind])
        for s in list(range(S)):
            camM_T_camXs[b,s] = torch.matmul(camM_T_origin, origin_T_camXs[b,s])
    return camM_T_camXs

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_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_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]
pp        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

    # 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

    # 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 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

utils/geom.py
aharley/pips2
import torch
import numpy as np
import utils.basic
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().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 = 3

    if together:
        reduced_emb = pca_embed_together(emb, keep)
    else:
        reduced_emb = pca_embed(emb, keep, valid) #not im

    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 convert_occ_to_height(occ, reduce_axis=3):
    B, C, D, H, W = list(occ.shape)
    assert(C==1)
    # note that height increases DOWNWARD in the tensor
    # (like pixel/camera coordinates)
    
    G = list(occ.shape)[reduce_axis]
    values = torch.linspace(float(G), 1.0, steps=G, dtype=torch.float32, device=occ.device)
    if reduce_axis==2:
        # fro view
        values = values.view(1, 1, G, 1, 1)
    elif reduce_axis==3:
        # top view
        values = values.view(1, 1, 1, G, 1)
    elif reduce_axis==4:
        # lateral view
        values = values.view(1, 1, 1, 1, G)
    else:
        assert(False) # you have to reduce one of the spatial dims (2-4)
    values = torch.max(occ*values, dim=reduce_axis)[0]/float(G)
    # values = values.view([B, C, D, W])
    return values

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):

    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
    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)
        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.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)
        

    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,N).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, scores, tids, linewidth=1):
        # all inputs are numpy tensors
        # rgb is H x W x 3
        # boxlist is N x 4
        # scores is N
        # tids 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)

        # color_map = cm.get_cmap('tab20')
        # color_map = cm.get_cmap('set1')
        color_map = cm.get_cmap('Accent')
        color_map = color_map.colors
        # print('color_map', color_map)

        # draw
        for ind, box in enumerate(boxlist):
            # box is 4
            if not np.isclose(scores[ind], 0.0):
                # box = utils.geom.scale_box2d(box, H, W)
                ymin, xmin, ymax, xmax = box

                # ymin, ymax = ymin*H, ymax*H
                # xmin, xmax = xmin*W, xmax*W
                
                # print 'score = %.2f' % scores[ind]
                # color_id = tids[ind] % 20
                color_id = tids[ind]
                color = color_map[color_id]
                color = np.array(color)*255.0
                color = color.round()
                # color = color.astype(np.uint8)
                # color = color[::-1]
                # print('color', color)
                
                # print 'tid = %d; score = %.3f' % (tids[ind], scores[ind])

                # if False:
                if scores[ind] < 1.0: # not gt
                    cv2.putText(rgb,
                                # '%d (%.2f)' % (tids[ind], scores[ind]), 
                                '%.2f' % (scores[ind]), 
                                (int(xmin), int(ymin)),
                                cv2.FONT_HERSHEY_SIMPLEX,
                                0.5, # font size
                                color),
                    #1) # font weight

                xmin = np.clip(int(xmin), 0,  W-1)
                xmax = np.clip(int(xmax), 0,  W-1)
                ymin = np.clip(int(ymin), 0,  H-1)
                ymax = np.clip(int(ymax), 0,  H-1)

                cv2.line(rgb, (xmin, ymin), (xmin, ymax), color, linewidth, cv2.LINE_AA)
                cv2.line(rgb, (xmin, ymin), (xmax, ymin), color, linewidth, cv2.LINE_AA)
                cv2.line(rgb, (xmax, ymin), (xmax, ymax), color, linewidth, cv2.LINE_AA)
                cv2.line(rgb, (xmax, ymax), (xmin, ymax), color, linewidth, cv2.LINE_AA)
                
        # 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, only_return=False, linewidth=2):
        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, only_return=only_return)
    
    def summ_rgbs(self, name, ims, frame_ids=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 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, only_return=False, halfres=False):
        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)

            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=50.0):
        """
        :param flow: Optical flow tensor.
        :return: RGB image normalized between 0 and 1.
        """

        # flow is B x C x H x W

        B, C, H, W = list(flow.size())

        flow = flow.clone().detach()
        
        abs_image = torch.abs(flow)
        flow_mean = abs_image.mean(dim=[1,2,3])
        flow_std = abs_image.std(dim=[1,2,3])

        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)

        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:], flow[:, 0:1]) / np.pi #B x 1 x H x W

        hue = torch.clamp((angle + 1.0) / 2.0, 0.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)
        return flow
    
    def summ_flow(self, name, im, clip=0.0, only_return=False, frame_id=None):
        # 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)
        else:
            return None
            
    def summ_oneds(self, name, ims, frame_ids=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 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, only_return=False):
        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, only_return=only_return)
            
    def summ_feats(self, name, feats, valids=None, pca=True, fro=False, only_return=False, frame_ids=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)

            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)

    def summ_feat(self, name, feat, valid=None, pca=True, only_return=False, bev=False, fro=False, frame_id=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)
            else:
                feat_pca = get_feat_pca(feat, valid)
                return self.summ_rgb(name, feat_pca, only_return=only_return, frame_id=frame_id)
            
    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 np.mod(self.global_step, self.scalar_freq)==0:
                self.writer.add_scalar(name, value, global_step=self.global_step)
                
    def summ_seg(self, name, seg, only_return=False, frame_id=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+1)
            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)
    
    def summ_pts_on_rgb(self, name, trajs, rgb, valids=None, frame_id=None, only_return=False, show_dots=True, cmap='coolwarm', linewidth=1):
        # 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]
        # print('trajs', trajs.shape)
        # print('valids', valids.shape)
        
        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

        rgb = rgb.astype(np.uint8).copy()
        
        for i in range(N):
            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

            color_map = cm.get_cmap(cmap)
            color = np.array(color_map(i)[:3]) * 255 # rgb
            for s in range(S):
                if valid[s]:
                    cv2.circle(rgb, (int(traj[s,0]), int(traj[s,1])), linewidth, color, -1)
        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)

    def summ_pts_on_rgbs(self, name, trajs, rgbs, valids=None, frame_ids=None, only_return=False, show_dots=True, cmap='coolwarm', linewidth=1):
        # 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]
        # print('trajs', trajs.shape)
        # print('valids', valids.shape)
        
        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

        rgbs_color = [rgb.astype(np.uint8).copy() for rgb in rgbs_color]
        
        for i in range(N):
            traj = trajs[:,i] # S,2
            valid = valids[:,i] # S

            color_map = cm.get_cmap(cmap)
            color = np.array(color_map(0)[:3]) * 255 # rgb
            for s in range(S):
                if valid[s]:
                    cv2.circle(rgbs_color[s], (traj[s,0], traj[s,1]), linewidth, color, -1)
        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)
    
    
    def summ_traj2ds_on_rgbs(self, name, trajs, rgbs, valids=None, frame_ids=None, only_return=False, show_dots=False, cmap='coolwarm', vals=None, linewidth=1):
        # 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]

        # print('trajs', trajs.shape)
        # print('valids', valids.shape)
        
        if vals is not None:
            vals = vals[0] # N
            # print('vals', vals.shape)
        
        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(N):
            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]:
                # traj_seq = traj[max(t-16,0):t+1]
                traj_seq = traj[max(t-8,0):t+1]
                val_seq = np.linspace(0,1,len(traj_seq))
                # if t<2:
                #     val_seq = np.zeros_like(val_seq)
                # print('val_seq', val_seq)
                # val_seq = 1.0
                # val_seq = np.arange(8)/8.0
                # val_seq = val_seq[-len(traj_seq):]
                # rgbs_color[t] = self.draw_traj_on_image_py(rgbs_color[t], traj_seq, S=S, show_dots=show_dots, cmap=cmap_, val=val_seq, linewidth=linewidth)
                rgbs_color[t] = self.draw_traj_on_image_py(rgbs_color[t], traj_seq, S=S, show_dots=show_dots, cmap=cmap_, val=val_seq, linewidth=linewidth)
            # input()

        for i in range(N):
            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
            vis = torch.ones_like(traj[:,0]) # S
            valid = valids[:,i] # S
            rgbs_color = self.draw_circ_on_images_py(rgbs_color, traj, vis, S=0, 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)

    def summ_traj2ds_on_rgbs2(self, name, trajs, visibles, rgbs, valids=None, frame_ids=None, only_return=False, show_dots=True, cmap=None, linewidth=1):
        # 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]
        # print('trajs', trajs.shape)
        # print('valids', valids.shape)
        
        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(N):
            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(N):
            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
            if valid[0]:
                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)

    def summ_traj2ds_on_rgb(self, name, trajs, rgb, valids=None, show_dots=False, show_lines=True, frame_id=None, only_return=False, cmap='coolwarm', linewidth=1):
        # 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
        
        for i in range(N):
            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)
    
    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:
                # if len(val) == S1:
                color = np.array(color_map(val[s])[:3]) * 255 # rgb
                # else:
                #     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, np.array(color_map(1)[:3])*255, -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, (xy[n,0], 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]))
            color = (int(color[2]),int(color[1]),int(color[0]))
            
        for s in range(S1):
            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+1)/max(1,float(S-1)))[: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+1, 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, only_return=False, show_circ=False, trajs_g=None, is_g=False):
        B, S, N, D = trajs_e.shape
        assert(N==1)
        assert(D==2)
        
        rgbs_vis = []
        n = 0
        pad_amount = 100
        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
        
        for s in range(S):
            rgb = rgbs[0,s].detach().cpu().numpy()
            # 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)
                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 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()

            rgb_ = rgb_.transpose(2,0,1)
            rgb_ = torch.from_numpy(rgb_)

            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, only_return=only_return)
        
    def summ_occ(self, name, occ, reduce_axes=[3], bev=False, fro=False, pro=False, frame_id=None, only_return=False):
        if self.save_this:
            B, C, D, H, W = list(occ.shape)
            if bev:
                reduce_axes = [3]
            elif fro:
                reduce_axes = [2]
            elif pro:
                reduce_axes = [4]
            for reduce_axis in reduce_axes:
                height = convert_occ_to_height(occ, reduce_axis=reduce_axis)
                if reduce_axis == reduce_axes[-1]:
                    return self.summ_oned(name=('%s_ax%d' % (name, reduce_axis)), im=height, norm=False, frame_id=frame_id, only_return=only_return)
                else:
                    self.summ_oned(name=('%s_ax%d' % (name, reduce_axis)), im=height, norm=False, frame_id=frame_id, only_return=only_return)
    
def erode2d(im, times=1, device='cuda'):
    weights2d = torch.ones(1, 1, 3, 3, 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='cuda', mode='square'):
    weights2d = torch.ones(1, 1, 3, 3, device=device)
    if mode=='cross':
        weights2d[:,:,0,0] = 0.0
        weights2d[:,:,0,2] = 0.0
        weights2d[:,:,2,0] = 0.0
        weights2d[:,:,2,2] = 0.0
    for time in range(times):
        im = F.conv2d(im, weights2d, padding=1).clamp(0, 1)
    return im



utils/improc.py
aharley/pips2
import torch
import numpy as np
import math
from prettytable import PrettyTable

def count_parameters(model):
    table = PrettyTable(["Modules", "Parameters"])
    total_params = 0
    for name, parameter in model.named_parameters():
        if not parameter.requires_grad:
            continue
        param = parameter.numel()
        if param > 100000:
            table.add_row([name, param])
        total_params+=param
    print(table)
    print('total params: %.2f M' % (total_params/1000000.0))
    return total_params

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

class SimplePool():
    def __init__(self, pool_size, version='pt'):
        self.pool_size = pool_size
        self.version = version
        self.items = []
        
        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=1):
        if 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 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

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).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    if deterministic:
        farthest = torch.randint(0, 1, (B,), dtype=torch.long).to(device)
    else:
        farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
    batch_indices = torch.arange(B, dtype=torch.long).to(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):
    N,C = xyz.shape
    inds = np.zeros(npoint, dtype=np.int32)
    distance = np.ones(N) * 1e10
    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
    

utils/misc.py
aharley/pips2
import torch
import utils.basic
import torch.nn.functional as F

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

def paste_crop_on_canvas(crop, box2d_unnorm, H, W, fast=True, mask=None, canvas=None):
    # this is the inverse of crop_and_resize_box2d
    B, C, Y, X = list(crop.shape)
    B2, D = list(box2d_unnorm.shape)
    assert(B == B2)
    assert(D == 4)

    # here, we want to place the crop into a bigger image,
    # at the location specified by the box2d.

    if canvas is None:
        canvas = torch.zeros((B, C, H, W), device=crop.device)
    else:
        B2, C2, H2, W2 = canvas.shape
        assert(B==B2)
        assert(C==C2)
        assert(H==H2)
        assert(W==W2)

    # box2d_unnorm = utils.geom.unnormalize_box2d(box2d, H, W)

    if fast:
        ymin = box2d_unnorm[:, 0].long()
        xmin = box2d_unnorm[:, 1].long()
        ymax = box2d_unnorm[:, 2].long()
        xmax = box2d_unnorm[:, 3].long()
        w = (xmax - xmin).float()
        h = (ymax - ymin).float()

        grids = utils.basic.gridcloud2d(B, H, W)
        grids_flat = grids.reshape(B, -1, 2)
        # grids_flat[:, :, 0] = (grids_flat[:, :, 0] - xmin.float().unsqueeze(1)) / w.unsqueeze(1) * X
        # grids_flat[:, :, 1] = (grids_flat[:, :, 1] - ymin.float().unsqueeze(1)) / h.unsqueeze(1) * Y

        # for each pixel in the main image,
        # grids_flat tells us where to sample in the crop image

        # print('grids_flat', grids_flat.shape)
        # print('crop', crop.shape)

        grids_flat[:, :, 0] = (grids_flat[:, :, 0] - xmin.float().unsqueeze(1)) / w.unsqueeze(1) * 2.0 - 1.0
        grids_flat[:, :, 1] = (grids_flat[:, :, 1] - ymin.float().unsqueeze(1)) / h.unsqueeze(1) * 2.0 - 1.0
        
        grid = grids_flat.reshape(B,H,W,2)

        canvas = F.grid_sample(crop, grid, align_corners=False)
        # print('canvas', canvas.shape)
        
        # if mask is None:
        #     crop_resamp, inb = bilinear_sample2d(crop, grids_flat[:, :, 0], grids_flat[:, :, 1], return_inbounds=True)
        #     crop_resamp = crop_resamp.reshape(B, C, H, W)
        #     inb = inb.reshape(B, 1, H, W)
        #     canvas = canvas * (1 - inb) + crop_resamp * inb
        # else:
        #     full_resamp = bilinear_sample2d(torch.cat([crop, mask], dim=1), grids_flat[:, :, 0], grids_flat[:, :, 1])
        #     full_resamp = full_resamp.reshape(B, C+1, H, W)
        #     crop_resamp = full_resamp[:,:3]
        #     mask_resamp = full_resamp[:,3:4]
        #     canvas = canvas * (1 - mask_resamp) + crop_resamp * mask_resamp
    else:
        for b in range(B):
            ymin = box2d_unnorm[b, 0].long()
            xmin = box2d_unnorm[b, 1].long()
            ymax = box2d_unnorm[b, 2].long()
            xmax = box2d_unnorm[b, 3].long()

            crop_b = F.interpolate(crop[b:b + 1], (ymax - ymin, xmax - xmin)).squeeze(0)

            # print('canvas[b,:,...', canvas[b,:,ymin:ymax,xmin:xmax].shape)
            # print('crop_b', crop_b.shape)

            canvas[b, :, ymin:ymax, xmin:xmax] = crop_b
    return canvas

utils/samp.py
aharley/pips2
import numpy as np
import torch
import torch.nn.functional as F

import utils.geom

class Vox_util(object):
    def __init__(self, Z, Y, X, scene_centroid, bounds, pad=None, assert_cube=False):
        self.XMIN, self.XMAX, self.YMIN, self.YMAX, self.ZMIN, self.ZMAX = bounds
        B, D = list(scene_centroid.shape)
        self.Z, self.Y, self.X = Z, Y, X

        scene_centroid = scene_centroid.detach().cpu().numpy()
        x_centroid, y_centroid, z_centroid = scene_centroid[0]
        self.XMIN += x_centroid
        self.XMAX += x_centroid
        self.YMIN += y_centroid
        self.YMAX += y_centroid
        self.ZMIN += z_centroid
        self.ZMAX += z_centroid

        self.default_vox_size_X = (self.XMAX-self.XMIN)/float(X)
        self.default_vox_size_Y = (self.YMAX-self.YMIN)/float(Y)
        self.default_vox_size_Z = (self.ZMAX-self.ZMIN)/float(Z)

        if pad:
            Z_pad, Y_pad, X_pad = pad
            self.ZMIN -= self.default_vox_size_Z * Z_pad
            self.ZMAX += self.default_vox_size_Z * Z_pad
            self.YMIN -= self.default_vox_size_Y * Y_pad
            self.YMAX += self.default_vox_size_Y * Y_pad
            self.XMIN -= self.default_vox_size_X * X_pad
            self.XMAX += self.default_vox_size_X * X_pad

        if assert_cube:
            # we assume cube voxels
            if (not np.isclose(self.default_vox_size_X, self.default_vox_size_Y)) or (not np.isclose(self.default_vox_size_X, self.default_vox_size_Z)):
                print('Z, Y, X', Z, Y, X)
                print('bounds for this iter:',
                      'X = %.2f to %.2f' % (self.XMIN, self.XMAX),
                      'Y = %.2f to %.2f' % (self.YMIN, self.YMAX),
                      'Z = %.2f to %.2f' % (self.ZMIN, self.ZMAX),
                )
                print('self.default_vox_size_X', self.default_vox_size_X)
                print('self.default_vox_size_Y', self.default_vox_size_Y)
                print('self.default_vox_size_Z', self.default_vox_size_Z)
            assert(np.isclose(self.default_vox_size_X, self.default_vox_size_Y))
            assert(np.isclose(self.default_vox_size_X, self.default_vox_size_Z))

    def Ref2Mem(self, xyz, Z, Y, X, assert_cube=False):
        # xyz is B x N x 3, in ref coordinates
        # transforms ref coordinates into mem coordinates
        B, N, C = list(xyz.shape)
        device = xyz.device
        assert(C==3)
        mem_T_ref = self.get_mem_T_ref(B, Z, Y, X, assert_cube=assert_cube, device=device)
        xyz = utils.geom.apply_4x4(mem_T_ref, xyz)
        return xyz

    def Mem2Ref(self, xyz_mem, Z, Y, X, assert_cube=False):
        # xyz is B x N x 3, in mem coordinates
        # transforms mem coordinates into ref coordinates
        B, N, C = list(xyz_mem.shape)
        ref_T_mem = self.get_ref_T_mem(B, Z, Y, X, assert_cube=assert_cube, device=xyz_mem.device)
        xyz_ref = utils.geom.apply_4x4(ref_T_mem, xyz_mem)
        return xyz_ref

    def get_mem_T_ref(self, B, Z, Y, X, assert_cube=False, device='cuda'):
        vox_size_X = (self.XMAX-self.XMIN)/float(X)
        vox_size_Y = (self.YMAX-self.YMIN)/float(Y)
        vox_size_Z = (self.ZMAX-self.ZMIN)/float(Z)

        if assert_cube:
            if (not np.isclose(vox_size_X, vox_size_Y)) or (not np.isclose(vox_size_X, vox_size_Z)):
                print('Z, Y, X', Z, Y, X)
                print('bounds for this iter:',
                      'X = %.2f to %.2f' % (self.XMIN, self.XMAX),
                      'Y = %.2f to %.2f' % (self.YMIN, self.YMAX),
                      'Z = %.2f to %.2f' % (self.ZMIN, self.ZMAX),
                )
                print('vox_size_X', vox_size_X)
                print('vox_size_Y', vox_size_Y)
                print('vox_size_Z', vox_size_Z)
            assert(np.isclose(vox_size_X, vox_size_Y))
            assert(np.isclose(vox_size_X, vox_size_Z))

        # translation
        # (this makes the left edge of the leftmost voxel correspond to XMIN)
        center_T_ref = utils.geom.eye_4x4(B, device=device)
        center_T_ref[:,0,3] = -self.XMIN-vox_size_X/2.0
        center_T_ref[:,1,3] = -self.YMIN-vox_size_Y/2.0
        center_T_ref[:,2,3] = -self.ZMIN-vox_size_Z/2.0

        # scaling
        # (this makes the right edge of the rightmost voxel correspond to XMAX)
        mem_T_center = utils.geom.eye_4x4(B, device=device)
        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 = utils.geom.matmul2(mem_T_center, center_T_ref)

        return mem_T_ref

    def get_ref_T_mem(self, B, Z, Y, X, assert_cube=False, device='cuda'):
        mem_T_ref = self.get_mem_T_ref(B, Z, Y, X, assert_cube=assert_cube, device=device)
        # note safe_inverse is inapplicable here,
        # since the transform is nonrigid
        ref_T_mem = mem_T_ref.inverse()
        return ref_T_mem

    def get_inbounds(self, xyz, Z, Y, X, already_mem=False, padding=0.0, assert_cube=False):
        # xyz is B x N x 3
        # padding should be 0 unless you are trying to account for some later cropping
        if not already_mem:
            xyz = self.Ref2Mem(xyz, Z, Y, X, assert_cube=assert_cube)

        x = xyz[:,:,0]
        y = xyz[:,:,1]
        z = xyz[:,:,2]

        x_valid = ((x-padding)>-0.5).byte() & ((x+padding)<float(X-0.5)).byte()
        y_valid = ((y-padding)>-0.5).byte() & ((y+padding)<float(Y-0.5)).byte()
        z_valid = ((z-padding)>-0.5).byte() & ((z+padding)<float(Z-0.5)).byte()
        nonzero = (~(z==0.0)).byte()

        inbounds = x_valid & y_valid & z_valid & nonzero
        return inbounds.bool()

    def voxelize_xyz(self, xyz_ref, Z, Y, X, already_mem=False, assert_cube=False, clean_eps=0):
        B, N, D = list(xyz_ref.shape)
        assert(D==3)
        if already_mem:
            xyz_mem = xyz_ref
        else:
            xyz_mem = self.Ref2Mem(xyz_ref, Z, Y, X, assert_cube=assert_cube)
            xyz_zero = self.Ref2Mem(xyz_ref[:,0:1]*0, Z, Y, X, assert_cube=assert_cube)
        vox = self.get_occupancy(xyz_mem, Z, Y, X, clean_eps=clean_eps, xyz_zero=xyz_zero)
        return vox

    def voxelize_xyz_and_feats(self, xyz_ref, feats, Z, Y, X, already_mem=False, assert_cube=False, clean_eps=0):
        B, N, D = list(xyz_ref.shape)
        B2, N2, D2 = list(feats.shape)
        assert(D==3)
        assert(B==B2)
        assert(N==N2)
        if already_mem:
            xyz_mem = xyz_ref
        else:
            xyz_mem = self.Ref2Mem(xyz_ref, Z, Y, X, assert_cube=assert_cube)
            xyz_zero = self.Ref2Mem(xyz_ref[:,0:1]*0, Z, Y, X, assert_cube=assert_cube)
        feats = self.get_feat_occupancy(xyz_mem, feats, Z, Y, X, clean_eps=clean_eps, xyz_zero=xyz_zero)
        return feats

    def get_occupancy(self, xyz, Z, Y, X, clean_eps=0, xyz_zero=None):
        # xyz is B x N x 3 and in mem coords
        # we want to fill a voxel tensor with 1's at these inds
        B, N, C = list(xyz.shape)
        assert(C==3)

        # these papers say simple 1/0 occupancy is ok:
        #  http://openaccess.thecvf.com/content_cvpr_2018/papers/Yang_PIXOR_Real-Time_3d_CVPR_2018_paper.pdf
        #  http://openaccess.thecvf.com/content_cvpr_2018/papers/Luo_Fast_and_Furious_CVPR_2018_paper.pdf
        # cont fusion says they do 8-neighbor interp
        # voxelnet does occupancy but with a bit of randomness in terms of the reflectance value i think

        inbounds = self.get_inbounds(xyz, Z, Y, X, already_mem=True)
        x, y, z = xyz[:,:,0], xyz[:,:,1], xyz[:,:,2]
        mask = torch.zeros_like(x)
        mask[inbounds] = 1.0

        if xyz_zero is not None:
            # only take points that are beyond a thresh of zero
            dist = torch.norm(xyz_zero-xyz, dim=2)
            mask[dist < 0.1] = 0

        if clean_eps > 0:
            # only take points that are already near centers
            xyz_round = torch.round(xyz) # B, N, 3
            dist = torch.norm(xyz_round - xyz, dim=2)
            mask[dist > clean_eps] = 0

        # set the invalid guys to zero
        # we then need to zero out 0,0,0
        # (this method seems a bit clumsy)
        x = x*mask
        y = y*mask
        z = z*mask

        x = torch.round(x)
        y = torch.round(y)
        z = torch.round(z)
        x = torch.clamp(x, 0, X-1).int()
        y = torch.clamp(y, 0, Y-1).int()
        z = torch.clamp(z, 0, Z-1).int()

        x = x.view(B*N)
        y = y.view(B*N)
        z = z.view(B*N)

        dim3 = X
        dim2 = X * Y
        dim1 = X * Y * Z

        base = torch.arange(0, B, dtype=torch.int32, device=xyz.device)*dim1
        base = torch.reshape(base, [B, 1]).repeat([1, N]).view(B*N)

        vox_inds = base + z * dim2 + y * dim3 + x
        voxels = torch.zeros(B*Z*Y*X, device=xyz.device).float()
        voxels[vox_inds.long()] = 1.0
        # zero out the singularity
        voxels[base.long()] = 0.0
        voxels = voxels.reshape(B, 1, Z, Y, X)
        # B x 1 x Z x Y x X
        return voxels

    def get_feat_occupancy(self, xyz, feat, Z, Y, X, clean_eps=0, xyz_zero=None):
        # xyz is B x N x 3 and in mem coords
        # feat is B x N x D
        # we want to fill a voxel tensor with 1's at these inds
        B, N, C = list(xyz.shape)
        B2, N2, D2 = list(feat.shape)
        assert(C==3)
        assert(B==B2)
        assert(N==N2)

        # these papers say simple 1/0 occupancy is ok:
        #  http://openaccess.thecvf.com/content_cvpr_2018/papers/Yang_PIXOR_Real-Time_3d_CVPR_2018_paper.pdf
        #  http://openaccess.thecvf.com/content_cvpr_2018/papers/Luo_Fast_and_Furious_CVPR_2018_paper.pdf
        # cont fusion says they do 8-neighbor interp
        # voxelnet does occupancy but with a bit of randomness in terms of the reflectance value i think

        inbounds = self.get_inbounds(xyz, Z, Y, X, already_mem=True)
        x, y, z = xyz[:,:,0], xyz[:,:,1], xyz[:,:,2]
        mask = torch.zeros_like(x)
        mask[inbounds] = 1.0

        if xyz_zero is not None:
            # only take points that are beyond a thresh of zero
            dist = torch.norm(xyz_zero-xyz, dim=2)
            mask[dist < 0.1] = 0

        if clean_eps > 0:
            # only take points that are already near centers
            xyz_round = torch.round(xyz) # B, N, 3
            dist = torch.norm(xyz_round - xyz, dim=2)
            mask[dist > clean_eps] = 0

        # set the invalid guys to zero
        # we then need to zero out 0,0,0
        # (this method seems a bit clumsy)
        x = x*mask # B, N
        y = y*mask
        z = z*mask
        feat = feat*mask.unsqueeze(-1) # B, N, D

        x = torch.round(x)
        y = torch.round(y)
        z = torch.round(z)
        x = torch.clamp(x, 0, X-1).int()
        y = torch.clamp(y, 0, Y-1).int()
        z = torch.clamp(z, 0, Z-1).int()

        # permute point orders
        perm = torch.randperm(N)
        x = x[:, perm]
        y = y[:, perm]
        z = z[:, perm]
        feat = feat[:, perm]

        x = x.view(B*N)
        y = y.view(B*N)
        z = z.view(B*N)
        feat = feat.view(B*N, -1)

        dim3 = X
        dim2 = X * Y
        dim1 = X * Y * Z

        base = torch.arange(0, B, dtype=torch.int32, device=xyz.device)*dim1
        base = torch.reshape(base, [B, 1]).repeat([1, N]).view(B*N)

        vox_inds = base + z * dim2 + y * dim3 + x
        feat_voxels = torch.zeros((B*Z*Y*X, D2), device=xyz.device).float()
        feat_voxels[vox_inds.long()] = feat
        # zero out the singularity
        feat_voxels[base.long()] = 0.0
        feat_voxels = feat_voxels.reshape(B, Z, Y, X, D2).permute(0, 4, 1, 2, 3)
        # B x C x Z x Y x X
        return feat_voxels

    def unproject_image_to_mem(self, rgb_camB, pixB_T_camA, camB_T_camA, Z, Y, X, assert_cube=False, xyz_camA=None):
        # rgb_camB is B x C x H x W
        # pixB_T_camA is B x 4 x 4

        # rgb lives in B pixel coords
        # we want everything in A memory coords

        # this puts each C-dim pixel in the rgb_camB
        # along a ray in the voxelgrid
        B, C, H, W = list(rgb_camB.shape)

        if xyz_camA is None:
            xyz_memA = utils.basic.gridcloud3d(B, Z, Y, X, norm=False, device=pixB_T_camA.device)
            xyz_camA = self.Mem2Ref(xyz_memA, Z, Y, X, assert_cube=assert_cube)

        xyz_camB = utils.geom.apply_4x4(camB_T_camA, xyz_camA)
        z = xyz_camB[:,:,2]

        xyz_pixB = utils.geom.apply_4x4(pixB_T_camA, xyz_camA)
        normalizer = torch.unsqueeze(xyz_pixB[:,:,2], 2)
        EPS=1e-6
        # z = xyz_pixB[:,:,2]
        xy_pixB = xyz_pixB[:,:,:2]/torch.clamp(normalizer, min=EPS)
        # this is B x N x 2
        # this is the (floating point) pixel coordinate of each voxel
        x, y = xy_pixB[:,:,0], xy_pixB[:,:,1]
        # these are B x N

        x_valid = (x>-0.5).bool() & (x<float(W-0.5)).bool()
        y_valid = (y>-0.5).bool() & (y<float(H-0.5)).bool()
        z_valid = (z>0.0).bool()
        valid_mem = (x_valid & y_valid & z_valid).reshape(B, 1, Z, Y, X).float()

        if (0):
            # handwritten version
            values = torch.zeros([B, C, Z*Y*X], dtype=torch.float32)
            for b in list(range(B)):
                values[b] = utils.samp.bilinear_sample_single(rgb_camB[b], x_pixB[b], y_pixB[b])
        else:
            # native pytorch version
            y_pixB, x_pixB = utils.basic.normalize_grid2d(y, x, H, W)
            # since we want a 3d output, we need 5d tensors
            z_pixB = torch.zeros_like(x)
            xyz_pixB = torch.stack([x_pixB, y_pixB, z_pixB], axis=2)
            rgb_camB = rgb_camB.unsqueeze(2)
            xyz_pixB = torch.reshape(xyz_pixB, [B, Z, Y, X, 3])
            values = F.grid_sample(rgb_camB, xyz_pixB, align_corners=False)

        values = torch.reshape(values, (B, C, Z, Y, X))
        values = values * valid_mem
        return values

    def warp_tiled_to_mem(self, rgb_tileB, pixB_T_camA, camB_T_camA, Z, Y, X, DMIN, DMAX, assert_cube=False):
        # rgb_tileB is B,C,D,H,W
        # pixB_T_camA is B,4,4
        # camB_T_camA is B,4,4

        # rgb_tileB lives in B pixel coords but it has been tiled across the Z dimension
        # we want everything in A memory coords

        # this resamples the so that each C-dim pixel in rgb_tilB
        # is put into its correct place in the voxelgrid
        # (using the pinhole camera model)
        
        B, C, D, H, W = list(rgb_tileB.shape)

        xyz_memA = utils.basic.gridcloud3d(B, Z, Y, X, norm=False, device=pixB_T_camA.device)

        xyz_camA = self.Mem2Ref(xyz_memA, Z, Y, X, assert_cube=assert_cube)

        xyz_camB = utils.geom.apply_4x4(camB_T_camA, xyz_camA)
        z_camB = xyz_camB[:,:,2]

        # rgb_tileB has depth=DMIN in tile 0, and depth=DMAX in tile D-1
        z_tileB = (D-1.0) * (z_camB-float(DMIN)) / float(DMAX-DMIN)

        xyz_pixB = utils.geom.apply_4x4(pixB_T_camA, xyz_camA)
        normalizer = torch.unsqueeze(xyz_pixB[:,:,2], 2)
        EPS=1e-6
        # z = xyz_pixB[:,:,2]
        xy_pixB = xyz_pixB[:,:,:2]/torch.clamp(normalizer, min=EPS)
        # this is B x N x 2
        # this is the (floating point) pixel coordinate of each voxel
        x, y = xy_pixB[:,:,0], xy_pixB[:,:,1]
        # these are B x N

        x_valid = (x>-0.5).bool() & (x<float(W-0.5)).bool()
        y_valid = (y>-0.5).bool() & (y<float(H-0.5)).bool()
        z_valid = (z_camB>0.0).bool()
        valid_mem = (x_valid & y_valid & z_valid).reshape(B, 1, Z, Y, X).float()

        z_tileB, y_pixB, x_pixB = utils.basic.normalize_grid3d(z_tileB, y, x, D, H, W)
        xyz_pixB = torch.stack([x_pixB, y_pixB, z_tileB], axis=2)
        xyz_pixB = torch.reshape(xyz_pixB, [B, Z, Y, X, 3])
        values = F.grid_sample(rgb_tileB, xyz_pixB, align_corners=False)

        values = torch.reshape(values, (B, C, Z, Y, X))
        values = values * valid_mem
        return values
    

    def apply_mem_T_ref_to_lrtlist(self, lrtlist_cam, Z, Y, X, assert_cube=False):
        # lrtlist is B x N x 19, in cam coordinates
        # transforms them into mem coordinates, including a scale change for the lengths
        B, N, C = list(lrtlist_cam.shape)
        assert(C==19)
        mem_T_cam = self.get_mem_T_ref(B, Z, Y, X, assert_cube=assert_cube, device=lrtlist_cam.device)

    def xyz2circles(self, xyz, radius, Z, Y, X, soft=True, already_mem=True, also_offset=False, grid=None):
        # xyz is B x N x 3
        # radius is B x N or broadcastably so
        # output is B x N x Z x Y x X
        B, N, D = list(xyz.shape)
        assert(D==3)
        if not already_mem:
            xyz = self.Ref2Mem(xyz, Z, Y, X)

        if grid is None:
            grid_z, grid_y, grid_x = utils.basic.meshgrid3d(B, Z, Y, X, stack=False, norm=False, device=xyz.device)
            # note the default stack is on -1
            grid = torch.stack([grid_x, grid_y, grid_z], dim=1)
            # this is B x 3 x Z x Y x X
            
        xyz = xyz.reshape(B, N, 3, 1, 1, 1)
        grid = grid.reshape(B, 1, 3, Z, Y, X)
        # this is B x N x Z x Y x X

        # round the xyzs, so that at least one value matches the grid perfectly,
        # and we get a value of 1 there (since exp(0)==1)
        xyz = xyz.round()

        if torch.is_tensor(radius):
            radius = radius.clamp(min=0.01)

        if soft:
            off = grid - xyz # B,N,3,Z,Y,X
            # interpret radius as sigma
            dist_grid = torch.sum(off**2, dim=2, keepdim=False)
            # this is B x N x Z x Y x X
            if torch.is_tensor(radius):
                radius = radius.reshape(B, N, 1, 1, 1)
            mask = torch.exp(-dist_grid/(2*radius*radius))
            # zero out near zero
            mask[mask < 0.001] = 0.0
            # h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
            # h[h < np.finfo(h.dtype).eps * h.max()] = 0
            # return h
            if also_offset:
                return mask, off
            else:
                return mask
        else:
            assert(False) # something is wrong with this. come back later to debug

            dist_grid = torch.norm(grid - xyz, dim=2, keepdim=False)
            # this is 0 at/near the xyz, and increases by 1 for each voxel away

            radius = radius.reshape(B, N, 1, 1, 1)

            within_radius_mask = (dist_grid < radius).float()
            within_radius_mask = torch.sum(within_radius_mask, dim=1, keepdim=True).clamp(0, 1)
            return within_radius_mask

    def xyz2circles_bev(self, xyz, radius, Z, Y, X, already_mem=True, also_offset=False):
        # xyz is B x N x 3
        # radius is B x N or broadcastably so
        # output is B x N x Z x Y x X
        B, N, D = list(xyz.shape)
        assert(D==3)
        if not already_mem:
            xyz = self.Ref2Mem(xyz, Z, Y, X)

        xz = torch.stack([xyz[:,:,0], xyz[:,:,2]], dim=2)

        grid_z, grid_x = utils.basic.meshgrid2d(B, Z, X, stack=False, norm=False, device=xyz.device)
        # note the default stack is on -1
        grid = torch.stack([grid_x, grid_z], dim=1)
        # this is B x 2 x Z x X

        xz = xz.reshape(B, N, 2, 1, 1)
        grid = grid.reshape(B, 1, 2, Z, X)
        # these are ready to broadcast to B x N x Z x X

        # round the points, so that at least one value matches the grid perfectly,
        # and we get a value of 1 there (since exp(0)==1)
        xz = xz.round()

        if torch.is_tensor(radius):
            radius = radius.clamp(min=0.01)

        off = grid - xz # B,N,2,Z,X
        # interpret radius as sigma
        dist_grid = torch.sum(off**2, dim=2, keepdim=False)
        # this is B x N x Z x X
        if torch.is_tensor(radius):
            radius = radius.reshape(B, N, 1, 1, 1)
        mask = torch.exp(-dist_grid/(2*radius*radius))
        # zero out near zero
        mask[mask < 0.001] = 0.0
        
        # add a Y dim
        mask = mask.unsqueeze(-2)
        off = off.unsqueeze(-2)
        # # B,N,2,Z,1,X
        
        if also_offset:
            return mask, off
        else:
            return mask
        

utils/vox.py
aharley/pips2