Official repository of "Towards Robust Mobile Digital-Twin Tracking via An RGBD-based Transformer Model and A Comprehensive Mobile Dataset"
import json
import os
import random
import sys
from audioop import avg
from collections import defaultdict

import numpy as np
import scipy.io as scio
import tqdm

sys.path.append("../")
from dttd.dataset import DTTDDataset
from ycb.dataset import YCBDataset

obj_list = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,17,18,19]

def readlines(fpath):
    if os.path.isfile(fpath):
        with open(fpath, 'r') as f:
            lines = f.read().splitlines()
            return lines
    else:
        print(f"[Error] got a wrong object classes path: {fpath}")
        return []
    
def DTTD_total_counter(root):
    scene_names = [name for name in os.listdir(os.path.join(root,"data")) if os.path.isdir(os.path.join(root,"data",name))]
    total_frame_count = 0
    total_annotation_count = 0
    total_scene_count = 0
    objects_occurance = defaultdict(int)
    for scene in scene_names:
        count_obj_flag = True
        frame_count = 0
        annotation_count = 0
        data_path = os.path.join(root, "data", scene, "data")
        file_names = os.listdir(data_path)
        for file in file_names:
            if file.endswith(".json"):
                frame_count += 1
                with open(os.path.join(data_path, file), "r") as f:
                    meta = json.load(f)
                    if "synthetic" not in scene:
                        annotation_count += len(meta["objects"])
                        if count_obj_flag:
                            for o in meta["objects"]:
                                objects_occurance[o] += 1
                            count_obj_flag = False
                    else:
                        annotation_count += len([obj for obj in meta["objects"] if obj in obj_list])
        print(f"Scene {scene} has {frame_count} frames and {annotation_count} annotations!")
        total_frame_count += frame_count
        total_annotation_count += annotation_count
        total_scene_count += 1
    print(f"Total frame number is {total_frame_count}, total annotation number is {total_annotation_count}")
    print(f"Total scene count is {total_scene_count}")
    print(objects_occurance)
        

def DTTD_analyzer(root, config_dir, sample_num = 500):
    dataset = DTTDDataset(root=root, mode="train", add_noise=False)
    # get objects
    classes = dataset.classes
    dist = {}
    dist[0] = []
    for cls in range(len(classes)):
        dist[cls+1] = []
    # get data list
    train_path = dataset.Trainlist
    test_path = dataset.Testlist
    data = readlines(train_path)+readlines(test_path)
    # analyse scene per frame
    prefix = "data"
    annotations_record = {}
    for dt in tqdm.tqdm(data):
        if dt.startswith('synthetic'):
            continue
        scene_name = dt.split("/")[0]
        with open(os.path.join(root, prefix, f"{dt}_meta.json"), "r") as f:
            meta = json.load(f) # dict_keys(['objects', 'object_poses', 'intrinsic', 'distortion'])
        if scene_name not in annotations_record:
            annotations_record[scene_name] = {}
            annotations_record[scene_name]["unique_object_count"] = [len(meta["objects"])]
        else:
            if len(meta["objects"]) not in annotations_record[scene_name]["unique_object_count"]:
                annotations_record[scene_name]["unique_object_count"].append(len(meta["objects"]))
                print(f'find inconsistent object counts scene: {dt}')
        objs = np.array(meta['objects']).flatten().astype(np.int32) # id of objects that appear in the image
        if len(objs) == 0:
            print(f"find zero object data: {dt}")
        for obj_idx in objs:
            R = np.array(meta['object_poses'][str(obj_idx)])[0:3, 0:3] # (3, 3)
            T = np.array(meta['object_poses'][str(obj_idx)])[0:3, 3:4].T # (1, 3)
            points = dataset.model_points[obj_idx]
            # sample points
            sample_list = list(range(len(points)))
            sample_list = sorted(random.sample(sample_list, min(len(points), sample_num)))
            points = points[sample_list, :]
            # calculate translation
            points = np.add(np.dot(points, R.T), T)
            avg_dist = np.mean(points, axis=0)[2] # get z
            dist[obj_idx].append(avg_dist)
            dist[0].append(avg_dist)
    # print avg distance:
    # print(annotations_record)
    print("Evaluation DTTD average distance")
    print(f"Total object average distance: {sum(dist[0])/max(len(dist[0]), 1)},  max: {max(dist[0])}, min: {min(dist[0])}")
    for cls in range(len(classes)):
        if len(dist[cls+1]) == 0:
            dist[cls+1] = [0]
        print(f"object {classes[cls]}: {sum(dist[cls+1])/len(dist[cls+1])},  max: {max(dist[cls+1])}, min: {min(dist[cls+1])}")
        

def YCB_analyzer(root, config_dir, sample_num = 500):
    dataset = YCBDataset(root=root, mode="train", add_noise=False, config_path=config_dir)
    # get objects
    classes = dataset.classes
    dist = {}
    dist[0] = []
    for cls in range(len(classes)):
        dist[cls+1] = []
    # get data list
    train_path = dataset.Trainlist
    test_path = dataset.Testlist
    data = dataset.real_data_dirs+readlines(test_path)
    
    max_v = 0
    max_dt = None
    max_obj = None
    # analyse scene per frame
    for dt in tqdm.tqdm(data):
        meta = scio.loadmat(os.path.join(root, f"{dt}-meta.mat"))
        objs = meta['cls_indexes'].flatten().astype(np.int32) # id of objects that appear in the image
        for i, obj_idx in enumerate(objs):
            R = np.array(meta['poses'][:, :, i][:, 0:3]) # (3, 3)
            T = np.array([meta['poses'][:, :, i][:, 3:4].flatten()]) # (1, 3)
            points = dataset.model_points[obj_idx]
            # sample points
            sample_list = list(range(len(points)))
            sample_list = sorted(random.sample(sample_list, min(len(points), sample_num)))
            points = points[sample_list, :]
            # calculate translation
            points = np.add(np.dot(points, R.T), T)
            avg_dist = np.mean(points, axis=0)[2] # get z
            if avg_dist > max_v:
                max_v = avg_dist
                max_dt = dt
                max_obj = obj_idx
            dist[obj_idx].append(avg_dist)
            dist[0].append(avg_dist)
    # print avg distance:
    print(max_v, max_dt, max_obj)
    print("Evaluation YCB average distance")
    print(f"Total object average distance: {sum(dist[0])/max(len(dist[0]), 1)},  max: {max(dist[0])}, min: {min(dist[0])}")
    for cls in range(len(classes)):
        if len(dist[cls+1]) == 0:
            dist[cls+1] = [0]
        print(f"object {classes[cls]}: {sum(dist[cls+1])/len(dist[cls+1])},  max: {max(dist[cls+1])}, min: {min(dist[cls+1])}")

if __name__ == "__main__":
    dttd_root = "./dttd_iphone/DTTD_IPhone_Dataset/root"
    dttd_config_dir = "./dttd/DTTD_IPhone_Dataset/root"
    DTTD_total_counter(dttd_root)
    #DTTD_analyzer(dttd_root, dttd_config_dir)
    
    # ycb_root = "./ycb/YCB_Video_Dataset"
    # ycb_config_dir = "./ycb/dataset_config"
    # YCB_analyzer(ycb_root, ycb_config_dir)
dataset/analyse.py
augcog/DTTD2
from cgi import test
import os
import cv2
import time
import random
import json
import numpy as np
import numpy.ma as ma
from PIL import Image

import torch
import torch.utils.data as data
import torchvision.transforms as transforms

from utils.image import get_discrete_width_bbox

# TODO: add synthetic noise ??

Borderlist = [-1] + list(range(40, 1320, 40))

class DTTDDataset(data.Dataset):
    def __init__(self, root, mode, refine=False, add_noise=True):
        self.Trainlist = os.path.join(root, 'train_data_list.txt')
        self.Testlist = os.path.join(root, 'test_data_list.txt')
        self.Classlist = os.path.join(root, 'classes.txt')
        self.mode = mode
        self.root = root
        self.refine = refine
        self.add_noise = add_noise
        self._init_config()
        self._init_data_info()
        self.model_points = self.load_points(self.root, self.classes) # {1:[[x1, y1, z1], [x2, y2, z2], ...], 2:[...], 3:[...], ...}
        self.model_frontv = self.load_frontv(self.root, self.classes) # # {1:[vx1, vy1, vz1], 2:[vx2, vy2, vz2], 3:[...], ...}

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

    def __getitem__(self, index):
        # load raw data
        img = Image.open(os.path.join(self.root, self.prefix, f"{self.all_data_dirs[index]}_color.jpg"))   # PIL, size: (1280, 720)
        depth = np.array(Image.open(os.path.join(self.root, self.prefix, f"{self.all_data_dirs[index]}_depth.png")), dtype=np.uint16) # shape: (720, 1280)
        label = np.array(Image.open(os.path.join(self.root, self.prefix, f"{self.all_data_dirs[index]}_label.png"))) # shape: (720, 1280)
       
        with open(os.path.join(self.root, self.prefix, f"{self.all_data_dirs[index]}_meta.json"), "r") as f:
            meta = json.load(f) # dict_keys(['objects', 'object_poses', 'intrinsic', 'distortion'])
        
        # color jittor (noise) for img
        if self.add_noise:
            img = self.trancolor(img)
        img = np.array(img) # shape: (720, 1280, 3)

        objs = np.array(meta['objects']).flatten().astype(np.int32) # id of objects that appear in the image
        
        if len(objs) == 0:
            print(self.all_data_dirs[index])
        # randomly choose one object and get its mask
        obj_indices = list(range(len(objs)))
        random.shuffle(obj_indices)
        for obj_idx in obj_indices:
            mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
            mask_label = ma.getmaskarray(ma.masked_equal(label, objs[obj_idx]))
            mask = mask_label * mask_depth # consider both label (where = objs[obj_idx]) and valid depth (where > 0)
            # cv2.imwrite(f"test_{objs[obj_idx]}.png", mask.astype(np.int32)*255) # for debug
            if len(mask.nonzero()[0]) > self.minimum_px_num:
                break

        # get object's bounding box
        img_h, img_w = label.shape
        rmin, rmax, cmin, cmax = get_discrete_width_bbox(mask_label, Borderlist, img_w, img_h)
        
        # set sample points (2D) on depth/point cloud
        sample = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0] # non-zero positions on flattened mask, 1-D array
        if len(sample) >= self.pt_num:
            sample = np.array(sorted(np.random.choice(sample, self.pt_num))) # randomly choose pt_num points (idx)
        elif len(sample) == 0:
            sample = np.pad(sample, (0, self.pt_num - len(sample)), 'constant')
        else:
            sample = np.pad(sample, (0, self.pt_num - len(sample)), 'wrap')
        
        # crop image, depth and xy map with bbox, take the sample points
        img_crop = np.transpose(img[:, :, :3], (2, 0, 1))[:, rmin:rmax, cmin:cmax]
        depth_crop = depth[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, )
        xmap_crop = self.xmap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store y for sample points
        ymap_crop = self.ymap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store x for sample points
        
        # set camera focus and center
        cam = np.array(meta['intrinsic'])
        cam_cx, cam_cy, cam_fx, cam_fy =  cam[0][2], cam[1][2], cam[0][0], cam[1][1]
        
        # get point cloud [[px, py, pz], ...]
        cam_scale = 1000 # uint16 * 1000
        pz = depth_crop / cam_scale
        px = (ymap_crop - cam_cx) * pz / cam_fx
        py = (xmap_crop - cam_cy) * pz / cam_fy
        point_cloud = np.concatenate((px, py, pz), axis=1) # (pt_num, 3) store XYZ point cloud value for sample points
        
        # get ground truth rotation and translation
        R_gt = np.array(meta['object_poses'][str(objs[obj_idx])])[0:3, 0:3] # (3, 3)
        T_gt = np.array(meta['object_poses'][str(objs[obj_idx])])[0:3, 3:4].T # (1, 3)
        
        # get sample points (3D) on model
        model_sample_list = list(range(len(self.model_points[objs[obj_idx]])))
        if self.refine:
            model_sample_list = sorted(random.sample(model_sample_list, self.pt_num_mesh_large))
        else:
            model_sample_list = sorted(random.sample(model_sample_list, self.pt_num_mesh_small))
        sampled_model_pt = np.array(self.model_points[objs[obj_idx]][model_sample_list, :]) # (pt_num_mesh*, 3) 

        # get model points in the world coordinate
        sampled_model_pt_world = np.add(np.dot(sampled_model_pt, R_gt.T), T_gt) # (pt_num_mesh*, 3)
        
        # add noise for pointcloud and model points
        if self.add_noise:
            add_noise_t = np.array([random.uniform(-self.noise_trans, self.noise_trans) for i in range(3)])
            point_cloud = np.add(point_cloud, add_noise_t)
            sampled_model_pt_world = np.add(sampled_model_pt_world, add_noise_t)

        # get model front vector in the world coordinate
        frontv = np.array(self.model_frontv[objs[obj_idx]]) # (3,)
        frontv_world = np.dot(frontv, R_gt.T) # (3,)
        
        return {"img": torch.from_numpy(img),
               "label": torch.from_numpy(label),
               "depth": torch.from_numpy(depth.astype(np.float32)),
               "point_cloud": torch.from_numpy(point_cloud.astype(np.float32)), \
               "sample": torch.LongTensor(sample.astype(np.int32)), \
               "img_crop": self.norm(torch.from_numpy(img_crop.astype(np.float32))), \
               "sampled_model_pt_world": torch.from_numpy(sampled_model_pt_world.astype(np.float32)), \
               "sampled_model_pt": torch.from_numpy(sampled_model_pt.astype(np.float32)), \
               "obj_id": torch.LongTensor([int(objs[obj_idx]) - 1]), \
               "R": torch.from_numpy(R_gt.astype(np.float32)), \
               "T": torch.from_numpy(T_gt.astype(np.float32)), \
               "frontv": torch.from_numpy(frontv.astype(np.float32)), \
               "frontv_world": torch.from_numpy(frontv_world.astype(np.float32))} 

    def _init_config(self):
        self.trancolor = transforms.ColorJitter(0.2, 0.2, 0.2, 0.05)
        self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        
        self.xmap = np.array([[j for i in range(1280)] for j in range(720)])
        self.ymap = np.array([[i for i in range(1280)] for j in range(720)])
        
        # self.noise_img_loc = 0.0
        # self.noise_img_scale = 7.0
        self.noise_trans = 0.02
        self.minimum_px_num = 50
        self.symmetry_obj_cls = []
        self.pt_num = 1000
        self.pt_num_mesh_small = 500
        self.pt_num_mesh_large = 2600
        self.num_obj = 20
        self.prefix = "data"
    
    def _init_data_info(self):
        # set data list, store in local path
        if self.mode == 'train':
            self.data_list_path = self.Trainlist
        elif self.mode == 'test':
            self.data_list_path = self.Testlist
        else:
            raise NotImplementedError
        self.class_path = self.Classlist
        
        # read data information from files
        with open(self.data_list_path, 'r') as f:
            self.all_data_dirs = f.read().splitlines()
            self.real_data_dirs = self.all_data_dirs # real-world data path list
            self.syn_data_dirs = [d for d in self.all_data_dirs if d.startswith('data_syn/')] # synthetic data path list
            
        with open(self.class_path, 'r') as f:
            self.classes = f.read().splitlines()
            
    def load_points(self, root, classes):
        ''' load points of each model from points.xyz and save as dict(key: class id, value: points array) '''
        model_points = {} # {1:[[x1, y1, z1], [x2, y2, z2], ...], 2:[...], 3:[...], ...}
        class_id = 1
        for cls in classes:
            cls_filepath = os.path.join(root, 'objects', cls, 'points.xyz')
            if os.path.isfile(cls_filepath):
                model_points[class_id] = np.loadtxt(cls_filepath) # (2621, 3) float
            else:
                print(f"[Warning] {cls} doesn't exist, load model points for {cls} failed")
            class_id += 1
        return model_points
    
    def load_frontv(self, root, classes):
        ''' load the front vector of each model from front.xyz and save as dict(key: class id, value: vector array) '''
        model_frontv = {} # {1:[vx1, vy1, vz1], 2:[vx2, vy2, vz2], 3:[...], ...}
        class_id = 1
        for cls in classes:
            cls_filepath = os.path.join(root, 'objects', cls, 'front.xyz')
            if os.path.isfile(cls_filepath):
                model_frontv[class_id] = np.loadtxt(cls_filepath) # (3, ) float
            else:
                print(f"[Warning] {cls} doesn't exist, load model front vector for {cls} failed")
            class_id += 1
        return model_frontv
    
    def get_object_num(self):
        return self.num_obj
    
    def get_sym_list(self):
        return self.symmetry_obj_cls
    
    def get_point_mesh(self):
        if self.refine:
            return self.pt_num_mesh_large
        else:
            return self.pt_num_mesh_small
        
    def get_sample_point_num(self):
        return self.pt_num

        


dataset/dttd/dataset.py
augcog/DTTD2
import json
import os
import random
import time
from cgi import test

import cv2
import numpy as np
import numpy.ma as ma
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
from PIL import Image

Borderlist = [-1] + list(range(40, 1960, 40))

class DTTDDataset(data.Dataset):
    def __init__(self, root, mode, refine=False, add_noise=True, config_path="/ssd/zixun/Documents/DigitalTwin-6DPose_cleaned/estimation/dataset/dttd_iphone/dataset_config"):
        self.Trainlist = os.path.join(config_path, 'train_data_list.txt')
        self.Testlist = os.path.join(config_path, 'test_data_list.txt')
        self.Classlist = os.path.join(config_path, 'objectids.csv')
        self.mode = mode
        self.root = root
        self.refine = refine
        self.add_noise = add_noise
        self._init_config()
        self._init_data_info()
        self.model_points = self.load_points(self.root, self.classes) # {1:[[x1, y1, z1], [x2, y2, z2], ...], 2:[...], 3:[...], ...}

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

    def __getitem__(self, index):
        # load raw data
        img = Image.open(os.path.join(self.root, self.prefix, f"{self.all_data_dirs[index]}_color.jpg")) 
        depth = np.array(Image.open(os.path.join(self.root, self.prefix, f"{self.all_data_dirs[index]}_depth.png")), dtype=np.uint16) 
        label = np.array(Image.open(os.path.join(self.root, self.prefix, f"{self.all_data_dirs[index]}_label.png"))) 
        with open(os.path.join(self.root, self.prefix, f"{self.all_data_dirs[index]}_meta.json"), "r") as f:
            # dict_keys(['objects', 'object_poses', 'intrinsic', 'distortion'])
            meta = json.load(f)
        
        # set camera focus and center
        cam = np.array(meta['intrinsic'])
        cam_cx, cam_cy, cam_fx, cam_fy =  cam[0][2], cam[1][2], cam[0][0], cam[1][1]
        
        # color jittor (noise) for img
        if self.add_noise:
            img = self.trancolor(img)
        img = np.array(img) # shape: (1440, 1920, 3)

        # Get id of objects that appear in the image
        objs = np.array(meta['objects']).flatten().astype(np.int32)
        
        if len(objs) == 0:
            print(self.all_data_dirs[index])
        # randomly choose one object and get its mask
        obj_indices = list(range(len(objs)))
        random.shuffle(obj_indices)
        for obj_idx in obj_indices:
            # consider both label (where = objs[obj_idx]) and valid depth (where > 0)
            mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
            mask_label = ma.getmaskarray(ma.masked_equal(label, objs[obj_idx]))
            mask = mask_label * mask_depth if self.use_labelmask else mask_depth
            if self.debug_mode: 
                cv2.imwrite(f"test_{objs[obj_idx]}.png", mask.astype(np.int32)*255)
            if len((mask_label * mask_depth).nonzero()[0]) > self.minimum_px_num:
                break

        # get ground truth rotation and translation
        R_gt = np.array(meta['object_poses'][str(objs[obj_idx])])[0:3, 0:3] # (3, 3)
        T_gt = np.array(meta['object_poses'][str(objs[obj_idx])])[0:3, 3:4].T # (1, 3)

        # get sample points (3D) on model
        model_sample_list = list(range(len(self.model_points[objs[obj_idx]])))
        if self.refine:
            model_sample_list = sorted(random.sample(model_sample_list, self.pt_num_mesh_large))
        else:
            model_sample_list = sorted(random.sample(model_sample_list, self.pt_num_mesh_small))
        sampled_model_pt = np.array(self.model_points[objs[obj_idx]][model_sample_list, :]) # (pt_num_mesh*, 3) 

        # get model points in the world coordinate
        sampled_model_pt_world = np.add(np.dot(sampled_model_pt, R_gt.T), T_gt) # (pt_num_mesh*, 3)
        
        # get object's bounding box
        img_h, img_w = label.shape
        rmin, rmax, cmin, cmax = get_discrete_width_bbox(mask_label, Borderlist, img_w, img_h)
        # set sample points (2D) on depth/point cloud
        sample2D = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0] # non-zero positions on flattened mask, 1-D array
        if len(sample2D) >= self.sample_2d_pt_num:
            sample2D = np.array(sorted(np.random.choice(sample2D, self.sample_2d_pt_num))) # randomly choose sample_2d_pt_num points (idx)
        elif len(sample2D) == 0:
            sample2D = np.pad(sample2D, (0, self.sample_2d_pt_num - len(sample2D)), 'constant')
        else:
            sample2D = np.pad(sample2D, (0, self.sample_2d_pt_num - len(sample2D)), 'wrap')
        
        # crop image, depth and xy map with bbox, take the sample points
        img_crop = np.transpose(img[:, :, :3], (2, 0, 1))[:, rmin:rmax, cmin:cmax]
        depth_crop = depth[rmin:rmax, cmin:cmax].flatten()[sample2D][:, np.newaxis].astype(np.float32) # (sample_2d_pt_num, )
        xmap_crop = self.xmap[rmin:rmax, cmin:cmax].flatten()[sample2D][:, np.newaxis].astype(np.float32) # (sample_2d_pt_num, ) store y for sample points
        ymap_crop = self.ymap[rmin:rmax, cmin:cmax].flatten()[sample2D][:, np.newaxis].astype(np.float32) # (sample_2d_pt_num, ) store x for sample points
        
        # get point cloud [[px, py, pz], ...]
        cam_scale = 1000 # uint16 * 1000
        pz = depth_crop / cam_scale
        px = (ymap_crop - cam_cx) * pz / cam_fx
        py = (xmap_crop - cam_cy) * pz / cam_fy
        point_cloud = np.concatenate((px, py, pz), axis=1) # (sample_2d_pt_num, 3) store XYZ point cloud value for sample points
        
        # add noise for pointcloud and model points
        if self.add_noise:
            add_noise_t = np.array([random.uniform(-self.noise_trans, self.noise_trans) for i in range(3)])
            point_cloud = np.add(point_cloud, add_noise_t)
            sampled_model_pt_world = np.add(sampled_model_pt_world, add_noise_t)
        
        return {"img": torch.from_numpy(img),
               "label": torch.from_numpy(label),
               "depth": torch.from_numpy(depth.astype(np.float32)),
               "point_cloud": torch.from_numpy(point_cloud.astype(np.float32)), \
               "sample_2d": torch.LongTensor(sample2D.astype(np.int32)), \
               "img_crop": self.norm(torch.from_numpy(img_crop.astype(np.float32))), \
               "sampled_model_pt_camera": torch.from_numpy(sampled_model_pt_world.astype(np.float32)), \
               "sampled_model_pt": torch.from_numpy(sampled_model_pt.astype(np.float32)), \
               "obj_id": torch.LongTensor([int(objs[obj_idx]) - 1]), \
               "R": torch.from_numpy(R_gt.astype(np.float32)), \
               "T": torch.from_numpy(T_gt.astype(np.float32)), \
               } 

    def _init_config(self):
        self.trancolor = transforms.ColorJitter(0.2, 0.2, 0.2, 0.05)
        self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        
        self.xmap = np.array([[j for i in range(1920)] for j in range(1440)])
        self.ymap = np.array([[i for i in range(1920)] for j in range(1440)])

        self.noise_trans = 0.02
        self.minimum_px_num = 50
        self.symmetry_obj_cls = []
        self.sample_2d_pt_num = 1000
        self.pt_num_mesh_small = 500
        self.pt_num_mesh_large = 2600
        self.num_obj = 20

        self.prefix = "data"
        self.use_labelmask = True
        self.debug_mode = False
    
    def _init_data_info(self):
        # set data list, store in local path
        if self.mode == 'train':
            self.data_list_path = self.Trainlist
        elif self.mode == 'test':
            self.data_list_path = self.Testlist
        else:
            raise NotImplementedError
        self.class_path = self.Classlist
        
        # read data information from files
        with open(self.data_list_path, 'r') as f:
            self.all_data_dirs = f.read().splitlines()
            self.real_data_dirs = [d for d in self.all_data_dirs if d.startswith('scene')] # real-world data path list
            self.syn_data_dirs = [d for d in self.all_data_dirs if d.startswith('synthetic/data')] # synthetic data path list
        with open(self.class_path, 'r') as f:
            import pandas as pd
            self.classes = pd.read_csv(self.class_path, index_col='id')
            
    def load_points(self, root, classes):
        ''' load points of each model from points.xyz and save as dict(key: class id, value: points array) '''
        model_points = {} # {1:[[x1, y1, z1], [x2, y2, z2], ...], 2:[...], 3:[...], ...}
        for idx, cls in classes.iterrows():
            cls_filepath = os.path.join(root, 'objects', cls['name'], 'points.xyz')
            if os.path.isfile(cls_filepath):
                model_points[idx] = np.loadtxt(cls_filepath) # (2621, 3) float
            else:
                print(f"[Warning] {cls} doesn't exist, load model points for {cls} failed")
        return model_points
    
    def get_object_num(self):
        return self.num_obj
    
    def get_sym_list(self):
        return self.symmetry_obj_cls
    
    def get_model_point_num(self):
        if self.refine:
            return self.pt_num_mesh_large
        else:
            return self.pt_num_mesh_small
        
    def get_2d_sample_num(self):
        return self.sample_2d_pt_num

        

######################## dataset utils #############################

def get_discrete_width_bbox(label, border_list, img_w, img_h):
    rows = np.any(label, axis=1)
    cols = np.any(label, axis=0)
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    rmax += 1
    cmax += 1
    r_b = border_list[binary_search(border_list, rmax - rmin)]
    c_b = border_list[binary_search(border_list, cmax - cmin)]
    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > img_h:
        delt = rmax - img_h
        rmax = img_h
        rmin -= delt
    if cmax > img_w:
        delt = cmax - img_w
        cmax = img_w
        cmin -= delt
    return rmin, rmax, cmin, cmax

def discretize_bbox(rmin, rmax, cmin, cmax, border_list, img_w, img_h):
    rmax += 1
    cmax += 1
    r_b = border_list[binary_search(border_list, rmax - rmin)]
    c_b = border_list[binary_search(border_list, cmax - cmin)]
    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > img_h:
        delt = rmax - img_h
        rmax = img_h
        rmin -= delt
    if cmax > img_w:
        delt = cmax - img_w
        cmax = img_w
        cmin -= delt
    return rmin, rmax, cmin, cmax
    
def binary_search(sorted_list, target):
    l = 0
    r = len(sorted_list)-1
    while l!=r:
        mid = (l+r)>>1
        if sorted_list[mid] > target:
            r = mid
        elif sorted_list[mid] < target:
            l = mid + 1
        else:
            return mid
    return l


if __name__ == "__main__":
    dataset = DTTDDataset(root="./DTTD_IPhone_Dataset/root", mode="test", config_path="./dataset_config")
    dt = dataset[2]
    print("img crop shape: ", dt["img_crop"].size())
    # print("xyzmap crop shape: ", dt["xyzmap_crop"].size())
dataset/dttd_iphone/dataset.py
augcog/DTTD2
import os
import cv2
import time
import json
import random
import numpy as np
import numpy.ma as ma
from PIL import Image
import scipy.io as scio

import torch
import torch.utils.data as data
import torchvision.transforms as transforms

Borderlist = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]

class YCBDataset(data.Dataset):
    def __init__(self, root, mode, add_noise=True, config_path="dataset/ycb/dataset_config"):
        self.config_path = config_path
        self.Trainlist = os.path.join(config_path, 'train_data_list.txt')
        self.Testlist = os.path.join(config_path, 'test_data_list.txt')
        self.Classlist = os.path.join(config_path, 'classes.txt')
        self.mode = mode
        self.root = root
        self.add_noise = add_noise
        self._init_config()
        self._init_data_info()
        self.model_points = self.load_points(self.root, self.classes) # {1:[[x1, y1, z1], [x2, y2, z2], ...], 2:[...], 3:[...], ...}
        self.prim_groups = self.load_primitives()

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

    def __getitem__(self, index):
        # load raw data
        img = Image.open(os.path.join(self.root, f"{self.all_data_dirs[index]}-color.png"))   # PIL, size: (640, 480)
        depth = np.array(Image.open(os.path.join(self.root, f"{self.all_data_dirs[index]}-depth.png"))) # shape: (480, 640)
        label = np.array(Image.open(os.path.join(self.root, f"{self.all_data_dirs[index]}-label.png"))) # shape: (480, 640)
        meta = scio.loadmat(os.path.join(self.root, f"{self.all_data_dirs[index]}-meta.mat"))
        
        # set camera focus and center
        cam = np.array(meta['intrinsic_matrix'])
        cam_scale = meta["factor_depth"][0][0]
        cam_cx, cam_cy, cam_fx, cam_fy =  cam[0][2], cam[1][2], cam[0][0], cam[1][1]
        
        # color jittor (noise) for img
        if self.add_noise:
            img = self.trancolor(img)
        img = np.array(img) # shape: (480, 640, 3)

        # Get id of objects that appear in the image
        objs = meta['cls_indexes'].flatten().astype(np.int32)
        
        # randomly choose one object and get its mask
        obj_indices = list(range(len(objs)))
        random.shuffle(obj_indices)
        for obj_idx in obj_indices:
            mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
            mask_label = ma.getmaskarray(ma.masked_equal(label, objs[obj_idx]))
            mask = mask_depth # consider both label (where = objs[obj_idx]) and valid depth (where > 0)
            if len((mask_label * mask_depth).nonzero()[0]) > self.minimum_px_num:
                break
        
        # get ground truth rotation and translation
        R_gt = np.array(meta['poses'][:, :, obj_idx][:, 0:3]) # (3, 3)
        T_gt = np.array([meta['poses'][:, :, obj_idx][:, 3:4].flatten()]) # (1, 3)
        
        # get sample points (3D) on model
        model_sample_list = list(range(len(self.model_points[objs[obj_idx]])))
        model_sample_list = sorted(random.sample(model_sample_list, self.sample_model_pt_num))
        sampled_model_pt = np.array(self.model_points[objs[obj_idx]][model_sample_list, :]) # (sample_model_pt_num, 3) 

        # get model points in the camera coordinate
        sampled_model_pt_camera = np.add(np.dot(sampled_model_pt, R_gt.T), T_gt) # (sample_model_pt_num, 3)
        
        # projection and get bbox
        proj_x = sampled_model_pt_camera[:, 0] * cam_fx / sampled_model_pt_camera[:, 2] + cam_cx
        proj_y = sampled_model_pt_camera[:, 1] * cam_fy / sampled_model_pt_camera[:, 2] + cam_cy
        cmin, cmax = min(proj_x), max(proj_x)
        rmin, rmax = min(proj_y), max(proj_y)
        img_h, img_w = label.shape
        rmin, rmax, cmin, cmax = discretize_bbox(rmin, rmax, cmin, cmax, Borderlist, img_w, img_h)
        
        # set sample points (2D) on depth/point cloud
        sample2D = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0] # non-zero positions on flattened mask, 1-D array
        if len(sample2D) >= self.sample_2d_pt_num:
            sample2D = np.array(sorted(np.random.choice(sample2D, self.sample_2d_pt_num))) # randomly choose pt_num points (idx)
        elif len(sample2D) == 0:
            sample2D = np.pad(sample2D, (0, self.sample_2d_pt_num - len(sample2D)), 'constant')
        else:
            sample2D = np.pad(sample2D, (0, self.sample_2d_pt_num - len(sample2D)), 'wrap')
            
        depth_crop = depth[rmin:rmax, cmin:cmax].flatten()[sample2D][:, np.newaxis].astype(np.float32) # (pt_num, )
        xmap_crop = self.xmap[rmin:rmax, cmin:cmax].flatten()[sample2D][:, np.newaxis].astype(np.float32) # (pt_num, ) store y for sample points
        ymap_crop = self.ymap[rmin:rmax, cmin:cmax].flatten()[sample2D][:, np.newaxis].astype(np.float32) # (pt_num, ) store x for sample points
        
        # get point cloud [[px, py, pz], ...]
        cam_scale = meta["factor_depth"][0][0]
        pz = depth_crop / cam_scale
        px = (xmap_crop - cam_cx) * pz / cam_fx
        py = (ymap_crop - cam_cy) * pz / cam_fy
        point_cloud = np.concatenate((px, py, pz), axis=1) # (pt_num, 3) store XYZ point cloud value for sample points
        
        # get rgb image
        img_crop = np.transpose(img[:, :, :3], (2, 0, 1))[:, rmin:rmax, cmin:cmax] # shape (3, H, W)

        return {"img": torch.from_numpy(img),
               "label": torch.from_numpy(label),
               "depth": torch.from_numpy(depth.astype(np.float32)),
               "img_crop": self.norm(torch.from_numpy(img_crop.astype(np.float32))), \
               "point_cloud": torch.from_numpy(point_cloud.astype(np.float32)), \
               "sample_2d": torch.LongTensor(sample2D.astype(np.int32)), \
               "sampled_model_pt_camera": torch.from_numpy(sampled_model_pt_camera.astype(np.float32)), \
               "sampled_model_pt": torch.from_numpy(sampled_model_pt.astype(np.float32)), \
               "obj_id": torch.LongTensor([int(objs[obj_idx]) - 1]), \
               "R": torch.from_numpy(R_gt.astype(np.float32)), \
               "T": torch.from_numpy(T_gt.astype(np.float32))} 

    def _init_config(self):
        self.trancolor = transforms.ColorJitter(0.2, 0.2, 0.2, 0.05)
        self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        
        self.xmap = np.array([[i for i in range(640)] for j in range(480)])
        self.ymap = np.array([[j for i in range(640)] for j in range(480)])
        
        self.noise_trans = 0.02
        self.minimum_px_num = 50
        self.symmetry_obj_cls = [12, 15, 18, 19, 20]
        self.sample_model_pt_num = 500
        self.sample_2d_pt_num = 1000 # pt_num
        self.num_obj = 21
    
    def _init_data_info(self):
        # set data list, store in local path
        if self.mode == 'train':
            self.data_list_path = self.Trainlist
        elif self.mode == 'test':
            self.data_list_path = self.Testlist
        else:
            raise NotImplementedError
        self.class_path = self.Classlist
        
        # read data information from files
        with open(self.data_list_path, 'r') as f:
            self.all_data_dirs = f.read().splitlines()
            self.real_data_dirs = [d for d in self.all_data_dirs if d.startswith('data/')] # real-world data path list
            self.syn_data_dirs = [d for d in self.all_data_dirs if d.startswith('data_syn/')] # synthetic data path list
            
        with open(self.class_path, 'r') as f:
            self.classes = f.read().splitlines()
            
    def load_points(self, root, classes):
        ''' load points of each model from points.xyz and save as dict(key: class id, value: points array) '''
        model_points = {} # {1:[[x1, y1, z1], [x2, y2, z2], ...], 2:[...], 3:[...], ...}
        class_id = 1
        for cls in classes:
            cls_filepath = os.path.join(root, 'models', cls, 'points.xyz')
            if os.path.isfile(cls_filepath):
                model_points[class_id] = np.loadtxt(cls_filepath) # (2621, 3) float
            else:
                print(f"[Warning] {cls} doesn't exist, load model points for {cls} failed")
            class_id += 1
        return model_points
    
    def load_primitives(self):
        obj_radius = [0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15]
        primitives = []
        with open(os.path.join(self.config_path, "ycb_gp.json"), 'r') as f:
            prim_groups = json.load(f)
            for i, prim in enumerate(prim_groups):
                tmp = []
                for grp in prim['groups']:
                    tmp.append(torch.tensor(grp, dtype=torch.float).permute(1, 0).contiguous() / obj_radius[i])
                primitives.append(tmp)
        return primitives
    
    def get_object_num(self):
        return self.num_obj
    
    def get_sym_list(self):
        return self.symmetry_obj_cls
    
    def get_model_point_num(self):
        return self.sample_model_pt_num
    
    def get_2d_sample_num(self):
        return self.sample_2d_pt_num
    
    def get_models_xyz(self):
        return self.model_points


######################## dataset utils #############################

def get_discrete_width_bbox(label, border_list, img_w, img_h):
    rows = np.any(label, axis=1)
    cols = np.any(label, axis=0)
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    rmax += 1
    cmax += 1
    r_b = border_list[binary_search(border_list, rmax - rmin)]
    c_b = border_list[binary_search(border_list, cmax - cmin)]
    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > img_h:
        delt = rmax - img_h
        rmax = img_h
        rmin -= delt
    if cmax > img_w:
        delt = cmax - img_w
        cmax = img_w
        cmin -= delt
    return rmin, rmax, cmin, cmax

def discretize_bbox(rmin, rmax, cmin, cmax, border_list, img_w, img_h):
    rmax += 1
    cmax += 1
    r_b = border_list[binary_search(border_list, rmax - rmin)]
    c_b = border_list[binary_search(border_list, cmax - cmin)]
    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > img_h:
        delt = rmax - img_h
        rmax = img_h
        rmin -= delt
    if cmax > img_w:
        delt = cmax - img_w
        cmax = img_w
        cmin -= delt
    return rmin, rmax, cmin, cmax
    
def binary_search(sorted_list, target):
    l = 0
    r = len(sorted_list)-1
    while l!=r:
        mid = (l+r)>>1
        if sorted_list[mid] > target:
            r = mid
        elif sorted_list[mid] < target:
            l = mid + 1
        else:
            return mid
    return l
    
    
if __name__ == "__main__":
    dataset = YCBDataset(root="./YCB_Video_Dataset", mode="train", config_path="./dataset_config")
    dt = dataset[0]
    print("img crop shape: ", dt["img_crop"].size())
    print("xyzmap crop shape: ", dt["xyzmap_crop"].size())
dataset/ycb/dataset.py
augcog/DTTD2
import os
import sys
import cv2
import json
import argparse
from datetime import datetime
import numpy as np
import numpy.ma as ma
from PIL import Image
import scipy.io as scio
import torch
import torch.nn.parallel
import torch.utils.data
import torchvision.transforms as transforms
from torch.autograd import Variable

sys.path.append("../../")

from dataset.dttd.dataset import DTTDDataset, Borderlist
from model.posenet import PoseNet
from utils.log import Logger
from utils.file import get_checkpoint
from utils.visualizer import visualize
from utils.transformations import quaternion_matrix
from utils.image import get_discrete_width_bbox

def parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_root', type=str, default = '../../datasets/dttd/DTTD_Dataset', help='dataset root dir')
    parser.add_argument('--output', type=str, default = './eval_results', help='output directory to save results')
    parser.add_argument('--model', type=str, default = '',  help='path to resume model file')
    parser.add_argument('--result', type=str, default = 'eval_results',  help='Directory to save results')
    parser.add_argument('--visualize', action='store_true')
    return parser.parse_args()

def eval():
    opt = parse()
    # get meta parameters
    dataset = DTTDDataset(opt.dataset_root, "test", add_noise=False)
    # set dirs
    if not os.path.exists(opt.output):
        os.makedirs(opt.output)
        os.mkdir(os.path.join(opt.output, "mats"))
    # set visualize
    if opt.visualize:
        if not os.path.exists(os.path.join(opt.output, "visualize")):
            os.mkdir(os.path.join(opt.output, "visualize"))
        color_list = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255), (255, 255, 255), (123, 10, 265), (245, 163, 101), (100, 100, 178)]
    # set model
    estimator = PoseNet(num_points = dataset.pt_num, num_obj = dataset.num_obj) # pt_num:1000, num_obj: 20
    estimator.cuda()
    opt.model = get_checkpoint(opt.model)
    estimator.load_state_dict(torch.load(opt.model, map_location=torch.device('cuda')))
    estimator.eval()
    # set logger
    logger = Logger(os.path.join(opt.output, "log.txt"))
    logger.log("evaluation of DTTDDataset...")
    logger.log(f"eval date and time: {datetime.now().strftime('%d/%m/%Y %H:%M:%S')}")
    logger.log(f"mode: use groudtruth label")
    logger.log(f"dataset root: {opt.dataset_root}")
    logger.log(f"load model: {opt.model}")
    logger.log(f"output directory: {opt.output}")
    logger.log(f"if visualize: {opt.visualize}\n")
    # set data
    testlist = dataset.all_data_dirs
    logger.log(f"len of evaluation data list: {len(testlist)}")
    
    # start eval
    for i in range(len(testlist)):
        img = np.array(Image.open(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_color.jpg")))
        depth = np.array(Image.open(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_depth.png")), dtype=np.uint16)
        label = np.array(Image.open(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_label.png")))
        
        with open(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_meta.json"), "r") as f:
            meta = json.load(f) # dict_keys(['objects', 'object_poses', 'intrinsic', 'distortion'])
        
        lst = np.array(meta['objects']).flatten().astype(np.int32) # id of objects that appear in the image
        results = []
        
        for idx in range(len(lst)): # for each object in the scene
            itemid = lst[idx]
            try:
                data = process_data(itemid, img, depth, label, meta, Borderlist, num_points=dataset.pt_num, norm=dataset.norm)
                
                # put data to cuda
                cloud = Variable(data["cloud"]).cuda()
                sample = Variable(data["sample"]).cuda()
                img_crop = Variable(data["img_crop"]).cuda()
                index = Variable(data["index"]).cuda()
                
                # add batch dim
                cloud = cloud.view(1, dataset.pt_num, 3)
                img_crop = img_crop.view(1, 3, img_crop.size()[1], img_crop.size()[2])
                
                # predict
                pred_r, pred_t, pred_c, _ = estimator(img_crop, cloud, sample, index)
                
                pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, dataset.pt_num, 1)

                pred_c = pred_c.view(1, dataset.pt_num)
                pred_t = pred_t.view(1 * dataset.pt_num, 1, 3)
                points = cloud.view(1 * dataset.pt_num, 1, 3)
                
                how_max, which_max = torch.max(pred_c, 1)

                pred_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy()
                pred_t = (points + pred_t)[which_max[0]].view(-1).cpu().data.numpy()
                pred_concat = np.append(pred_r, pred_t)
                results.append(pred_concat.tolist())
            except Exception as e:
                logger.log(f"Detector Lost {itemid} at No.{i} keyframe. Error message: {e}")
                results.append([0.0 for i in range(7)])
        
        scio.savemat(os.path.join(opt.output, "mats", '%04d.mat' % i), {'poses':results})
        
        # visualize
        if opt.visualize:
            try:
                img = cv2.imread(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_color.jpg"))
                for idx in range(len(lst)):
                    itemid = lst[idx]
                    model_pts = np.array(dataset.model_points[itemid])
                    R, T = results[idx][0:4], results[idx][4:7]
                    R = quaternion_matrix(R)[0:3, 0:3]
                    T = np.array(T).reshape(1,3)
                    img = visualize(img=img, model_pts=model_pts, R=R, T=T, intrinsics=np.array(meta['intrinsic']), color=color_list[idx])
                cv2.imwrite(os.path.join(opt.output, "visualize", '%04d.png' % i), img)
            except Exception as e:
                logger.log(f"Visualization fail at No.{i} keyframe. Error message: {e}")
                
        logger.log("Finish No.{0} keyframe".format(i))
        
            
def process_data(itemid, img, depth, label, meta, border_list, num_points=1000, norm=None):
    img_h, img_w, _ = img.shape
    xmap = np.array([[j for i in range(img_w)] for j in range(img_h)])
    ymap = np.array([[i for i in range(img_w)] for j in range(img_h)])
    
    # mask valid depth (!= 0) and valid label (== itemid)
    mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
    mask_label = ma.getmaskarray(ma.masked_equal(label, itemid))
    mask = mask_label * mask_depth
    
    # get object's discretized bounding box
    rmin, rmax, cmin, cmax = get_discrete_width_bbox(mask_label, border_list, img_w, img_h)
    
    # set sample points (2D) on depth/point cloud
    sample = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0] # non-zero positions on flattened mask, 1-D array
    if len(sample) >= num_points:
        sample = np.array(sorted(np.random.choice(sample, num_points))) # randomly choose pt_num points (idx)
    elif len(sample) == 0:
        sample = np.pad(sample, (0, num_points - len(sample)), 'constant')
    else:
        sample = np.pad(sample, (0, num_points - len(sample)), 'wrap')
        
    # crop image, depth and xy map with bbox, take the sample points
    img_crop = np.transpose(img[:, :, :3], (2, 0, 1))[:, rmin:rmax, cmin:cmax]
    depth_crop = depth[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, )
    xmap_crop = xmap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store y for sample points
    ymap_crop = ymap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store x for sample points
    
    # set camera focus and center
    cam = np.array(meta['intrinsic'])
    cam_cx, cam_cy, cam_fx, cam_fy =  cam[0][2], cam[1][2], cam[0][0], cam[1][1]
    
    # get point cloud [[px, py, pz], ...]
    cam_scale = 1000 # uint16 * 1000
    pz = depth_crop / cam_scale
    px = (ymap_crop - cam_cx) * pz / cam_fx
    py = (xmap_crop - cam_cy) * pz / cam_fy
    point_cloud = np.concatenate((px, py, pz), axis=1) # (pt_num, 3) store XYZ point cloud value for sample points
    
    # trans to torch tensor
    if norm == None:
        norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    
    return {
        "cloud": torch.from_numpy(point_cloud.astype(np.float32)),
        "sample": torch.LongTensor(sample.astype(np.int32)),
        "img_crop": norm(torch.from_numpy(img_crop.astype(np.float32))),
        "index": torch.LongTensor([itemid - 1])
    }
    
if __name__ == "__main__":
    eval()
eval/dttd/eval_dttd_gt.py
augcog/DTTD2
import os
import sys
import cv2
import json
import argparse
from datetime import datetime
import numpy as np
import numpy.ma as ma
from PIL import Image
import scipy.io as scio
import torch
import torch.nn.parallel
import torch.utils.data
import torchvision.transforms as transforms
from torch.autograd import Variable

sys.path.append("../../")

from dataset.dttd_iphone.dataset import DTTDDataset, Borderlist
from model.posefusion import PoseNet
from utils.log import Logger
from utils.file import get_checkpoint
from utils.visualizer import visualize
from utils.transformations import quaternion_matrix
from utils.image import get_discrete_width_bbox

def parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_root', type=str, default = '../../dataset/dttd_iphone/DTTD_IPhone_Dataset/root', help='dataset root dir')
    parser.add_argument('--output', type=str, default = './eval_results', help='output directory to save results')
    parser.add_argument('--model', type=str, default = '',  help='path to resume model file')
    parser.add_argument('--result', type=str, default = 'eval_results',  help='Directory to save results')
    parser.add_argument('--visualize', action='store_true')
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--filter', action='store_true')
    parser.add_argument('--base_latent', type=int, default = 256, help='base latent dim for unimodal encoder')
    parser.add_argument('--embed_dim', type=int, default = 512, help='embedding dim for transformer encoder')
    parser.add_argument('--fusion_block_num', type=int, default = 1, help='number of fusion block')
    parser.add_argument('--layer_num_m', type=int, default = 4, help='layer num for modality fusion per block')
    parser.add_argument('--layer_num_p', type=int, default = 2, help='layer num for point-to-point fuison per block')
    return parser.parse_args()

def eval():
    opt = parse()
    # get meta parameters
    dataset = DTTDDataset(opt.dataset_root, "test", add_noise=False)
    # set dirs
    if not os.path.exists(opt.output):
        os.makedirs(opt.output)
        os.mkdir(os.path.join(opt.output, "mats"))
    # set visualize
    if opt.visualize:
        if not os.path.exists(os.path.join(opt.output, "visualize")):
            os.mkdir(os.path.join(opt.output, "visualize"))
        color_list = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255), (255, 255, 255), (123, 10, 265), (245, 163, 101), (100, 100, 178)]
    # set model
    estimator = PoseNet(num_points = dataset.sample_2d_pt_num, 
                        num_obj = dataset.num_obj,  # pt_num:1000, num_obj: 20 
                        base_latent=opt.base_latent, 
                        embedding_dim=opt.embed_dim, 
                        fusion_block_num=opt.fusion_block_num, 
                        layer_num_m=opt.layer_num_m, layer_num_p=opt.layer_num_p, 
                        filter_enhance=opt.filter) 
    estimator.cuda()
    opt.model = get_checkpoint(opt.model)
    estimator.load_state_dict(torch.load(opt.model, map_location=torch.device('cuda')), strict=True)
    estimator.eval()

    # set logger
    logger = Logger(os.path.join(opt.output, "log.txt"))
    logger.log("evaluation of DTTDDataset...")
    logger.log(f"eval date and time: {datetime.now().strftime('%d/%m/%Y %H:%M:%S')}")
    logger.log(f"mode: use groudtruth label")
    logger.log(f"dataset root: {opt.dataset_root}")
    logger.log(f"load model: {opt.model}")
    logger.log(f"output directory: {opt.output}")
    logger.log(f"if visualize: {opt.visualize}\n")
    # set data
    testlist = dataset.all_data_dirs
    logger.log(f"len of evaluation data list: {len(testlist)}")
    
    # start eval
    for i in range(len(testlist)):
        img = np.array(Image.open(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_color.jpg")))
        depth = np.array(Image.open(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_depth.png")), dtype=np.uint16)
        label = np.array(Image.open(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_label.png")))
        
        with open(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_meta.json"), "r") as f:
            meta = json.load(f) # dict_keys(['objects', 'object_poses', 'intrinsic', 'distortion'])
        
        lst = np.array(meta['objects']).flatten().astype(np.int32) # id of objects that appear in the image
        results = []
        
        for idx in range(len(lst)): # for each object in the scene
            itemid = lst[idx]
            try:
                data = process_data(itemid, img, depth, label, meta, Borderlist, num_points=dataset.sample_2d_pt_num, norm=dataset.norm)
                
                # put data to cuda
                cloud = Variable(data["cloud"]).cuda()
                sample = Variable(data["sample"]).cuda()
                img_crop = Variable(data["img_crop"]).cuda()
                index = Variable(data["index"]).cuda()
                
                # add batch dim
                cloud = cloud.view(1, dataset.sample_2d_pt_num, 3)
                img_crop = img_crop.view(1, 3, img_crop.size()[1], img_crop.size()[2])
                
                # predict
                pred_r, pred_t, pred_c, _, pt_recon, _ = estimator(img_crop, cloud, sample, index)

                # debug
                if opt.debug:
                    save_att(estimator, img_crop, cloud, sample, f'debug/{i}_{idx}_')
                    save_freq(estimator, cloud)

                pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, dataset.sample_2d_pt_num, 1)
                pred_c = pred_c.view(1, dataset.sample_2d_pt_num)
                pred_t = pred_t.view(1 * dataset.sample_2d_pt_num, 1, 3)
                points = cloud.view(1 * dataset.sample_2d_pt_num, 1, 3)
                
                how_max, which_max = torch.max(pred_c, 1)

                pred_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy()
                pred_t = (points + pred_t)[which_max[0]].view(-1).cpu().data.numpy()
                pred_concat = np.append(pred_r, pred_t)
                results.append(pred_concat.tolist())
            
            except Exception as e:
                logger.log(f"Detector Lost {itemid} at No.{i} keyframe. Error message: {e}")
                results.append([0.0 for i in range(7)])
        scio.savemat(os.path.join(opt.output, "mats", '%04d.mat' % i), {'poses':results})
        
        # visualize
        if opt.visualize:
            try:
                img = cv2.imread(os.path.join(opt.dataset_root, dataset.prefix, f"{testlist[i]}_color.jpg"))
                for idx in range(len(lst)):
                    itemid = lst[idx]
                    model_pts = np.array(dataset.model_points[itemid])
                    R, T = results[idx][0:4], results[idx][4:7]
                    R = quaternion_matrix(R)[0:3, 0:3]
                    T = np.array(T).reshape(1,3)
                    img = visualize(img=img, model_pts=model_pts, R=R, T=T, intrinsics=np.array(meta['intrinsic']), color=color_list[idx])
                cv2.imwrite(os.path.join(opt.output, "visualize", '%04d.png' % i), img)
            except Exception as e:
                logger.log(f"Visualization fail at No.{i} keyframe. Error message: {e}")
                
        logger.log("Finish No.{0} keyframe".format(i))
        
            
def process_data(itemid, img, depth, label, meta, border_list, num_points=1000, norm=None):
    img_h, img_w, _ = img.shape
    xmap = np.array([[j for i in range(img_w)] for j in range(img_h)])
    ymap = np.array([[i for i in range(img_w)] for j in range(img_h)])
    
    # mask valid depth (!= 0) and valid label (== itemid)
    mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
    mask_label = ma.getmaskarray(ma.masked_equal(label, itemid))
    mask = mask_label * mask_depth
    
    # get object's discretized bounding box
    rmin, rmax, cmin, cmax = get_discrete_width_bbox(mask_label, border_list, img_w, img_h)
    
    # set sample points (2D) on depth/point cloud
    sample = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0] # non-zero positions on flattened mask, 1-D array
    if len(sample) >= num_points:
        sample = np.array(sorted(np.random.choice(sample, num_points))) # randomly choose pt_num points (idx)
    elif len(sample) == 0:
        sample = np.pad(sample, (0, num_points - len(sample)), 'constant')
    else:
        sample = np.pad(sample, (0, num_points - len(sample)), 'wrap')
        
    # crop image, depth and xy map with bbox, take the sample points
    img_crop = np.transpose(img[:, :, :3], (2, 0, 1))[:, rmin:rmax, cmin:cmax]
    # Image.fromarray(np.transpose(img_crop, (1,2,0))).convert('RGB').save(f'rgb{itemid}.png')
    depth_crop = depth[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, )
    # Image.fromarray(depth[rmin:rmax, cmin:cmax]).convert('RGB').save(f'depth{itemid}.png')
    xmap_crop = xmap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store y for sample points
    ymap_crop = ymap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store x for sample points
    
    # set camera focus and center
    cam = np.array(meta['intrinsic'])
    cam_cx, cam_cy, cam_fx, cam_fy =  cam[0][2], cam[1][2], cam[0][0], cam[1][1]
    
    # get point cloud [[px, py, pz], ...]
    cam_scale = 1000 # uint16 * 1000
    pz = depth_crop / cam_scale
    px = (ymap_crop - cam_cx) * pz / cam_fx
    py = (xmap_crop - cam_cy) * pz / cam_fy
    point_cloud = np.concatenate((px, py, pz), axis=1) # (pt_num, 3) store XYZ point cloud value for sample points
    
    # trans to torch tensor
    if norm == None:
        norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    
    return {
        "cloud": torch.from_numpy(point_cloud.astype(np.float32)),
        "sample": torch.LongTensor(sample.astype(np.int32)),
        "img_crop": norm(torch.from_numpy(img_crop.astype(np.float32))),
        "index": torch.LongTensor([itemid - 1])
    }
    
def save_ptcld(xyz, fn):
    with open(fn, "w") as f:
        for i in range(len(xyz)):
            f.write(f"v {xyz[i][0]} {xyz[i][1]} {xyz[i][2]} \n")

def save_freq(estimator, x):
    estimator.get_freq_domain(x)

def save_att(estimator, img_crop, cloud, sample, fn_):
    estimator.eval()
    attn1, attn2 = estimator.get_attention_map(img_crop, cloud, sample)
    hn1 = attn1.shape[0]
    hn2 = attn2.shape[0]

    draw_attn_map(fn_+'m_avg.png', attn1)
    for i in range(hn1):
        fn = fn_+f'm_{i}'+'.png'
        draw_attn_map(fn, attn1, i)

    draw_attn_map(fn_+'p_avg.png', attn2)
    for i in range(hn2):
        fn = fn_+f'p_{i}'+'.png'
        draw_attn_map(fn, attn2, i)

def draw_attn_map(fn, attn, head_idx=None):
    if head_idx is not None:
        attn = attn[head_idx]
    else:
        attn = torch.mean(attn, dim=0)
    attn = np.asarray(attn.detach().cpu().numpy())
    attn = 255*(attn - np.min(attn))/(np.max(attn) - np.min(attn))
    image = Image.fromarray(attn)
    if image.mode == 'F':
        image = image.convert('RGB')
    image.save(fn)

if __name__ == "__main__":
    eval()
eval/dttd_iphone/eval.py
augcog/DTTD2
import os
import glob
import re
from moviepy.editor import *

base_dir = os.path.realpath("./dttd_iphone/eval_results_densefusion_gtmask/visualize")
# print(base_dir)

gif_name = 'pic'
fps = 30

# file_list = glob.glob(base_dir, '*.png')  # Get all the pngs in the current directory
# file_list_sorted = sorted(file_list,reverse=False)  # Sort the images

files = [os.path.join(base_dir, f) for f in sorted(os.listdir(base_dir))]
# print(files)
nf = []
for i in range(len(files)):
    res = re.match(r"(\d+)\.png", os.path.basename(files[i]))
    if res != None:
        nf.append((int(str(res.group(1))), files[i]))
    print(int(str(res.group(1))))

# nf = sorted(nf, key=lambda x: x[0])


clips = [ImageClip(m[1]).set_duration(1/fps) for m in nf]
concat_clip = concatenate_videoclips(clips)
concat_clip.write_videofile("eval.mp4", fps=fps)
eval/generate_video.py
augcog/DTTD2
import os
import sys
import cv2
import json
import argparse
from datetime import datetime
import numpy as np
import numpy.ma as ma
from PIL import Image
import scipy.io as scio
import torch
import torch.nn.parallel
import torch.utils.data
import torchvision.transforms as transforms
from torch.autograd import Variable

sys.path.append("../../")

from dataset.ycb.dataset import YCBDataset, Borderlist
from model.posenet import PoseNet
from utils.log import Logger
from utils.file import get_checkpoint
from utils.visualizer import visualize, Draw3DBBox
from utils.transformations import quaternion_matrix
from utils.image import get_discrete_width_bbox

def parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_root', type=str, default = '../../dataset/ycb/YCB_Video_Dataset', help='dataset root dir')
    parser.add_argument('--dataset_config', type=str, default = '../../dataset/ycb/dataset_config', help='dataset config dir')
    parser.add_argument('--ycb_toolbox', type=str, default = './YCB_Video_toolbox', help='dir to ycb video toolbox')
    parser.add_argument('--output', type=str, default = './eval_results', help='output directory to save results')
    parser.add_argument('--model', type=str, default = '',  help='path to resume model file')
    parser.add_argument('--visualize', action='store_true')
    return parser.parse_args()

def eval():
    opt = parse()
    # get meta parameters
    dataset = YCBDataset(opt.dataset_root, "test", add_noise=False, config_path=opt.dataset_config)
    # set dirs
    if not os.path.exists(opt.output):
        os.makedirs(opt.output)
        os.mkdir(os.path.join(opt.output, "mats"))
        os.mkdir(os.path.join(opt.output, "json"))
    # set visualize
    if opt.visualize:
        if not os.path.exists(os.path.join(opt.output, "visualize")):
            os.mkdir(os.path.join(opt.output, "visualize"))
        color_list = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255), (255, 255, 255), (123, 10, 265), (245, 163, 101), (100, 100, 178)]
    # set model
    estimator = PoseNet(num_points = dataset.get_2d_sample_num(), num_obj = dataset.get_object_num()) # pt_num:1000, num_obj: 21
    estimator.cuda()
    opt.model = get_checkpoint(opt.model)
    if not opt.model:
        print(f"[Error] Check model checkpoint path: {opt.model}, load model failed!")
        exit()
    estimator.load_state_dict(torch.load(opt.model, map_location=torch.device('cuda')))
    estimator.eval()
    # set logger
    logger = Logger(os.path.join(opt.output, "log.txt"))
    logger.log("evaluation of YCBDataset...")
    logger.log(f"eval date and time: {datetime.now().strftime('%d/%m/%Y %H:%M:%S')}")
    logger.log(f"mode: use groudtruth label")
    logger.log(f"dataset root: {opt.dataset_root}")
    logger.log(f"dataset config: {opt.dataset_config}")
    logger.log(f"YCB toolbox path: {opt.ycb_toolbox}")
    logger.log(f"load model: {opt.model}")
    logger.log(f"output directory: {opt.output}")
    logger.log(f"if visualize: {opt.visualize}\n")
    # set data
    testlist = dataset.all_data_dirs
    logger.log(f"len of evaluation data list: {len(testlist)}")
    
    # start eval
    for i in range(len(testlist)):
        img = np.array(Image.open(os.path.join(opt.dataset_root, f"{testlist[i]}-color.png")))
        depth = np.array(Image.open(os.path.join(opt.dataset_root, f"{testlist[i]}-depth.png")))
        label = np.array(Image.open(os.path.join(opt.dataset_root, f"{testlist[i]}-label.png")))
        meta = scio.loadmat(os.path.join(opt.dataset_root, f"{testlist[i]}-meta.mat"))
        
        lst = meta['cls_indexes'].flatten().astype(np.int32) # obj index list of the cuurent scene
        results = []
        
        for idx in range(len(lst)): # for each object in the frame
            itemid = lst[idx]
            try:
                data = process_data(itemid, img, depth, label, meta, Borderlist, num_points=dataset.get_2d_sample_num(), norm=dataset.norm)
                
                # put data to cuda
                cloud = Variable(data["cloud"]).cuda()
                sample = Variable(data["sample"]).cuda()
                img_crop = Variable(data["img_crop"]).cuda()
                index = Variable(data["index"]).cuda()
                
                # add batch dim
                cloud = cloud.view(1, dataset.get_2d_sample_num(), 3)
                img_crop = img_crop.view(1, 3, img_crop.size()[1], img_crop.size()[2])
                
                # predict
                pred_r, pred_t, pred_c, _ = estimator(img_crop, cloud, sample, index)
                
                pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, dataset.get_2d_sample_num(), 1)

                pred_c = pred_c.view(1, dataset.get_2d_sample_num())
                pred_t = pred_t.view(1 * dataset.get_2d_sample_num(), 1, 3)
                points = cloud.view(1 * dataset.get_2d_sample_num(), 1, 3)
                
                how_max, which_max = torch.max(pred_c, 1)

                pred_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy()
                pred_t = (points + pred_t)[which_max[0]].view(-1).cpu().data.numpy()
                pred_concat = np.append(pred_r, pred_t)
                results.append(pred_concat.tolist())
            except Exception as e:
                logger.log(f"Detector Lost {itemid} at No.{i} keyframe. Error message: {e}")
                results.append([0.0 for i in range(7)])
        
        scio.savemat(os.path.join(opt.output, "mats", '%04d.mat' % i), {'poses':results})
        
        with open(os.path.join(opt.output, "json", '%04d.json' % i), 'w') as f:
            json.dump({'poses':results}, f) 
        
        # visualize
        if opt.visualize:
            try:
                img = cv2.imread('{0}/{1}-color.png'.format(opt.dataset_root, testlist[i]))
                meta = scio.loadmat(('{0}/{1}-meta.mat'.format(opt.dataset_root, testlist[i])))
                for idx in range(len(lst)):
                    itemid = lst[idx]
                    model_pts = np.array(dataset.model_points[itemid])
                    R, T = results[idx][0:4], results[idx][4:7]
                    R = quaternion_matrix(R)[0:3, 0:3]
                    T = np.array(T).reshape(1,3)
                    img = visualize(img=img, model_pts=model_pts, R=R, T=T, intrinsics=np.array(meta['intrinsic_matrix']), color=color_list[idx])
                    img = Draw3DBBox(img=img, model_pts=model_pts, R=R, T=T, intrinsics=np.array(meta['intrinsic_matrix']), color=color_list[idx], thickness=1)
                cv2.imwrite(os.path.join(opt.output, "visualize", '%04d.png' % i), img)
            except Exception as e:
                logger.log(f"Visualization fail at No.{i} keyframe. Error message: {e}")
                
        logger.log("Finish No.{0} keyframe".format(i))
        
            
def process_data(itemid, img, depth, label, meta, border_list, num_points=1000, norm=None):
    img_h, img_w, _ = img.shape
    xmap = np.array([[j for i in range(img_w)] for j in range(img_h)])
    ymap = np.array([[i for i in range(img_w)] for j in range(img_h)])
    
    # mask valid depth (!= 0) and valid label (== itemid)
    mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
    mask_label = ma.getmaskarray(ma.masked_equal(label, itemid))
    # mask = mask_label * mask_depth
    mask = mask_depth
    
    # get object's discretized bounding box
    rmin, rmax, cmin, cmax = get_discrete_width_bbox(mask_label, border_list, img_w, img_h)
    
    # set sample points (2D) on depth/point cloud
    sample = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0] # non-zero positions on flattened mask, 1-D array
    if len(sample) >= num_points:
        sample = np.array(sorted(np.random.choice(sample, num_points))) # randomly choose pt_num points (idx)
    elif len(sample) == 0:
        sample = np.pad(sample, (0, num_points - len(sample)), 'constant')
    else:
        sample = np.pad(sample, (0, num_points - len(sample)), 'wrap')
        
    # crop image, depth and xy map with bbox, take the sample points
    img_crop = np.transpose(img[:, :, :3], (2, 0, 1))[:, rmin:rmax, cmin:cmax]
    depth_crop = depth[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, )
    xmap_crop = xmap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store y for sample points
    ymap_crop = ymap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store x for sample points
    
    # set camera focus and center
    cam = np.array(meta['intrinsic_matrix'])
    cam_cx, cam_cy, cam_fx, cam_fy =  cam[0][2], cam[1][2], cam[0][0], cam[1][1]
    
    # get point cloud [[px, py, pz], ...]
    cam_scale = meta["factor_depth"][0][0]
    pz = depth_crop / cam_scale
    px = (ymap_crop - cam_cx) * pz / cam_fx
    py = (xmap_crop - cam_cy) * pz / cam_fy
    point_cloud = np.concatenate((px, py, pz), axis=1) # (pt_num, 3) store XYZ point cloud value for sample points
    
    # trans to torch tensor
    if norm == None:
        norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    
    return {
        "cloud": torch.from_numpy(point_cloud.astype(np.float32)),
        "sample": torch.LongTensor(sample.astype(np.int32)),
        "img_crop": norm(torch.from_numpy(img_crop.astype(np.float32))),
        "index": torch.LongTensor([itemid - 1])
    }
    
if __name__ == "__main__":
    eval()
eval/ycb/eval_ycb_gt.py
augcog/DTTD2
import os
import sys
import cv2
import argparse
import numpy as np
from PIL import Image
import scipy.io as scio
import numpy.ma as ma
import torch
import torch.nn.parallel
import torch.utils.data
import torchvision.transforms as transforms
from torch.autograd import Variable

sys.path.append("../../")

from dataset.ycb.dataset import YCBDataset
from model.posenet import PoseNet
from utils.log import Logger
from utils.visualizer import visualize
from utils.transformations import quaternion_matrix
from utils.file import get_checkpoint

parser = argparse.ArgumentParser()
parser.add_argument('--dataset_root', type=str, default = '', help='dataset root dir')
parser.add_argument('--model', type=str, default = '',  help='resume PoseNet model')
parser.add_argument('--output', type=str, default = 'eval_results',  help='Directory to save results')
parser.add_argument('--visualize', action='store_true')
opt = parser.parse_args()

norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]
xmap = np.array([[j for i in range(640)] for j in range(480)])
ymap = np.array([[i for i in range(640)] for j in range(480)])
cam_cx = 312.9869
cam_cy = 241.3109
cam_fx = 1066.778
cam_fy = 1067.487
cam_scale = 10000.0
num_obj = 21
img_width = 480
img_length = 640
num_points = 1000
num_points_mesh = 500
bs = 1
dataset_config_dir = '../../dataset/ycb/dataset_config'
ycb_toolbox_dir = './YCB_Video_toolbox'
result_dir = opt.output
if not os.path.exists(result_dir):
    os.makedirs(result_dir)
    os.mkdir(os.path.join(result_dir, "mats"))
if opt.visualize:
    if not os.path.exists(os.path.join(result_dir, "visualize")):
        os.mkdir(os.path.join(result_dir, "visualize"))
    dataset = YCBDataset(opt.dataset_root, "test", add_noise=False, config_path=dataset_config_dir)
    color_list = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255), (255, 255, 255), (123, 10, 265), (245, 163, 101), (100, 100, 178)]

def get_bbox(posecnn_rois):
    rmin = int(posecnn_rois[idx][3]) + 1
    rmax = int(posecnn_rois[idx][5]) - 1
    cmin = int(posecnn_rois[idx][2]) + 1
    cmax = int(posecnn_rois[idx][4]) - 1
    r_b = rmax - rmin
    for tt in range(len(border_list)):
        if r_b > border_list[tt] and r_b < border_list[tt + 1]:
            r_b = border_list[tt + 1]
            break
    c_b = cmax - cmin
    for tt in range(len(border_list)):
        if c_b > border_list[tt] and c_b < border_list[tt + 1]:
            c_b = border_list[tt + 1]
            break
    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > img_width:
        delt = rmax - img_width
        rmax = img_width
        rmin -= delt
    if cmax > img_length:
        delt = cmax - img_length
        cmax = img_length
        cmin -= delt
    return rmin, rmax, cmin, cmax

estimator = PoseNet(num_points = num_points, num_obj = num_obj)
estimator.cuda()
opt.model = get_checkpoint(opt.model)
estimator.load_state_dict(torch.load(opt.model, map_location=torch.device('cuda')))
estimator.eval()

logger = Logger(os.path.join(result_dir, "log.txt"))
logger.log("evaluation of YCBDataset Begins...")
logger.log(f"dataset root: {opt.dataset_root}")
logger.log(f"load model: {opt.model}\n")

testlist = []
input_file = open('{0}/test_data_list.txt'.format(dataset_config_dir))
while 1:
    input_line = input_file.readline()
    if not input_line:
        break
    if input_line[-1:] == '\n':
        input_line = input_line[:-1]
    testlist.append(input_line)
input_file.close()

logger.log(f"len of evaluation data list: {len(testlist)}")

class_file = open('{0}/classes.txt'.format(dataset_config_dir))
class_id = 1
cld = {}
while 1:
    class_input = class_file.readline()
    if not class_input:
        break
    class_input = class_input[:-1]

    input_file = open('{0}/models/{1}/points.xyz'.format(opt.dataset_root, class_input))
    cld[class_id] = []
    while 1:
        input_line = input_file.readline()
        if not input_line:
            break
        input_line = input_line[:-1]
        input_line = input_line.split(' ')
        cld[class_id].append([float(input_line[0]), float(input_line[1]), float(input_line[2])])
    input_file.close()
    cld[class_id] = np.array(cld[class_id])
    class_id += 1

for now in range(min(len(testlist), 2949)):
    img = Image.open('{0}/{1}-color.png'.format(opt.dataset_root, testlist[now]))
    depth = np.array(Image.open('{0}/{1}-depth.png'.format(opt.dataset_root, testlist[now])))
    posecnn_meta = scio.loadmat('{0}/results_PoseCNN_RSS2018/{1}.mat'.format(ycb_toolbox_dir, '%06d' % now))
    label = np.array(posecnn_meta['labels'])
    posecnn_rois = np.array(posecnn_meta['rois'])
    
    lst = posecnn_rois[:, 1:2].flatten()
    results = []
    
    for idx in range(len(lst)):
        itemid = lst[idx]
        try:
            rmin, rmax, cmin, cmax = get_bbox(posecnn_rois)

            mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
            mask_label = ma.getmaskarray(ma.masked_equal(label, itemid))
            mask = mask_label * mask_depth

            choose = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0]
            if len(choose) > num_points:
                c_mask = np.zeros(len(choose), dtype=int)
                c_mask[:num_points] = 1
                np.random.shuffle(c_mask)
                choose = choose[c_mask.nonzero()]
            elif len(choose) == 0:
                choose = np.pad(choose, (0, num_points - len(choose)), 'constant')
            else:
                choose = np.pad(choose, (0, num_points - len(choose)), 'wrap')

            depth_masked = depth[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
            xmap_masked = xmap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
            ymap_masked = ymap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
            choose = np.array([choose])

            pt2 = depth_masked / cam_scale
            pt0 = (ymap_masked - cam_cx) * pt2 / cam_fx
            pt1 = (xmap_masked - cam_cy) * pt2 / cam_fy
            cloud = np.concatenate((pt0, pt1, pt2), axis=1)

            img_masked = np.array(img)[:, :, :3]
            img_masked = np.transpose(img_masked, (2, 0, 1))
            img_masked = img_masked[:, rmin:rmax, cmin:cmax]

            cloud = torch.from_numpy(cloud.astype(np.float32))
            choose = torch.LongTensor(choose.astype(np.int32))
            img_masked = norm(torch.from_numpy(img_masked.astype(np.float32)))
            index = torch.LongTensor([itemid - 1])

            cloud = Variable(cloud).cuda()
            choose = Variable(choose).cuda()
            img_masked = Variable(img_masked).cuda()
            index = Variable(index).cuda()

            cloud = cloud.view(1, num_points, 3)
            img_masked = img_masked.view(1, 3, img_masked.size()[1], img_masked.size()[2])

            pred_r, pred_t, pred_c, emb = estimator(img_masked, cloud, choose, index)
            pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, num_points, 1)

            pred_c = pred_c.view(bs, num_points)
            how_max, which_max = torch.max(pred_c, 1)
            pred_t = pred_t.view(bs * num_points, 1, 3)
            points = cloud.view(bs * num_points, 1, 3)

            my_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy()
            my_t = (points + pred_t)[which_max[0]].view(-1).cpu().data.numpy()
            my_pred = np.append(my_r, my_t)
            results.append(my_pred.tolist())

            # Here 'my_pred' is the final pose estimation result after refinement ('my_r': quaternion, 'my_t': translation)
        except ZeroDivisionError:
            print("PoseCNN Detector Lost {0} at No.{1} keyframe".format(itemid, now))
            results.append([0.0 for i in range(7)])

    scio.savemat(os.path.join(result_dir, "mats", '%04d.mat' % now), {'poses':results})
    
    # visualize
    if opt.visualize:
        img = cv2.imread('{0}/{1}-color.png'.format(opt.dataset_root, testlist[now]))
        meta = scio.loadmat(('{0}/{1}-meta.mat'.format(opt.dataset_root, testlist[now])))
        for idx in range(len(lst)):
            itemid = lst[idx]
            model_pts = np.array(dataset.model_points[itemid])
            R, T = results[idx][0:4], results[idx][4:7]
            R = quaternion_matrix(R)[0:3, 0:3]
            T = np.array(T).reshape(1,3)
            img = visualize(img=img, model_pts=model_pts, R=R, T=T, intrinsics=np.array(meta['intrinsic_matrix']), color=color_list[idx])
        cv2.imwrite(os.path.join(result_dir, "visualize", '%04d.png' % now), img)
    logger.log("Finish No.{0} keyframe".format(now))
eval/ycb/eval_ycb_posecnn.py
augcog/DTTD2
import os
import sys
import cv2
import json
import argparse
from datetime import datetime
import numpy as np
import numpy.ma as ma
from PIL import Image
import scipy.io as scio
import torch
import torch.nn.parallel
import torch.utils.data
import torchvision.transforms as transforms
from torch.autograd import Variable

sys.path.append("../../")

from dataset.ycb.dataset import YCBDataset, Borderlist
from model.poseformer import PoseNet

from _utils_.log import Logger
from _utils_.file import get_checkpoint
from _utils_.visualizer import visualize, Draw3DBBox
from _utils_.transformations import quaternion_matrix
from _utils_.image import get_discrete_width_bbox

def parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_root', type=str, default = '../../dataset/ycb/YCB_Video_Dataset', help='dataset root dir')
    parser.add_argument('--dataset_config', type=str, default = '../../dataset/ycb/dataset_config', help='dataset config dir')
    parser.add_argument('--ycb_toolbox', type=str, default = './YCB_Video_toolbox', help='dir to ycb video toolbox')
    parser.add_argument('--output', type=str, default = './eval_results', help='output directory to save results')
    parser.add_argument('--model', type=str, default = '',  help='path to resume model file')
    parser.add_argument('--visualize', action='store_true')
    parser.add_argument('--nano_test', action='store_true')

    return parser.parse_args()

def eval():
    opt = parse()
    # get meta parameters
    dataset = YCBDataset(opt.dataset_root, "test", add_noise=False, config_path=opt.dataset_config)
    # set dirs
    if not os.path.exists(opt.output):
        os.makedirs(opt.output)
        os.mkdir(os.path.join(opt.output, "mats"))
        os.mkdir(os.path.join(opt.output, "json"))
    # set visualize
    if opt.visualize:
        if not os.path.exists(os.path.join(opt.output, "visualize")):
            os.mkdir(os.path.join(opt.output, "visualize"))
        color_list = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255), (255, 255, 255), (123, 10, 265), (245, 163, 101), (100, 100, 178)]
    # set model
    estimator = PoseNet(num_points = dataset.get_2d_sample_num(), num_obj = dataset.get_object_num()) # pt_num:1000, num_obj: 21
    estimator.cuda()
    opt.model = get_checkpoint(opt.model)
    if not opt.model:
        print(f"[Error] Check model checkpoint path: {opt.model}, load model failed!")
        exit()
    estimator.load_state_dict(torch.load(opt.model, map_location=torch.device('cuda')))
    estimator.eval()
    # set logger
    logger = Logger(os.path.join(opt.output, "log.txt"))
    logger.log("evaluation of YCBDataset...")
    logger.log(f"eval date and time: {datetime.now().strftime('%d/%m/%Y %H:%M:%S')}")
    logger.log(f"mode: use groudtruth label")
    logger.log(f"dataset root: {opt.dataset_root}")
    logger.log(f"dataset config: {opt.dataset_config}")
    logger.log(f"YCB toolbox path: {opt.ycb_toolbox}")
    logger.log(f"load model: {opt.model}")
    logger.log(f"output directory: {opt.output}")
    logger.log(f"if visualize: {opt.visualize}\n")
    # set data
    testlist = dataset.all_data_dirs
    logger.log(f"len of evaluation data list: {len(testlist)}")
    
    # start eval
    for i in range(0,len(testlist)):
        img = np.array(Image.open(os.path.join(opt.dataset_root, f"{testlist[i]}-color.png")))
        depth = np.array(Image.open(os.path.join(opt.dataset_root, f"{testlist[i]}-depth.png")))
        label = np.array(Image.open(os.path.join(opt.dataset_root, f"{testlist[i]}-label.png")))
        meta = scio.loadmat(os.path.join(opt.dataset_root, f"{testlist[i]}-meta.mat"))
        
        lst = meta['cls_indexes'].flatten().astype(np.int32) # obj index list of the cuurent scene
        results = []
        
        for idx in range(len(lst)): # for each object in the frame
            itemid = lst[idx]
            try:
                data = process_data(itemid, img, depth, label, meta, Borderlist, num_points=dataset.get_2d_sample_num(), norm=dataset.norm)
                
                # put data to cuda
                cloud = Variable(data["cloud"]).cuda()
                sample = Variable(data["sample"]).cuda()
                img_crop = Variable(data["img_crop"]).cuda()
                index = Variable(data["index"]).cuda()
                
                # add batch dim
                cloud = cloud.view(1, dataset.get_2d_sample_num(), 3)
                img_crop = img_crop.view(1, 3, img_crop.size()[1], img_crop.size()[2])
                
                # predict
                pred_r, pred_t, pred_c, _, _, recon_loss = estimator(img_crop, cloud, sample, index)
                
                pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, dataset.get_2d_sample_num(), 1)

                pred_c = pred_c.view(1, dataset.get_2d_sample_num())
                pred_t = pred_t.view(1 * dataset.get_2d_sample_num(), 1, 3)
                points = cloud.view(1 * dataset.get_2d_sample_num(), 1, 3)
                
                how_max, which_max = torch.max(pred_c, 1)

                pred_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy()
                pred_t = (points + pred_t)[which_max[0]].view(-1).cpu().data.numpy()
                pred_concat = np.append(pred_r, pred_t)
                results.append(pred_concat.tolist())
            except Exception as e:
                logger.log(f"Detector Lost {itemid} at No.{i} keyframe. Error message: {e}")
                results.append([0.0 for i in range(7)])

        if opt.nano_test:
            break

        scio.savemat(os.path.join(opt.output, "mats", '%04d.mat' % i), {'poses':results})
        
        with open(os.path.join(opt.output, "json", '%04d.json' % i), 'w') as f:
            json.dump({'poses':results}, f) 
        
        
        # visualize
        if opt.visualize:
            try:
                img = cv2.imread('{0}/{1}-color.png'.format(opt.dataset_root, testlist[i]))
                meta = scio.loadmat(('{0}/{1}-meta.mat'.format(opt.dataset_root, testlist[i])))
                for idx in range(len(lst)):
                    itemid = lst[idx]
                    model_pts = np.array(dataset.model_points[itemid])
                    R, T = results[idx][0:4], results[idx][4:7]
                    R = quaternion_matrix(R)[0:3, 0:3]
                    T = np.array(T).reshape(1,3)
                    img = visualize(img=img, model_pts=model_pts, R=R, T=T, intrinsics=np.array(meta['intrinsic_matrix']), color=color_list[idx])
                    img = Draw3DBBox(img=img, model_pts=model_pts, R=R, T=T, intrinsics=np.array(meta['intrinsic_matrix']), color=color_list[idx], thickness=1)
                cv2.imwrite(os.path.join(opt.output, "visualize", '%04d.png' % i), img)
            except Exception as e:
                logger.log(f"Visualization fail at No.{i} keyframe. Error message: {e}")
                
        logger.log("Finish No.{0} keyframe".format(i))
        
            
def process_data(itemid, img, depth, label, meta, border_list, num_points=1000, norm=None):
    img_h, img_w, _ = img.shape
    xmap = np.array([[j for i in range(img_w)] for j in range(img_h)])
    ymap = np.array([[i for i in range(img_w)] for j in range(img_h)])
    
    # mask valid depth (!= 0) and valid label (== itemid)
    mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
    mask_label = ma.getmaskarray(ma.masked_equal(label, itemid))
    # mask = mask_label * mask_depth
    mask = mask_depth
    
    # get object's discretized bounding box
    rmin, rmax, cmin, cmax = get_discrete_width_bbox(mask_label, border_list, img_w, img_h)
    
    # set sample points (2D) on depth/point cloud
    sample = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0] # non-zero positions on flattened mask, 1-D array
    if len(sample) >= num_points:
        sample = np.array(sorted(np.random.choice(sample, num_points))) # randomly choose pt_num points (idx)
    elif len(sample) == 0:
        sample = np.pad(sample, (0, num_points - len(sample)), 'constant')
    else:
        sample = np.pad(sample, (0, num_points - len(sample)), 'wrap')
        
    # crop image, depth and xy map with bbox, take the sample points
    def normalize(im):
        """
        normalize image into floats
        :param: im: the img need to be normalized. shape[H, W, C]
        """
        return (im - np.min(im)) / (np.max(im) - np.min(im))
    img_crop = np.transpose(normalize(img[:, :, :3]), (2, 0, 1))[:, rmin:rmax, cmin:cmax]
    depth_crop = depth[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, )
    xmap_crop = xmap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store y for sample points
    ymap_crop = ymap[rmin:rmax, cmin:cmax].flatten()[sample][:, np.newaxis].astype(np.float32) # (pt_num, ) store x for sample points
    
    # set camera focus and center
    cam = np.array(meta['intrinsic_matrix'])
    cam_cx, cam_cy, cam_fx, cam_fy =  cam[0][2], cam[1][2], cam[0][0], cam[1][1]
    
    # get point cloud [[px, py, pz], ...]
    cam_scale = meta["factor_depth"][0][0]
    pz = depth_crop / cam_scale
    px = (ymap_crop - cam_cx) * pz / cam_fx
    py = (xmap_crop - cam_cy) * pz / cam_fy
    point_cloud = np.concatenate((px, py, pz), axis=1) # (pt_num, 3) store XYZ point cloud value for sample points
    
    # trans to torch tensor
    if norm == None:
        norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    
    return {
        "cloud": torch.from_numpy(point_cloud.astype(np.float32)),
        "sample": torch.LongTensor(sample.astype(np.int32)),
        "img_crop": norm(torch.from_numpy(img_crop.astype(np.float32))),
        "index": torch.LongTensor([itemid - 1])
    }
    
if __name__ == "__main__":
    eval()
eval/ycb/eval_ycb_poseformer.py
augcog/DTTD2
import torch
from torch.nn.modules.loss import _Loss

def loss(pred_r, pred_t, pred_c, model_target, model, point_cloud, w, num_point_mesh):
    # get batch size and number of sample points
    bs, num_p, _ = pred_c.size() # (1, 1000, 1)
    
    pred_r = pred_r / (torch.norm(pred_r, dim=2).view(bs, num_p, 1)) # (1, 1000, 4)
    
    base = torch.cat(((1.0 - 2.0*(pred_r[:, :, 2]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1),\
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] - 2.0*pred_r[:, :, 0]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 3]*pred_r[:, :, 0]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 2]**2)).view(bs, num_p, 1)), dim=2).contiguous().view(bs * num_p, 3, 3)
    
    # get predicted object points
    base = base.contiguous().transpose(2, 1).contiguous()
    model = model.view(bs, 1, num_point_mesh, 3).repeat(1, num_p, 1, 1).view(bs * num_p, num_point_mesh, 3)
    model_target = model_target.view(bs, 1, num_point_mesh, 3).repeat(1, num_p, 1, 1).view(bs * num_p, num_point_mesh, 3)
    pred_t = pred_t.contiguous().view(bs * num_p, 1, 3)
    point_cloud = point_cloud.contiguous().view(bs * num_p, 1, 3)
    pred = torch.add(torch.bmm(model, base), point_cloud + pred_t) # (bs*num_p, num_point_mesh, 3)
    
    # get distance and loss
    dis = torch.mean(torch.sum(torch.abs(pred - model_target), dim=2), dim=1)
    
    pred_c = pred_c.contiguous().view(bs * num_p)
    loss = torch.mean((dis * pred_c - w * torch.log(pred_c)), dim=0)
    
    pred_c = pred_c.view(bs, num_p)
    _, which_max = torch.max(pred_c, 1)
    dis = dis.view(bs, num_p)
    
    return loss, dis[0][which_max[0]]


class ABC_Loss(_Loss):

    def __init__(self, num_points_mesh):
        super(ABC_Loss, self).__init__(True)
        self.num_pt_mesh = num_points_mesh

    def forward(self, pred_r, pred_t, pred_c, target, model_points, points, w):
        return loss(pred_r, pred_t, pred_c, target, model_points, points, w, self.num_pt_mesh)
loss/abc.py
augcog/DTTD2
import torch
from torch.nn.modules.loss import _Loss

def loss(pred_r, pred_t, pred_c, model_target, model, point_cloud, w, num_point_mesh):
    # get batch size and number of sample points
    bs, num_p, _ = pred_c.size() # (1, 1000, 1)
    
    pred_r = pred_r / (torch.norm(pred_r, dim=2).view(bs, num_p, 1)) # (1, 1000, 4)
    
    base = torch.cat(((1.0 - 2.0*(pred_r[:, :, 2]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1),\
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] - 2.0*pred_r[:, :, 0]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 3]*pred_r[:, :, 0]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 2]**2)).view(bs, num_p, 1)), dim=2).contiguous().view(bs * num_p, 3, 3)
    
    # get predicted object points
    base = base.contiguous().transpose(2, 1).contiguous()
    model = model.view(bs, 1, num_point_mesh, 3).repeat(1, num_p, 1, 1).view(bs * num_p, num_point_mesh, 3)
    model_target = model_target.view(bs, 1, num_point_mesh, 3).repeat(1, num_p, 1, 1).view(bs * num_p, num_point_mesh, 3)
    pred_t = pred_t.contiguous().view(bs * num_p, 1, 3)
    point_cloud = point_cloud.contiguous().view(bs * num_p, 1, 3)
    pred = torch.add(torch.bmm(model, base), point_cloud + pred_t) # (bs*num_p, num_point_mesh, 3)
    
    # get distance and loss
    dis = torch.mean(torch.sum(torch.abs(pred - model_target), dim=2), dim=1)
    
    pred_c = pred_c.contiguous().view(bs * num_p)
    loss = torch.mean((dis * pred_c - w * torch.log(pred_c)), dim=0)
    
    pred_c = pred_c.view(bs, num_p)
    _, which_max = torch.max(pred_c, 1)
    dis = dis.view(bs, num_p)
    
    return loss, dis[0][which_max[0]]


class ABC_Loss(_Loss):

    def __init__(self, num_points_mesh):
        super(ABC_Loss, self).__init__(True)
        self.num_pt_mesh = num_points_mesh

    def forward(self, pred_r, pred_t, pred_c, target, model_points, points, w):
        return loss(pred_r, pred_t, pred_c, target, model_points, points, w, self.num_pt_mesh)
loss/add_l1.py
augcog/DTTD2
import torch
from torch.nn.modules.loss import _Loss
from knn_cuda import KNN

def loss(pred_r, pred_t, pred_c, model_target, model, idx, point_cloud, w, num_point_mesh, knn, sym_list):
    # get batch size and number of sample points
    bs, num_p, _ = pred_c.size() # (1, 1000, 1)
    
    pred_r = pred_r / (torch.norm(pred_r, dim=2).view(bs, num_p, 1)) # (1, 1000, 4)
    
    base = torch.cat(((1.0 - 2.0*(pred_r[:, :, 2]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1),\
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] - 2.0*pred_r[:, :, 0]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 3]*pred_r[:, :, 0]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 2]**2)).view(bs, num_p, 1)), dim=2).contiguous().view(bs * num_p, 3, 3)
    
    # get predicted object points
    base = base.contiguous().transpose(2, 1).contiguous()
    model = model.view(bs, 1, num_point_mesh, 3).repeat(1, num_p, 1, 1).view(bs * num_p, num_point_mesh, 3)
    model_target = model_target.view(bs, 1, num_point_mesh, 3).repeat(1, num_p, 1, 1).view(bs * num_p, num_point_mesh, 3)
    pred_t = pred_t.contiguous().view(bs * num_p, 1, 3)
    point_cloud = point_cloud.contiguous().view(bs * num_p, 1, 3)
    pred = torch.add(torch.bmm(model, base), point_cloud + pred_t) # (bs*num_p, num_point_mesh, 3)
    
    # if the current object is symmetrical, find the nearest point
    # TODO: batch size should be restricted to 1
    # TODO: Should run on CUDA environment, CPU is not supported
    if idx[0].item() in sym_list:
        model_target = model_target[0].contiguous().view(-1, 3)
        pred = pred.contiguous().view(-1, 3)
        _, inds = knn(model_target.unsqueeze(0), pred.unsqueeze(0))
        model_target = torch.index_select(model_target, 0, inds.view(-1))
        model_target = model_target.view(bs*num_p, num_point_mesh, 3)
        pred = pred.view(bs*num_p, num_point_mesh, 3)
    
    # get distance and loss
    dis = torch.mean(torch.norm((pred - model_target), dim=2), dim=1)
    
    pred_c = pred_c.contiguous().view(bs * num_p)
    loss = torch.mean((dis * pred_c - w * torch.log(pred_c)), dim=0)
    
    pred_c = pred_c.view(bs, num_p)
    _, which_max = torch.max(pred_c, 1)
    dis = dis.view(bs, num_p)
    
    return loss, dis[0][which_max[0]]


class ADDS_Loss(_Loss):

    def __init__(self, num_points_mesh, sym_list):
        super(ADDS_Loss, self).__init__(True)
        self.num_pt_mesh = num_points_mesh
        self.sym_list = sym_list
        self.knn = KNN(k=1, transpose_mode=True)

    def forward(self, pred_r, pred_t, pred_c, target, model_points, idx, points, w):
        return loss(pred_r, pred_t, pred_c, target, model_points, idx, points, w, self.num_pt_mesh, self.knn, self.sym_list)
loss/add_s.py
augcog/DTTD2
import torch
from torch.nn.modules.loss import _Loss

def loss(pred_r, pred_t, pred_c, frontv_target, frontv, point_cloud, w):
    # get batch size and number of sample points
    bs, num_p, _ = pred_c.size() # (1, 1000, 1)
    
    pred_r = pred_r / (torch.norm(pred_r, dim=2).view(bs, num_p, 1)) # (1, 1000, 4)
    
    base = torch.cat(((1.0 - 2.0*(pred_r[:, :, 2]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1),\
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] - 2.0*pred_r[:, :, 0]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 3]*pred_r[:, :, 0]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 2]**2)).view(bs, num_p, 1)), dim=2).contiguous().view(bs * num_p, 3, 3)
    
    # get predicted object points
    base = base.contiguous().transpose(2, 1).contiguous()
    frontv = frontv.view(bs, 1, 1, 3).repeat(1, num_p, 1, 1).view(bs * num_p, 1, 3)
    frontv_target = frontv_target.view(bs, 1, 1, 3).repeat(1, num_p, 1, 1).view(bs * num_p, 1, 3)
    pred_t = pred_t.contiguous().view(bs * num_p, 1, 3)
    point_cloud = point_cloud.contiguous().view(bs * num_p, 1, 3)
    pred = torch.bmm(frontv, base) # (bs*num_p, 1, 3)
    
    # get distance and loss
    dis = torch.mean(torch.norm((pred - frontv_target), dim=2), dim=1)
    
    pred_c = pred_c.contiguous().view(bs * num_p)
    loss = torch.mean((dis * pred_c - w * torch.log(pred_c)), dim=0)
    
    pred_c = pred_c.view(bs, num_p)
    _, which_max = torch.max(pred_c, 1)
    dis = dis.view(bs, num_p)
    
    return loss, dis[0][which_max[0]]


class FV_Loss(_Loss):

    def __init__(self):
        super(FV_Loss, self).__init__(True)

    def forward(self, pred_r, pred_t, pred_c, frontv_target, frontv, point_cloud, w):
        return loss(pred_r, pred_t, pred_c, frontv_target, frontv, point_cloud, w)
loss/frontv.py
augcog/DTTD2
import torch
from torch.nn.modules.loss import _Loss

def loss(pred_r, pred_t, pred_c, tar_r, tar_t, prim_groups, obj_idx, point_cloud, w, device=torch.device("cuda")):
    bs, num_p, _ = pred_c.size() # (bs, 1000, 1)
    pred_r = pred_r / (torch.norm(pred_r, dim=2).view(bs, num_p, 1)) # (bs, 1000, 4)
    pred_r = torch.cat(((1.0 - 2.0*(pred_r[:, :, 2]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1),\
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] - 2.0*pred_r[:, :, 0]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 3]*pred_r[:, :, 0]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 2]**2)).view(bs, num_p, 1)), dim=2).contiguous().view(bs, num_p, 3, 3) # (bs, 1000, 3, 3)
    pred_t = (pred_t+point_cloud).contiguous().view(bs, num_p, 1, 3)
    
    # get prim group
    sub_dist = torch.zeros(bs).to(device)
    sub_loss = torch.zeros(bs).to(device)
    for i in range(bs):
        obj_groups = prim_groups[obj_idx[i]]
        tar_r_i = tar_r[i].view(1, 3, 3).repeat(num_p, 1, 1).contiguous() # [num_p, 3, 3]
        tar_t_i = tar_t[i].view(1, 3, 1).repeat(num_p, 1, 1).contiguous() # [num_p, 3, 1]
        min_dists = []
        for j, grp in enumerate(obj_groups):
            _, pt_num = grp.size()
            grp = grp.to(device)
            grp = grp.view(1, 3, pt_num).repeat(num_p, 1, 1) # (num_p, 3, pt_num)
            pred = torch.bmm(pred_r[i], grp) + pred_t[i].view(num_p, 3, 1) # (num_p, 3, pt_num)
            tar = torch.bmm(tar_r_i, grp) + tar_t_i.view(num_p, 3, 1) # (num_p, 3, pt_num)
            pred = pred.unsqueeze(dim=1).repeat(1, pt_num, 1, 1).contiguous() # (num_p, pt_num, 3, pt_num)
            tar = tar.permute(0, 2, 1).unsqueeze(dim=3).repeat(1, 1, 1, pt_num).contiguous() # (num_p, pt_num, 3, pt_num)
            min_dist, _ = torch.min(torch.norm(pred-tar, dim=2), dim=2)
            
            if len(obj_groups) == 3 and j == 2:
                min_dist = torch.max(min_dist, dim=1)[0]  # [num_p]
            else:
                min_dist = torch.mean(min_dist, dim=1)  # [num_p]  
            min_dists.append(min_dist.unsqueeze(0))
            
        min_dists = torch.concat(min_dists, dim=0)
        if len(obj_groups) == 3 and obj_groups[2].size()[1] > 1:
            min_dists = torch.max(min_dists, dim=0)[0]  # [num_p]
        else:
            min_dists = torch.mean(min_dists, dim=0)  # [num_p]
        
        sub_dist[i] = torch.mean(min_dists[torch.max(pred_c[i].view(num_p), dim=0)[1]])
        sub_loss[i] = torch.mean(min_dists * pred_c[i].view(num_p) - w * torch.log(pred_c[i].view(num_p)))
        
    gadd = torch.mean(sub_dist)
    gadd_loss = torch.mean(sub_loss)
    
    gadd_loss = torch.where(torch.isinf(gadd_loss), torch.full_like(gadd_loss, 0), gadd_loss)
    gadd_loss = torch.where(torch.isnan(gadd_loss), torch.full_like(gadd_loss, 0), gadd_loss)
    
    return gadd_loss, gadd

class GADD_Loss(_Loss):

    def __init__(self, device):
        super(GADD_Loss, self).__init__(True)
        self.device = device

    def forward(self, pred_r, pred_t, pred_c, tar_r, tar_t, prim_groups, obj_idx, point_cloud, w):
        return loss(pred_r, pred_t, pred_c, tar_r, tar_t, prim_groups, obj_idx, point_cloud, w, self.device)
    
if __name__ == "__main__":
    import sys
    sys.path.append("..")
    from dataset.ycb.dataset import YCBDataset
    dataset = YCBDataset("../dataset/ycb/YCB_Video_Dataset/", "train", True, "../dataset/ycb/dataset_config")
    bs = 2
    pred_r = torch.rand((bs, 1000, 4))
    pred_t = torch.rand((bs, 1000, 3))
    pred_c = torch.rand((bs, 1000, 1))
    point_cloud = torch.rand((bs, 1000, 1))
    tar_r = torch.rand((bs, 3, 3))
    tar_t = torch.rand((bs, 1, 3))
    obj_idx = torch.zeros((bs,), dtype=torch.int32)
    l, _ = loss(pred_r, pred_t, pred_c, tar_r, tar_t, dataset.prim_groups, obj_idx, point_cloud, 0.15, torch.device("cpu"))
    print(l)
loss/gadd.py
augcog/DTTD2

import torch

from torch.utils.cpp_extension import load
cd = load(name="cd",
          sources=["/ssd/zixun/Documents/DigitalTwin-6DPose_cleaned/estimation/model/chamfer_distance/chamfer_distance_cpu.cpp"
                   #,"./model/chamfer_distance/chamfer_distance.cu"
                   ], verbose=True)

class ChamferDistanceFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, xyz1, xyz2):
        batchsize, n, _ = xyz1.size()
        _, m, _ = xyz2.size()
        xyz1 = xyz1.contiguous()
        xyz2 = xyz2.contiguous()
        dist1 = torch.zeros(batchsize, n)
        dist2 = torch.zeros(batchsize, m)

        idx1 = torch.zeros(batchsize, n, dtype=torch.int)
        idx2 = torch.zeros(batchsize, m, dtype=torch.int)

        if not xyz1.is_cuda:
            cd.forward(xyz1, xyz2, dist1, dist2, idx1, idx2)
        # else:
        #     dist1 = dist1.cuda()
        #     dist2 = dist2.cuda()
        #     idx1 = idx1.cuda()
        #     idx2 = idx2.cuda()
        #     cd.forward_cuda(xyz1, xyz2, dist1, dist2, idx1, idx2)

        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)

        return dist1, dist2

    @staticmethod
    def backward(ctx, graddist1, graddist2):
        xyz1, xyz2, idx1, idx2 = ctx.saved_tensors

        graddist1 = graddist1.contiguous()
        graddist2 = graddist2.contiguous()

        gradxyz1 = torch.zeros(xyz1.size())
        gradxyz2 = torch.zeros(xyz2.size())

        if not graddist1.is_cuda:
            cd.backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2)
        # else:
        #     gradxyz1 = gradxyz1.cuda()
        #     gradxyz2 = gradxyz2.cuda()
        #     cd.backward_cuda(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2)

        return gradxyz1, gradxyz2


class ChamferDistance(torch.nn.Module):
    def forward(self, xyz1, xyz2):
        return ChamferDistanceFunction.apply(xyz1, xyz2)

model/chamfer_distance/chamfer_distance_cpu.py
augcog/DTTD2

import torch

from torch.utils.cpp_extension import load
cd = load(name="cd",
          sources=["/ssd/zixun/Documents/DigitalTwin-6DPose_cleaned/estimation/model/chamfer_distance/chamfer_distance_gpu.cpp"
                   ,"/ssd/zixun/Documents/DigitalTwin-6DPose_cleaned/estimation/model/chamfer_distance/chamfer_distance.cu"
                   ])

class ChamferDistanceFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, xyz1, xyz2):
        batchsize, n, _ = xyz1.size()
        _, m, _ = xyz2.size()
        xyz1 = xyz1.contiguous()
        xyz2 = xyz2.contiguous()
        dist1 = torch.zeros(batchsize, n)
        dist2 = torch.zeros(batchsize, m)

        idx1 = torch.zeros(batchsize, n, dtype=torch.int)
        idx2 = torch.zeros(batchsize, m, dtype=torch.int)

        if not xyz1.is_cuda:
            cd.forward(xyz1, xyz2, dist1, dist2, idx1, idx2)
        else:
            dist1 = dist1.cuda()
            dist2 = dist2.cuda()
            idx1 = idx1.cuda()
            idx2 = idx2.cuda()
            cd.forward_cuda(xyz1, xyz2, dist1, dist2, idx1, idx2)

        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)

        return dist1, dist2

    @staticmethod
    def backward(ctx, graddist1, graddist2):
        xyz1, xyz2, idx1, idx2 = ctx.saved_tensors

        graddist1 = graddist1.contiguous()
        graddist2 = graddist2.contiguous()

        gradxyz1 = torch.zeros(xyz1.size())
        gradxyz2 = torch.zeros(xyz2.size())

        if not graddist1.is_cuda:
            cd.backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2)
        else:
            gradxyz1 = gradxyz1.cuda()
            gradxyz2 = gradxyz2.cuda()
            cd.backward_cuda(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2)

        return gradxyz1, gradxyz2


class ChamferDistance(torch.nn.Module):
    def forward(self, xyz1, xyz2):
        return ChamferDistanceFunction.apply(xyz1, xyz2)

model/chamfer_distance/chamfer_distance_gpu.py
augcog/DTTD2
#@title Import packages
import torch
import torch.nn as nn
import torch.nn.functional as F 
import einops
import numpy as np


class Transformer(nn.Module):
    """Transformer Encoder 
    Args:
        embedding_dim: dimension of embedding
        n_heads: number of attention heads
        n_layers: number of attention layers
        feedforward_dim: hidden dimension of MLP layer
    Returns:
        Transformer embedding of input
    """
    def __init__(self, embedding_dim=256, n_heads=4, n_layers=3, feedforward_dim=1024, dropout=0.0):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.feedforward_dim = feedforward_dim
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=embedding_dim,
                nhead=self.n_heads,
                dim_feedforward=self.feedforward_dim,
                activation=F.gelu,
                batch_first=True,
                dropout=dropout,
            ),
            num_layers=n_layers,
        )

    def forward(self, x):
        return self.transformer(x)
    


class MLPs(nn.Module):
    def __init__(self, dim_in, dim_out):
        super(MLPs, self).__init__()
        self.mlp = nn.Sequential(
            nn.Conv1d(dim_in, 640, 1),
            nn.ReLU(),
            nn.Conv1d(640, 256, 1),
            nn.ReLU(),
            nn.Conv1d(256, 128, 1),
            nn.ReLU(),
            nn.Conv1d(128, dim_out, 1)
            )
    def forward(self, x):
        return self.mlp(x)
    

class LayerNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-12):
        """
            Construct a layernorm module in the TF style (epsilon inside the square root).
        """
        super(LayerNorm, self).__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.bias = nn.Parameter(torch.zeros(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x):
        u = x.mean(-1, keepdim=True)
        s = (x - u).pow(2).mean(-1, keepdim=True)
        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
        return self.weight * x + self.bias

class FilterLayer(nn.Module):
    '''
        https://arxiv.org/abs/2202.13556
    '''
    def __init__(self, max_seq_length, hidden_size, hidden_dropout_prob):
        super(FilterLayer, self).__init__()
        self.complex_weight = nn.Parameter(torch.randn(1, max_seq_length//2 + 1, hidden_size, 2, dtype=torch.float32) * 0.02)
        self.out_dropout = nn.Dropout(hidden_dropout_prob)
        self.LayerNorm = LayerNorm(hidden_size, eps=1e-12)

    def forward(self, input_tensor, return_filtered=False):
        input_tensor = input_tensor.transpose(2, 1).contiguous()
        batch, seq_len, hidden = input_tensor.shape
        x = torch.fft.rfft(input_tensor, dim=1, norm='ortho')
        weight = torch.view_as_complex(self.complex_weight)
        x = x * weight
        sequence_emb_fft = torch.fft.irfft(x, n=seq_len, dim=1, norm='ortho')
        if return_filtered: return sequence_emb_fft
        hidden_states = self.out_dropout(sequence_emb_fft)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states.transpose(2, 1).contiguous()
    
    def visualize_frequency_domain(self, input_tensor):
        import matplotlib.pyplot as plt
        import matplotlib
        font = {'family' : 'Times New Roman',
                'weight' : 'bold',
                'size'   : 22}

        matplotlib.rc('font', **font)
        output_tensor = self.forward(input_tensor, return_filtered=True)
        print(output_tensor.shape)

        input_tensor = input_tensor.cpu().detach().numpy()[0].T
        output_tensor = output_tensor.cpu().detach().numpy()[0]

        def save(input_tensor, fn, title):   
            # Perform SVD
            U, s, V = np.linalg.svd(input_tensor, full_matrices=False)

            projected_tensor = input_tensor@U[0,:]

            plt.figure(figsize=(10, 6))
            plt.hist(projected_tensor, bins=1000, edgecolor='red')

            plt.xlabel('Geometric Features Reduced to 1-dimensional')
            plt.ylabel('Probability Density')
            plt.title(title)
            plt.grid(True)
            plt.savefig(fn)

        save(input_tensor, 'before_filtered.png', 'Before')
        save(output_tensor, 'after_filtered.png', 'After')


#  Transformer Customization
future_mask = torch.triu(torch.zeros([1024, 1024]).fill_(float("-inf")), 1)
def scaled_dot_product_attention(q, k, v, key_padding_mask=None, causal=False, return_attn = False):
    d_head = q.size(-1)
    s = einops.einsum(q, k, "n tl dh, n sl dh -> n tl sl") / d_head ** 0.5
    if key_padding_mask is not None:
        s = s.masked_fill(
            key_padding_mask.unsqueeze(1).to(torch.bool),
            float("-inf"),
        )
    if causal:
        attn_mask = future_mask[: s.size(1), : s.size(2)].to(s)
        s += attn_mask.unsqueeze(0)
    a = F.softmax(s, dim=-1, dtype=torch.float32).type_as(s)
    if return_attn:
        return a
    return einops.einsum(a, v, "n tl sl, n sl dh -> n tl dh")

class MultiheadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
                
        self.n_heads = n_heads
        self.d_model = d_model
        self.d_head = d_model // n_heads

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.o_proj = nn.Linear(d_model, d_model)
        self.apply(self._init_weights)
        self._init_mimetic()

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=1.0)
            if module.bias is not None:
                module.bias.data.zero_()

    def _init_mimetic(self):
        self.q_proj /= self.d_head
        self.k_proj /= self.d_head
        self.v_proj *= (1/self.d_model)
        self.o_proj *= (-1/self.d_model)

    
    def forward(self, q, k, v, key_padding_mask=None, causal=False):
        q = self.q_proj(q)
        k = self.k_proj(k)
        v = self.v_proj(v)
        q = einops.rearrange(q, "b tl (nh dh) -> (b nh) tl dh", nh=self.n_heads)
        k = einops.rearrange(k, "b sl (nh dh) -> (b nh) sl dh", nh=self.n_heads)
        v = einops.rearrange(v, "b sl (nh dh) -> (b nh) sl dh", nh=self.n_heads)
        if key_padding_mask is not None:
            key_padding_mask = einops.repeat(
                key_padding_mask, "b sl -> (b nh) sl", nh=self.n_heads)
        o = scaled_dot_product_attention(q, k, v, key_padding_mask, causal)
        o = einops.rearrange(o, "(b nh) tl dh -> b tl (nh dh)", nh=self.n_heads)
        return self.o_proj(o)

class HandmadeTransformerLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ffn, p_drop, is_decoder=False):
        super().__init__()
        self.is_decoder = is_decoder
        self.self_attn = MultiheadAttention(d_model, n_heads)
        self.self_attn_drop = nn.Dropout(p_drop)
        self.self_attn_ln = nn.LayerNorm(d_model)
        if is_decoder:
            self.cross_attn = MultiheadAttention(d_model, n_heads)
            self.cross_attn_drop = nn.Dropout(p_drop)
            self.cross_attn_ln = nn.LayerNorm(d_model)
        self.fc1 = nn.Linear(d_model, d_ffn)
        self.fc2 = nn.Linear(d_ffn, d_model)
        self.ffn_drop = nn.Dropout(p_drop)
        self.ffn_ln = nn.LayerNorm(d_model)
    
    def forward(self, x, padding_mask=None, encoder_out=None, encoder_padding_mask=None):
        residual = x
        x = self.self_attn(x, x, x, padding_mask, causal=self.is_decoder)
        x = self.self_attn_drop(x)
        x = self.self_attn_ln(x + residual)

        if self.is_decoder:
            residual = x
            x = self.cross_attn(x, encoder_out, encoder_out, encoder_padding_mask)
            x = self.cross_attn_drop(x)
            x = self.cross_attn_ln(x + residual)

        residual = x
        x = self.fc2(F.relu(self.fc1(x)))
        x = self.ffn_drop(x)
        x = self.ffn_ln(x + residual)
        return x



from typing import  Optional, Tuple
Tensor = torch.Tensor

def simplified_attention_forward(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    num_heads: int,
    in_proj_weight: Optional[Tensor],
    in_proj_bias: Optional[Tensor],
    only_attn: bool=True
    ) -> Tuple[Tensor, Optional[Tensor]]:

    # set up shape vars
    tgt_len, bsz, embed_dim = query.shape
    head_dim = embed_dim // num_heads

    # compute in-projection
    q, k, v = F._in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)

    # reshape q, k, v for multihead attention and make em batch first
    q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
    k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
    v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)

    # (deep breath) calculate attention
    attn_map = scaled_dot_product_attention(q, k, v, return_attn=True)
    if only_attn:
        return attn_map
    embeddings = einops.einsum(attn_map, v, "n tl sl, n sl dh -> n tl dh")
    embeddings = einops.rearrange(embeddings, "(b n) tl dh -> b tl (n dh)", b = bsz)
    return embeddings, attn_map


def attn_diverse_loss(attn, th=2):
    sim_sum = 0
    counter = 1e-6
    for i in range(len(attn)):
        mask0 = attn[i].mean(dim=1).squeeze()
        n_tokens = mask0.shape[-1]
        threshold = th/n_tokens
        score0 = torch.mean(mask0, dim=1, keepdim=True)
        mask0 = (mask0 > threshold) * (mask0)
        score0 = (score0 > threshold) * (score0)
        sim = F.cosine_similarity(score0, mask0, dim=-1)
        sim = sim.mean()
        sim_sum += sim
        counter += 1
    sim_sum = sim_sum / counter
    return sim_sum


model/model_utils.py
augcog/DTTD2
import torch
import torch.nn as nn
import torch.nn.functional as F
import sys
sys.path.append("../../")

if torch.cuda.is_available():
    from model.chamfer_distance.chamfer_distance_gpu import ChamferDistance # https://github.com/chrdiller/pyTorchChamferDistance
else:
    from model.chamfer_distance.chamfer_distance_cpu import ChamferDistance # https://github.com/chrdiller/pyTorchChamferDistance

class PointCloudEncoder(nn.Module):
    def __init__(self, emb_dim):
        super(PointCloudEncoder, self).__init__()
        self.conv1 = nn.Conv1d(3, 64, 1)
        self.conv2 = nn.Conv1d(64, 128, 1)
        self.conv3 = nn.Conv1d(256, 1024, 1)
        self.fc = nn.Linear(1024, emb_dim)
        # TODO: add layer norm, avg pool

    def forward(self, xyz):
        """
        Args:
            xyz: (B, 3, N)
        """
        np = xyz.size()[2]
        x = F.relu(self.conv1(xyz))
        x = F.relu(self.conv2(x))
        global_feat = F.adaptive_max_pool1d(x, 1)
        x0 = torch.cat((x, global_feat.repeat(1, 1, np)), dim=1)
        x = F.relu(self.conv3(x0))
        x = torch.squeeze(F.adaptive_max_pool1d(x, 1), dim=2)
        embedding = self.fc(x)
        return embedding, x0


class PointCloudDecoder(nn.Module):
    def __init__(self, emb_dim, n_pts, out_channels=3):
        super(PointCloudDecoder, self).__init__()
        self.out_channels = out_channels
        self.fc1 = nn.Linear(emb_dim, 512)
        self.fc2 = nn.Linear(512, 1024)
        self.fc3 = nn.Linear(1024, out_channels*n_pts)
        self.norm = nn.LayerNorm(out_channels*n_pts)

    def forward(self, embedding):
        """
        Args:
            embedding: (B, 512)

        """
        bs = embedding.size()[0]
        out = F.relu(self.fc1(embedding))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)
        out = self.norm(out)
        out_pc = out.view(bs, -1, self.out_channels)
        return out_pc


class PointCloudAE(nn.Module):
    def __init__(self, emb_dim=256, n_pts=1024, latent_dim=None):
        super(PointCloudAE, self).__init__()
        self.num_points = n_pts
        self.encoder = PointCloudEncoder(emb_dim)
        self.decoder = PointCloudDecoder(emb_dim, n_pts)
        if latent_dim is not None:
            self.ptfeat = nn.Conv1d(256+emb_dim, latent_dim, 1)
            # self.ptfeat = nn.Conv1d(256, latent_dim, 1)

    def latent(self, feat, emb):
        return self.ptfeat(torch.cat([feat, emb[:,:,None].repeat(1, 1, self.num_points)], dim=1))
        # return self.ptfeat(feat + emb[:,:,None].repeat(1, 1, self.num_points))

    def forward(self, in_pc, emb=None, recon_ref=None):
        """
        Args:
            in_pc: (B, N, 3)
            emb: (B, 512)

        Returns:
            emb: (B, emb_dim)
            out_pc: (B, n_pts, 3)

        """
        if emb is None:
            xyz = in_pc.permute(0, 2, 1)
            emb, feat = self.encoder(xyz)
        if recon_ref is None:
            recon_ref = in_pc
        out_pc = self.decoder(emb)
        loss = recon_loss(recon_ref, out_pc)
        return feat, emb, out_pc, loss

def recon_loss(points, recon_points):
    chamfer_dist = ChamferDistance()
    points = points.transpose(1,2)
    reconstructed_points = recon_points.transpose(1,2)
    dist1, dist2 = chamfer_dist(points, reconstructed_points)   # calculate loss
    loss = (torch.mean(dist1)) + (torch.mean(dist2))
    return loss
model/pointae.py
augcog/DTTD2
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
from model.pspnet import PSPNet
from model.pointae import PointCloudAE
from model.model_utils import *
from model.model_utils import simplified_attention_forward, attn_diverse_loss
from einops import rearrange

class ModifiedResnet(nn.Module):
    def __init__(self, out_dim=128):
        super(ModifiedResnet, self).__init__()
        self.model = PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet18', pretrained=True, out_dim=out_dim)

    def forward(self, x):
        x = self.model(x)
        return x

class BaseFormer(nn.Module):
    def __init__(self, feat_dim, embedding_dim, n_heads=2, n_layers=1, encoder_type='pytorch', require_adl=False):
        super().__init__()
        self.encoder_type = encoder_type
        self.n_layers = n_layers
        self.transformer = Transformer(
            embedding_dim, 
            n_heads=n_heads, 
            n_layers=n_layers, 
            feedforward_dim=512,
            dropout=0.1
        )
        self.linear_projection = nn.Conv1d(feat_dim , embedding_dim, 1)
        # A Layernorm and a Linear layer are applied on the encoder embeddings
        self.norm = nn.LayerNorm(embedding_dim)
        self.require_adl = require_adl

    def forward(self, feat):
        embeddings = self.linear_projection(feat).transpose(2, 1).contiguous() 
        if not self.require_adl:
            embeddings = self.transformer(embeddings)
        else:
            attn_map_list = []
            for i, encoder_layer in enumerate(self.transformer.transformer.layers): 
                attn_layer = encoder_layer.self_attn
                query = key = value = embeddings.transpose(1, 0)
                x, attn_map = simplified_attention_forward(
                    query, key, value, attn_layer.num_heads,
                    attn_layer.in_proj_weight, attn_layer.in_proj_bias, only_attn=False)
                embeddings = encoder_layer.norm1(embeddings + encoder_layer.dropout1(x))
                embeddings = encoder_layer.norm2(embeddings + encoder_layer._ff_block(embeddings))

                attn_map_list.append(attn_map)

            adl = attn_diverse_loss(attn_map_list)

        embeddings = self.norm(embeddings) #B,seq_len,emb_dim 
        return embeddings if not self.require_adl else (embeddings, adl)

    def get_attention_map(self, feat, layer_id=None):
        if self.encoder_type == 'customized':
            raise NotImplementedError
        if self.encoder_type == 'pytorch':
            if layer_id is None:
                layer_id = self.n_layers - 1
            if layer_id >= self.n_layers or layer_id < 0:
                raise ValueError(
                    f"Provided layer_id: {layer_id} is not valid. 0 <= {layer_id} < {self.n_layers}."
                )
            x = self.linear_projection(feat).transpose(2, 1).contiguous()
            for i, encoder_layer in enumerate(self.transformer.transformer.layers): 
                if i < layer_id:
                    x = encoder_layer(x)
                else:
                    attn_layer = encoder_layer.self_attn
                    query = key = value = x.transpose(1, 0)
                    attn_map = simplified_attention_forward(
                        query, key, value, attn_layer.num_heads,
                        attn_layer.in_proj_weight, attn_layer.in_proj_bias)
                    return attn_map



class PosePredictor(nn.Module):
    def __init__(self, feat_dim, num_points, num_obj):
        super(PosePredictor, self).__init__()
        self.num_points = num_points
        self.num_obj = num_obj
        #pose predictor
        self.conv_r = MLPs(feat_dim, num_obj*4) #quaternion
        self.conv_t = MLPs(feat_dim, num_obj*3) #translation
        self.conv_c = MLPs(feat_dim, num_obj*1) #confidence
        
    def forward(self, ap_x, obj):
        bs, _, _ = ap_x.size()
        rx = self.conv_r(ap_x).view(bs, self.num_obj, 4, self.num_points)
        tx = self.conv_t(ap_x).view(bs, self.num_obj, 3, self.num_points)
        cx = torch.sigmoid(self.conv_c(ap_x)).view(bs, self.num_obj, 1, self.num_points)
        
        b = 0
        out_rx = torch.index_select(rx[b], 0, obj[b]).contiguous().transpose(2, 1).contiguous()
        out_tx = torch.index_select(tx[b], 0, obj[b]).contiguous().transpose(2, 1).contiguous()
        out_cx = torch.index_select(cx[b], 0, obj[b]).contiguous().transpose(2, 1).contiguous()
        return out_rx, out_tx, out_cx

######## PoseFusion ########
class FusionBlock(nn.Module):
    def __init__(self, base_latent, embed_dim, n_layer1, n_layer2, require_adl=False):
        super(FusionBlock, self).__init__()
        self.embed_dim = embed_dim
        self.require_adl = require_adl

        if n_layer1>0:
            self.modality_fusion = BaseFormer(base_latent, base_latent, 4, n_layer1, require_adl=require_adl)
            modality_dim = 4*base_latent
        else:
            self.modality_fusion = None
            modality_dim = 2*base_latent
        self.point_fusion = BaseFormer(modality_dim, embed_dim, 8, n_layer2, require_adl=require_adl) if n_layer2>0 else None
        assert (self.modality_fusion is not None or self.point_fusion is not None), \
                "<Error Message> Either the number of modality or point-wise fusion layer \
                    should be larger than 0. \
                    Otherwise, the fusion block would be empty."

    def forward(self, rgb_emb, pt_emb, hidden_state=None, require_attn=False):
        cross_feat = None
        global_feat = None

        if self.modality_fusion is not None:
            if hidden_state is None:
                hidden_state = torch.cat([rgb_emb, pt_emb], dim=2)
            else:
                hidden_state = rearrange(hidden_state, 'b (k d) n -> b d (k n)', k=2)
                hidden_state += torch.cat([rgb_emb, pt_emb], dim=2)
            feat = self.modality_fusion(hidden_state)
            if self.require_adl:
                adl = feat[1]
                feat = feat[0]
            if require_attn: 
                attn1 = self.modality_fusion.get_attention_map(hidden_state)
            cross_feat = rearrange(feat, 'b (k n) d -> b n (k d)', k=2).transpose(2, 1).contiguous()

        if self.point_fusion is not None:
            feat = torch.cat([rgb_emb, pt_emb, cross_feat], dim=1) if cross_feat is not None else torch.cat([rgb_emb, pt_emb], dim=1)
            global_feat = self.point_fusion(feat)
            if self.require_adl:
                adl += global_feat[1]
                global_feat = global_feat[0]
                
            global_feat = global_feat.transpose(2, 1).contiguous()
            if require_attn: 
                attn2 = self.point_fusion.get_attention_map(feat)
        if require_attn:
            return cross_feat, global_feat, attn1, attn2
        if self.require_adl:
            return cross_feat, global_feat, adl
        return cross_feat, global_feat 
    
class PoseFusion(nn.Module):
    def __init__(self, base_latent, embed_dim, n_block, n_layer1, n_layer2, require_adl=False):
        super(PoseFusion, self).__init__()
        self.n_block = n_block
        self.fusion_mode = (n_layer1 * n_layer2 > 0)
        self.require_adl = require_adl
        self.layers = self._make_layer(base_latent, embed_dim, n_layer1, n_layer2, require_adl)
        assert embed_dim == 2 * base_latent
        
    def forward(self, rgb_emb, pt_emb):
        if self.fusion_mode:
            result = None
            adl = 0
            for i, layer in enumerate(self.layers):
                if not self.require_adl:
                    cross_feat, global_feat = layer(rgb_emb, pt_emb)
                else:
                    cross_feat, global_feat, adl_ = layer(rgb_emb, pt_emb)
                    adl += adl_
                result = torch.cat([result, cross_feat, global_feat], dim=1) if result is not None else torch.cat([cross_feat, global_feat], dim=1)
        else:
            assert self.n_block == 1, \
                "<Warning> You only have 1 fusion stage. Please merge all layers into 1 block. To fix it, you can set n_block=1"
            cross_feat, global_feat = self.layers[0](rgb_emb, pt_emb)
            result = cross_feat if global_feat is None else global_feat
        return result if not self.require_adl else (result, adl)
    
    def _make_layer(self, base_latent, embed_dim, n_layer1, n_layer2, require_adl):
        layers = []
        for i in range(0, self.n_block):
            layers.append(FusionBlock(base_latent, embed_dim, n_layer1, n_layer2, require_adl=require_adl))
        return nn.Sequential(*layers)


class PoseNet(nn.Module):
    def __init__(self, num_points, num_obj, \
                 base_latent=256, embedding_dim=512, fusion_block_num=1, layer_num_m=2, layer_num_p=4, \
                 recon_choice='depth', filter_enhance=True, require_adl=False):
        super(PoseNet, self).__init__()
        self.num_points = num_points
        self.num_obj = num_obj
        self.base_latent = base_latent
        self.embedding_dim = embedding_dim
        self.fusion_block_num = fusion_block_num
        self.layer_num_m = layer_num_m
        self.layer_num_p = layer_num_p
        self.fusion_mode = (layer_num_m*layer_num_p>0)
        self._init_config_check()
        
        # unimodal embedding
        self.cnn = ModifiedResnet(base_latent)
        self.recon_choice = recon_choice
        self.ptnet = PointCloudAE(256, num_points, base_latent)
        self.modelnet = PointCloudAE(256, num_points, base_latent) if recon_choice=='both' else None
        self.filter_enhance = FilterLayer(num_points, base_latent, 0.0) if filter_enhance else None
        
        # modality and position interaction
        self.fusion = PoseFusion(base_latent, embedding_dim, \
                                 fusion_block_num, layer_num_m, layer_num_p, \
                                 require_adl)
        self.require_adl = require_adl
        # prediction
        if self.fusion_mode:
            self.posepred = PosePredictor(base_latent * 4, num_points, num_obj)
        else:
            self.posepred = PosePredictor(base_latent * 2, num_points, num_obj)

    def _init_config_check(self):
        print("ResNet &PtNet Output Dim:", self.base_latent, '\n',\
              "Fusion Block Num:", self.fusion_block_num, '\n',\
              "Modality Fusion Layer Num:", self.layer_num_m, '\n',\
              "Point2point Fusion Layer Num:", self.layer_num_p)
        
    def forward(self, img, x, choose, obj, recon_ref=None):
        # rgb color embedding
        out_img = self.cnn(img) 
        bs, di, _, _ = out_img.size()
        emb = out_img.view(bs, di, -1)
        robust_loss = 0

        # selection of rgb color embedding
        choose = choose.repeat(1, di, 1) 
        rgb_emb = torch.gather(emb, 2, choose).contiguous()

        # depth map / point cloud (embedding)
        if self.recon_choice == 'both':
            object_geo = recon_ref[1]
            recon_ref = recon_ref[0]
        pt_feat, pt_emb, pt_recon, cdl_0 = self.ptnet(x, None, recon_ref)
        robust_loss+=cdl_0
        pt_emb = self.ptnet.latent(pt_feat, pt_emb)
        if self.filter_enhance is not None:
            pt_emb = self.filter_enhance(pt_emb)
        if self.recon_choice == 'both':
            _, obj_emb, _, cdl_1 = self.modelnet(x, None, object_geo)
            robust_loss+=cdl_1
            pt_emb += obj_emb
        feat = self.fusion(rgb_emb, pt_emb) 

        if self.require_adl:
            adl = feat[1] 
            feat = feat[0]
            robust_loss+=adl
            
        out_rx, out_tx, out_cx = self.posepred(feat, obj)

        return out_rx, out_tx, out_cx, rgb_emb.detach(), pt_recon.detach(), robust_loss
    
    def get_attention_map(self, img, x, choose):
        # rgb color embedding
        out_img = self.cnn(img) 
        bs, di, _, _ = out_img.size()
        emb = out_img.view(bs, di, -1)

        # selection of rgb color embedding
        choose = choose.repeat(1, di, 1) 
        rgb_emb = torch.gather(emb, 2, choose).contiguous()

        # depth map / point cloud (embedding)
        pt_feat, pt_emb, _, _ = self.ptnet(x, None, None)
        pt_emb = self.ptnet.latent(pt_feat, pt_emb)
        if self.filter_enhance is not None:
            pt_emb = self.filter_enhance(pt_emb)
        # feat = torch.cat([rgb_emb, pt_emb], dim=2)
        _, _, attn1, attn2 = self.fusion.layers[0](rgb_emb, pt_emb, require_attn=True)
    
        return attn1, attn2
    
    def get_freq_domain(self, x):
        pt_feat, pt_emb, _, _ = self.ptnet(x, None, None)
        pt_emb = self.ptnet.latent(pt_feat, pt_emb)
        assert self.filter_enhance is not None, "filter enhanced MLP is not applied."
        self.filter_enhance.visualize_frequency_domain(pt_emb)

model/posefusion.py
augcog/DTTD2
import torch
from torch import nn
from torch.nn import functional as F
import model.resnet as resnets

class PSPModule(nn.Module):
    def __init__(self, features, out_features=1024, sizes=(1, 2, 3, 6)):
        super(PSPModule, self).__init__()
        self.stages = []
        self.stages = nn.ModuleList([self._make_stage(features, size) for size in sizes])
        self.bottleneck = nn.Conv2d(features * (len(sizes) + 1), out_features, kernel_size=1)
        self.relu = nn.ReLU()

    def _make_stage(self, features, size):
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        conv = nn.Conv2d(features, features, kernel_size=1, bias=False)
        return nn.Sequential(prior, conv)

    def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        priors = [F.interpolate(input=stage(feats), size=(h, w), mode='bilinear') for stage in self.stages] + [feats]
        bottle = self.bottleneck(torch.cat(priors, 1))
        return self.relu(bottle)


class PSPUpsample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(PSPUpsample, self).__init__()
        self.conv = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.PReLU()
        )

    def forward(self, x):
        return self.conv(x)


class PSPNet(nn.Module):
    def __init__(self, n_classes=21, sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet18',
                 pretrained=False, out_dim=32):
        super(PSPNet, self).__init__()
        self.feats = getattr(resnets, backend)(pretrained)
        self.psp = PSPModule(psp_size, 1024, sizes)
        self.drop_1 = nn.Dropout2d(p=0.3)
        # self.up_1 = PSPUpsample(1024, 256)
        # self.up_2 = PSPUpsample(256, 64)
        # self.up_3 = PSPUpsample(64, 64)
        self.up_1 = PSPUpsample(1024, 512)
        self.up_2 = PSPUpsample(512, 256)
        self.up_3 = PSPUpsample(256, 256)

        self.drop_2 = nn.Dropout2d(p=0.15)
        self.final = nn.Sequential(
            nn.Conv2d(256, out_dim, kernel_size=1),
            nn.LogSoftmax(dim=1)
        )

        self.classifier = nn.Sequential(
            nn.Linear(deep_features_size, 256),
            nn.ReLU(),
            nn.Linear(256, n_classes)
        )

    def forward(self, x):
        f, class_f = self.feats(x) 
        p = self.psp(f)
        p = self.drop_1(p)

        p = self.up_1(p)
        p = self.drop_2(p)

        p = self.up_2(p)
        p = self.drop_2(p)

        p = self.up_3(p)
        
        return self.final(p)

model/pspnet.py
augcog/DTTD2
import torch
import math
from collections import OrderedDict
import torch.nn as nn

def load_weights_sequential(target, source_state):
    new_dict = OrderedDict()
    for (k1, v1), (k2, v2) in zip(target.state_dict().items(), source_state.items()):
        new_dict[k1] = v2
    target.load_state_dict(new_dict)

def conv3x3(in_planes, out_planes, stride=1, dilation=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, dilation=dilation, bias=False)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride=stride, dilation=dilation)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes, stride=1, dilation=dilation)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.relu(out)

        out = self.conv2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    expansion = 4
    def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, dilation=dilation,
                               padding=dilation, bias=False)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.relu(out)

        out = self.conv3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class ResNet(nn.Module):
    def __init__(self, block, layers=(3, 4, 23, 3)):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False)
            )

        layers = [block(self.inplanes, planes, stride, downsample)]
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, dilation=dilation))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x_3 = self.layer3(x)
        x = self.layer4(x_3)

        return x, x_3

def load_pretrain(model, name):
    dict_url = {
        'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
        'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
        'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
        'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
        'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
        'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
        'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
        'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
        'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth'
    }
    if name in dict_url:
        pretrained_dict = torch.hub.load_state_dict_from_url(dict_url[name])
        state_dict = OrderedDict()
        for k in pretrained_dict:
            if k in model.state_dict():
                state_dict[k] = pretrained_dict[k]
        model.load_state_dict(state_dict)

def resnet18(pretrained=False):
    model = ResNet(BasicBlock, [2, 2, 2, 2])
    if pretrained:
        load_pretrain(model, "resnet18")
    return model

def resnet34(pretrained=False):
    model = ResNet(BasicBlock, [3, 4, 6, 3])
    return model

def resnet50(pretrained=False):
    model = ResNet(Bottleneck, [3, 4, 6, 3])
    return model

def resnet101(pretrained=False):
    model = ResNet(Bottleneck, [3, 4, 23, 3])
    return model

def resnet152(pretrained=False):
    model = ResNet(Bottleneck, [3, 8, 36, 3])
    return model
model/resnet.py
augcog/DTTD2
import argparse
import os
from datetime import datetime

import torch
import torch.optim as optim
from loss.add_s import ADDS_Loss
# from loss.gadd import GADD_Loss
from loss.abc import ABC_Loss

from model.posefusion import PoseNet
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from utils.log import Logger

parser = argparse.ArgumentParser()
parser.add_argument('--device', type=int, default = 0, help='device number, not valid when set data parallel')
parser.add_argument('--dataset', type=str, default = 'iphone', choices=['ycb', 'dttd', 'iphone'], help='Dataset name: ycb, dttd, ...')
parser.add_argument('--dataset_root', type=str, default = './dataset/dttd_iphone/DTTD_IPhone_Dataset/root', help='dataset root dir (''YCB_Video_Dataset'')')
parser.add_argument('--dataset_config', type=str, default = './dataset/dttd_iphone/dataset_config', help='dataset config dir')

parser.add_argument('--loss', type=str, default = 'adds', choices=['adds', 'abc'])
parser.add_argument('--recon_choice', type=str, default = 'depth', choices=['depth', 'model'])
parser.add_argument('--recon_w', default=0.3, type=float, help='recon_weight')
parser.add_argument('--filter_enhance', action='store_true')
parser.add_argument('--attn_diverse', action='store_true')

parser.add_argument('--base_latent', type=int, default = 128, help='base latent dim for unimodal encoder')
parser.add_argument('--embed_dim', type=int, default = 512, help='embedding dim for transformer encoder')
parser.add_argument('--fusion_block_num', type=int, default = 1, help='number of fusion block')
parser.add_argument('--layer_num_m', type=int, default = 3, help='layer num for modality fusion per block')
parser.add_argument('--layer_num_p', type=int, default = 3, help='layer num for point-to-point fuison per block')

parser.add_argument('--pretrain', type=str, default = None, help='path to pretrain model')
parser.add_argument('--frozen', type=str, default = None, choices=['ptnet', 'filter'])
parser.add_argument('--output_dir', type=str, default = './out', help='output directory, store output info')
parser.add_argument('--optim_batch', type=int, default = 4, help='number of batch per optim')
parser.add_argument('--workers', type=int, default = 8, help='number of data loading workers')
parser.add_argument('--data_parallel', action='store_true')
parser.add_argument('--batch_size', default=1, type=int, help='batch size')

parser.add_argument('--lr', default=1e-5, type=float, help='learning rate')
parser.add_argument('--min_lr', default=1e-6, type=float, help='min learning rate')
parser.add_argument('--warm_epoch', type=int, default=1, help='epoch for warmup')
parser.add_argument('--warm_iters', type=int, default=0, help='iters for warmup')
parser.add_argument('--lr_rate', default=0.3, type=float, help='learning rate decay rate')
parser.add_argument('--weight_decay', default=0.05, type=float, help='weight decay')
# parser.add_argument('--clip_grad', default=None, type=float, help='clip grad norm')

parser.add_argument('--w', default=0.015, type=float, help='balancing rate')
parser.add_argument('--w_rate', default=0.3, type=float, help='balancing rate decay rate')
parser.add_argument('--decay_margin', default=0.033, type=float, help='margin to decay lr & w')
parser.add_argument('--decay_rate', default=0.77, type=float, help='decay rate for decay margin')
parser.add_argument('--noise_trans', default=0.03, type=float, help='range of the random noise of translation added to the training data')
parser.add_argument('--nepoch', type=int, default=120, help='max number of epochs to train')
parser.add_argument('--resume_posenet', type=str, default = '',  help='resume PoseNet model')
parser.add_argument('--start_epoch', type=int, default = 0, help='which epoch to start')

parser.add_argument('--single_obj_test', type=int, default = -1, help='train based on a single obj')
parser.add_argument('--exclude_single_obj', type=int, default = -1, help='exclude a single obj')

opt = parser.parse_args()

import math

import numpy as np


def cosine_scheduler(lr, min_lr, epochs, steps_per_epoch, warmup_epochs=0,
                     start_warmup_value=0, warmup_steps=-1, cur_epoch=None):
    warmup_schedule = np.array([])
    warmup_iters = warmup_epochs * steps_per_epoch
    # set up warmup iterations
    if warmup_steps > 0:
        warmup_iters = warmup_steps
    if warmup_iters > 0:
        warmup_schedule = np.linspace(start_warmup_value, lr, warmup_iters)
    iters = np.arange(epochs * steps_per_epoch - warmup_iters)
    # set up schedule
    if cur_epoch is None:
        schedule = np.array([min_lr + 0.5 * (lr - min_lr) * (1 + math.cos(math.pi * i / (len(iters)))) for i in iters])
        schedule = np.concatenate((warmup_schedule, schedule))
        assert len(schedule) == epochs * steps_per_epoch
        return schedule
    elif cur_epoch*steps_per_epoch<warmup_iters:
        schedule = np.array([min_lr + 0.5 * (lr - min_lr) * (1 + math.cos(math.pi * i / (len(iters)))) for i in np.arange(steps_per_epoch)])    
        warmup_schedule = np.concatenate((warmup_schedule, schedule))
        return warmup_schedule[cur_epoch*steps_per_epoch : cur_epoch*steps_per_epoch+steps_per_epoch]
    else:
        return np.array(
            [min_lr + 0.5 * (lr - min_lr) * (1 + math.cos(math.pi * (i - warmup_iters) / (len(iters)))) for i in np.arange(cur_epoch*steps_per_epoch, cur_epoch*steps_per_epoch+steps_per_epoch)])    
    
def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = [p for p in parameters if p.grad is not None]
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    if norm_type == torch.inf:
        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
    else:
        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
    return total_norm

def main():
    # set device
    if not opt.data_parallel:
        torch.cuda.set_device(opt.device)

    # set output
    opt.output_dir = opt.output_dir + '_m' + str(opt.layer_num_m) + 'p' + str(opt.layer_num_p)\
        + '_' +opt.dataset + '_' + datetime.now().strftime('%m%d-%H%M')
    out_dir = opt.output_dir
    cp_dir = os.path.join(out_dir, "checkpoints")
    if not os.path.isdir(out_dir):
        os.makedirs(out_dir)
    if not os.path.isdir(cp_dir):
        os.makedirs(cp_dir)
    logger = Logger(os.path.join(out_dir, "log_train.txt"))
    writer = SummaryWriter(os.path.join(out_dir, "train_event"))

    # set dataset
    if opt.dataset == "ycb":
        from dataset.ycb.dataset import YCBDataset
        train_dataset = YCBDataset(root=opt.dataset_root, config_path=opt.dataset_config, mode='train', add_noise=True)
        test_dataset = YCBDataset(root=opt.dataset_root, config_path=opt.dataset_config, mode='test', add_noise=False)
        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers)
        test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=opt.workers)
    
    elif opt.dataset == "dttd":
        # from dataset.dttd.dataset import DTTDDataset
        raise NotImplementedError

    elif opt.dataset == "iphone":
        from dataset.dttd_iphone.dataset import DTTDDataset
        train_dataset = DTTDDataset(root=opt.dataset_root, mode='train', add_noise=True)
        test_dataset = DTTDDataset(root=opt.dataset_root, mode='test', add_noise=False)
        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers)
        test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=opt.workers)
    else:
        raise NotImplementedError
    
    # set model
    estimator = PoseNet(num_points=train_dataset.get_2d_sample_num(), \
                        num_obj = train_dataset.get_object_num(),\
                        base_latent=opt.base_latent, \
                        embedding_dim=opt.embed_dim, \
                        fusion_block_num=opt.fusion_block_num, \
                        layer_num_m=opt.layer_num_m, layer_num_p=opt.layer_num_p, \
                        filter_enhance = opt.filter_enhance,\
                        require_adl = opt.attn_diverse)
    if opt.pretrain:
        tar_params = estimator.state_dict()
        src_params = torch.load(opt.pretrain, map_location=torch.device('cpu'))
        diff_params = {k:v for k,v in tar_params.items() if k not in src_params.keys()}
        logger.log(f"Warnning: Params NOT be Found. Details: {diff_params.keys()}")
        estimator.load_state_dict(src_params, strict=False)
        if opt.frozen:
            if opt.frozen == 'filter':
                for param in estimator.filter_enhance.parameters():
                    param.requires_grad = False
            if opt.frozen == 'ptnet':
                for param in estimator.ptnet.parameters():
                    param.requires_grad = False
    estimator = estimator.cuda()
    adds_loss = ADDS_Loss(num_points_mesh=train_dataset.get_model_point_num(), sym_list=train_dataset.get_sym_list())
    abc_loss = ABC_Loss(num_points_mesh=train_dataset.get_model_point_num())

    # gadd_loss = GADD_Loss(torch.device("cuda"))
    optimizer = optim.AdamW(estimator.parameters(), lr=opt.lr, weight_decay=opt.weight_decay)
    
    logger.log(f"=========Object 6D Pose Estimation Training Log========")
    logger.log(f"training start date and time: {datetime.now().strftime('%d/%m/%Y %H:%M:%S')}")
    logger.log(f"dataset: {opt.dataset}")
    logger.log(f"data root: {opt.dataset_root}")
    logger.log(f"train set size: {len(train_dataset)}, test set size: {len(test_dataset)}")
    logger.log(f"output directory: {opt.output_dir}")
    logger.log(f"pretrain: {opt.pretrain}")
    logger.log(f"param frozen: {opt.frozen}")
    logger.log(f"resume posenet path: {opt.resume_posenet}")
    logger.log(f"device: {opt.device}")
    logger.log(f"network parallel: {opt.data_parallel}")
    logger.log(f"training on single obj: {opt.single_obj_test} exclude: {opt.exclude_single_obj}")

    logger.log(f"\n=========Object 6D Pose Estimation Model Architecture========")
    logger.log(f"base latent dim for unimodal encoder: {opt.base_latent}")
    logger.log(f"embedding dim for transformer encoder: {opt.embed_dim}")
    logger.log(f"number of fusion block: {opt.fusion_block_num}")
    logger.log(f"layer num for modality fusion per block: {opt.layer_num_m}")
    logger.log(f"layer num for point-to-point fuison per block: {opt.layer_num_p}")
    logger.log(f"loss type: {opt.loss}")
    logger.log(f"recon weight: {opt.recon_w}")
    logger.log(f"recon choice: {opt.recon_choice}")
    logger.log(f"filter enhance: {opt.filter_enhance}")
    logger.log(f"attn diverse: {opt.attn_diverse}")

    logger.log(f"\n=========Object 6D Pose Estimation Model Hyper Params========")
    logger.log(f"learning rate: {opt.lr}")
    logger.log(f"min learning rate: {opt.min_lr}")
    logger.log(f"warm lr: epoch {opt.warm_epoch} iters {opt.warm_iters}")
    # logger.log(f"clip grad: {opt.clip_grad}")
    logger.log(f"weight decay: {opt.weight_decay}")
    logger.log(f"decay margin: {opt.decay_margin}")
    logger.log(f"decay rate: {opt.decay_rate}")
    logger.log(f"optim batch: {opt.optim_batch}")
    logger.log(f"balancing rate: {opt.w}")
    logger.log(f"epoch num: {opt.nepoch}")
    logger.log(f"start epoch: {opt.start_epoch}")

    logger.log("\n\n")
    
    # start training
    best_test_dist = 1e10
    best_dist_2cm = 0
    for epoch in range(opt.start_epoch, opt.nepoch):
        # train
        logger.log(f"start training epoch {epoch}....")
        estimator.train()
        loss_list = []
        robust_loss_list = []
        dist_list = []
        # lr planning
        lr_schedule = cosine_scheduler(opt.lr, opt.min_lr, opt.nepoch, 
                                       len(train_dataset)//opt.optim_batch+1, 
                                       opt.warm_epoch, 0, opt.warm_iters, epoch)
        logger.log(f"learning rate planning: start {lr_schedule[0]} end {lr_schedule[-1]}\n")
        with tqdm(train_loader, unit="batch") as tepoch:
            for i, data in enumerate(tepoch):
                try:
                    if opt.single_obj_test>0:
                        assert data["obj_id"]==opt.single_obj_test
                    if opt.exclude_single_obj>0:
                        assert not data["obj_id"]==opt.exclude_single_obj
                except:
                    continue
                rgb, pt_cld, sample_2d, obj_id = data["img_crop"].cuda(), data["point_cloud"].cuda(), data["sample_2d"].cuda(), data["obj_id"].cuda()
                target, model_points = data["sampled_model_pt_camera"].cuda(), data["sampled_model_pt"].cuda()
                tar_r, tar_t = data["R"].cuda(), data["T"].cuda()

                if opt.recon_choice=='model':
                    recon_ref = model_points
                elif opt.recon_choice=='depth':
                    recon_ref = pt_cld
                else:
                    recon_ref = None

                rx, tx, cx, _, _, robust_loss = estimator(rgb, pt_cld, sample_2d, obj_id, recon_ref)
                
                if opt.loss == 'adds':
                    loss, dist = adds_loss(rx, tx, cx, target, model_points, obj_id, pt_cld, opt.w)
                elif opt.loss == 'abc':
                    loss, dist = abc_loss(rx, tx, cx, target, model_points, pt_cld, opt.w)
                else:
                    raise NotImplementedError
                    # loss, _ = gadd_loss(rx, tx, cx, tar_r, tar_t, train_dataset.prim_groups, obj_id, pt_cld, opt.w)
                
                loss_list.append(loss.item())
                dist_list.append(dist.item())
                robust_loss_list.append(robust_loss.item())
                
                # optimize
                loss = (loss + opt.recon_w*robust_loss)

                loss.backward()
                    
                if i % opt.optim_batch == 0:
                    for param_group in optimizer.param_groups:
                        param_group['lr'] = lr_schedule[i//opt.optim_batch]
                    grad_norm = get_grad_norm_(estimator.parameters())
                    optimizer.step()
                    # renew
                    tepoch.set_postfix(loss=loss.item(), grad_norm=grad_norm.item())
                    optimizer.zero_grad()
                    torch.cuda.synchronize()
                    torch.cuda.empty_cache()
                
        # save loss logs
        if len(loss_list) > 0:
            logger.log(f"Epoch: {epoch}, Avg Loss: {sum(loss_list)/len(loss_list)}, Avg Recon Loss: {sum(robust_loss_list)/len(robust_loss_list)}")
            writer.add_scalar("Loss/train", sum(loss_list)/len(loss_list), epoch)
        
        # test
        logger.log(f"start testing for epoch {epoch}....")
        estimator.eval()
        dist_list = []
        robust_loss_list = []
        torch.cuda.empty_cache()
        for i, data in enumerate(test_loader):
            try:
                if opt.single_obj_test>0:
                    assert data["obj_id"]==opt.single_obj_test
                if opt.exclude_single_obj>0:
                    assert not data["obj_id"]==opt.exclude_single_obj
            except:
                continue
            rgb, pt_cld, sample_2d, obj_id = data["img_crop"].cuda(), data["point_cloud"].cuda(), data["sample_2d"].cuda(), data["obj_id"].cuda()
            target, model_points = data["sampled_model_pt_camera"].cuda(), data["sampled_model_pt"].cuda()
            if opt.recon_choice=='model':
                recon_ref = model_points
            elif opt.recon_choice=='depth':
                recon_ref = pt_cld
            else:
                recon_ref = None
            rx, tx, cx, _, _, robust_loss = estimator(rgb, pt_cld, sample_2d, obj_id, recon_ref)
            _, dist = adds_loss(rx, tx, cx, target, model_points, obj_id, pt_cld, opt.w)
            dist_list.append(dist.item())
            robust_loss_list.append(robust_loss.item())
        dist_value = sum(dist_list)/len(dist_list)
        robust_loss = sum(robust_loss_list)/len(robust_loss_list)
        dist_2cm = 100.0 * len([d for d in dist_list if d<=0.02])/len(dist_list)
        dist_1cm = 100.0 * len([d for d in dist_list if d<=0.01])/len(dist_list)
        logger.log(f"testing done >>>> \n Average distance value: {dist_value}. \n Percentage of distance<1cm: {dist_1cm}.\n Percentage of distance<2cm: {dist_2cm}. \n Max dist value: {max(dist_list)}. \n Robust loss: {robust_loss} <<<<")
        writer.add_scalar("Distance/validation", dist_value, epoch)
        writer.add_scalar("Dist<1cm/validation", dist_1cm, epoch)
        writer.add_scalar("Dist<2cm/validation", dist_2cm, epoch)
        
        # save checkpoints
        saved = False
        if dist_value < best_test_dist:
            logger.log(f"Reach best test performance, save checkpoint to {cp_dir}")
            best_test_dist = dist_value
            saved = True
            torch.save(estimator.state_dict(), os.path.join(cp_dir, f"epoch_{epoch}_dist_{dist_value}.pth"))
        if dist_2cm > best_dist_2cm:
            logger.log(f"Reach best test performance, save checkpoint to {cp_dir}")
            best_dist_2cm = dist_2cm
            if not saved:
                torch.save(estimator.state_dict(), os.path.join(cp_dir, f"epoch_{epoch}_dist_{dist_value}.pth"))
            
        # update training settings
        if best_test_dist < opt.decay_margin:
            opt.lr *= opt.lr_rate
            opt.min_lr *= opt.lr_rate
            opt.w *= opt.w_rate
            opt.decay_margin *= opt.decay_rate
            logger.log(f"learning rate decay to {opt.lr}, minimum end learning rate decay to {opt.min_lr}, w decay to {opt.w}, decay_margin decay to {opt.decay_margin}")
            optimizer = optim.AdamW(estimator.parameters(), lr=opt.lr, weight_decay=opt.weight_decay)
        # end epoch
        logger.log(f"epoch {epoch} ends. \n\n")
    
    writer.close()
            

if __name__ == "__main__":
    main()
    
    

train.py
augcog/DTTD2
import os
import torch

def select_device(device='', batch_size=0, newline=True):
    # device = None or 'cpu' or 0 or '0' or '0,1,2,3'
    s = f'[Info] Getting Device: '
    device = str(device).strip().lower().replace('cuda:', '').replace('none', '')  # to string, 'cuda:0' to '0'
    cpu = device == 'cpu'
    mps = device == 'mps'  # Apple Metal Performance Shaders (MPS)
    if cpu or mps:
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # force torch.cuda.is_available() = False
    elif device:  # non-cpu device requested
        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable - must be before assert is_available()
        assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \
            f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"

    if not cpu and not mps and torch.cuda.is_available():  # prefer GPU if available
        devices = device.split(',') if device else '0'  # range(torch.cuda.device_count())  # i.e. 0,1,6,7
        n = len(devices)  # device count
        if n > 1 and batch_size > 0:  # check batch_size is divisible by device_count
            assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
        space = ' ' * (len(s) + 1)
        for i, d in enumerate(devices):
            p = torch.cuda.get_device_properties(i)
            s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n"  # bytes to MB
        arg = 'cuda:0'
    elif mps and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available():  # prefer MPS if available
        s += 'MPS\n'
        arg = 'mps'
    else:  # revert to CPU
        s += 'CPU\n'
        arg = 'cpu'

    if not newline:
        s = s.rstrip()
    
    print(s)
    return torch.device(arg)
utils/device.py
augcog/DTTD2
import os
import re

def get_checkpoint(model_path, isAcc=False):
    if os.path.isfile(model_path):
        return model_path
    elif os.path.isdir(model_path):
        checkpoints = sorted(os.listdir(model_path))
        record = []
        for cp in checkpoints:
            find = re.match(r'epoch_(\d+)_(.+)_([\d\.]+).pth', cp)
            if find == None:
                if cp == "best_save.pth":
                    return os.path.join(model_path, cp) # return best save directly
                print("[Warning] checkpoint name should be in epoch_(\d+)_dist_([\d\.]+).pth form")
            else:
                val = find.group(3)
                record.append((val, cp))
        if len(record) > 0:
            record = sorted(record, key=lambda x: x[0], reverse=isAcc)
            return os.path.join(model_path, record[0][1])
        else:
            return None
    else:
        return None
    
utils/file.py
augcog/DTTD2
import numpy as np

def get_discrete_width_bbox(label, border_list, img_w, img_h):
    rows = np.any(label, axis=1)
    cols = np.any(label, axis=0)
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    rmax += 1
    cmax += 1
    r_b = border_list[binary_search(border_list, rmax - rmin)]
    c_b = border_list[binary_search(border_list, cmax - cmin)]
    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > img_h:
        delt = rmax - img_h
        rmax = img_h
        rmin -= delt
    if cmax > img_w:
        delt = cmax - img_w
        cmax = img_w
        cmin -= delt
    return rmin, rmax, cmin, cmax

def discretize_bbox(rmin, rmax, cmin, cmax, border_list, img_w, img_h):
    rmax += 1
    cmax += 1
    r_b = border_list[binary_search(border_list, rmax - rmin)]
    c_b = border_list[binary_search(border_list, cmax - cmin)]
    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > img_h:
        delt = rmax - img_h
        rmax = img_h
        rmin -= delt
    if cmax > img_w:
        delt = cmax - img_w
        cmax = img_w
        cmin -= delt
    return rmin, rmax, cmin, cmax
    
def binary_search(sorted_list, target):
    l = 0
    r = len(sorted_list)-1
    while l!=r:
        mid = (l+r)>>1
        if sorted_list[mid] > target:
            r = mid
        elif sorted_list[mid] < target:
            l = mid + 1
        else:
            return mid
    return l
utils/image.py
augcog/DTTD2
class Logger(object):
    def __init__(self, filename):
        self.file = filename
        with open(self.file, 'w') as f:
            f.write("")
        
    def log(self, content):
        print(content)
        with open(self.file, 'a') as f:
            f.write(content + '\n')
utils/log.py
augcog/DTTD2
import cv2
import copy
import random
import numpy as np

def visualize(img, model_pts, R, T, intrinsics, color=(0, 255, 0), sample=1.0):
    """Render model points to the input image with (R,T) and camera intrinsics to visualize the results

    Args:
        img (ndarray [H,W,C]): input (original, not cropped) image that will be rendered on. Keep the intrinsics coherent.
        model_pts (ndarray [N, 3]): the points (x, y, z) on the model to be rendered.
        R (ndarray [3, 3]): Rotation of the model.
        T (ndarray [1, 3]): Translation of the model.
        intrinsics (tuple of size 4 | ndarray [3,3]): camera intrinsic parameters. Could be either tuple: (cam_cx, cam_cy, cam_fx, cam_fy) or intrinsic matrix.
        color (tuple of size 3): color of the rendered points.
        sample (float, 0.0~1.0): The sampling proportion of the model points.
        
    Return:
        rendered image
    """
    # valid check
    img, model_pts, R, T, intrinsics = _valid_check(img, model_pts, R, T, intrinsics)
    img = copy.deepcopy(img)
    
    # Transform model to world coordinates
    if sample < 1.0: # sample
        sample_list = list(range(len(model_pts)))
        sample_list = sorted(random.sample(sample_list, round(sample*len(model_pts))))
        model_pts = np.array(model_pts[sample_list, :])
    model_pts = np.add(np.dot(model_pts, R.T), T) # (N, 3)
    
    # Transform model to camera space and render
    cam_cx, cam_cy, cam_fx, cam_fy = intrinsics
    for pts in model_pts:
        x, y, z = pts
        px = round(x*cam_fx/z + cam_cx)
        py = round(y*cam_fy/z + cam_cy)
        if py<0 or py>=img.shape[0]:
            continue
        if px<0 or px>=img.shape[1]:
            continue
        img[py][px] = np.array(color)
    
    # save
    return img


def Draw3DBBox(img, model_pts, R, T, intrinsics, color=[255, 255, 255], thickness=2):
    """Render 3D bounding box of <model_pts> to the input image with (R,T) and camera intrinsics to visualize the results

    Args:
        img (ndarray [H,W,C]): input (original, not cropped) image that will be rendered on. Keep the intrinsics coherent.
        model_pts (ndarray [N, 3]): the points (x, y, z) of the target object CAD model (local coordinate).
        R (ndarray [3, 3]): Rotation of the model.
        T (ndarray [1, 3]): Translation of the model.
        intrinsics (tuple of size 4 | ndarray [3,3]): camera intrinsic parameters. Could be either tuple: (cam_cx, cam_cy, cam_fx, cam_fy) or intrinsic matrix.
        color (tuple of size 3): color of the rendered bbox cube frame.
        thickness (int): The thickness of the rendered bbox cube frame (pixel).
        
    Return:
        rendered image
    """
    # valid check
    _valid_check(img, model_pts, R, T, intrinsics)
    
    # get 3D BBox
    xmax = np.max(model_pts[:, 0])
    ymax = np.max(model_pts[:, 1])
    zmax = np.max(model_pts[:, 2])
    xmin = np.min(model_pts[:, 0])
    ymin = np.min(model_pts[:, 1])
    zmin = np.min(model_pts[:, 2])
    Vertices = np.array([[xmax, ymax, zmax], [xmax, ymax, zmin],\
                        [xmin, ymax, zmax], [xmin, ymax, zmin],\
                        [xmax, ymin, zmax], [xmax, ymin, zmin],\
                        [xmin, ymin, zmax], [xmin, ymin, zmin]])
    
    # Transform
    Vertices = np.add(np.dot(Vertices, R.T), T)
    cam = intrinsics
    cam_cx, cam_cy, cam_fx, cam_fy =  cam[0][2], cam[1][2], cam[0][0], cam[1][1]
    pixels = []
    for v in Vertices:
        x, y, z = v
        px = round(x*cam_fx/z + cam_cx)
        py = round(y*cam_fy/z + cam_cy)
        pixels.append([px, py])
    # Cube edges
    edges = [[0,1], [0,2], [0,4], [1,3], [1,5], [2,3], [2,6], [3,7], [4,5], [4,6], [5,7], [6,7]]
    # Draw Cube
    for edge in edges:
        img = cv2.line(img, pixels[edge[0]], pixels[edge[1]], color, thickness)
    return img


def _valid_check(img, model_pts, R, T, intrinsics):
    img = np.array(img)
    assert img.shape[2] == 3
    model_pts = np.array(model_pts)
    assert model_pts.shape[1] == 3
    R = np.array(R)
    assert R.shape == (3,3)
    T = np.array(T)
    assert T.shape == (1,3)
    if len(intrinsics) == 3 and len(intrinsics[0]) == 3:
        intrinsics = (intrinsics[0][2], intrinsics[1][2], intrinsics[0][0], intrinsics[1][1])
    intrinsics = np.array(intrinsics)
    assert(intrinsics.shape == (4, ))
    return img, model_pts, R, T, intrinsics
        
    
    
utils/visualizer.py
augcog/DTTD2