def __init__(self, checkpoint_dir, log_dir):
print('------ CREATING NEW MODEL ------')
print(hyp.name)
self.checkpoint_dir = checkpoint_dir
self.log_dir = log_dir
self.all_inputs = inputs.get_inputs()
print("------ Done getting inputs ------")
if hyp.moc:
self.poolvox_moc = MOC_QUEUE_NORMAL(hyp.moc_qsize)
if hyp.moc_2d:
self.poolvox_moc_2d = MOC_QUEUE_NORMAL(hyp.moc_qsize)
if hyp.offline_cluster:
self.cluster_pool = ClusterPool(hyp.offline_cluster_pool_size)
if hyp.offline_cluster_eval:
self.info_dict = defaultdict(lambda:[])
if hyp.max.hardmining or hyp.hard_eval or hyp.hard_vis:
self.mbr = cross_corr.meshgrid_based_rotation(hyp.BOX_SIZE-2*hyp.max.margin, hyp.BOX_SIZE-2*hyp.max.margin, hyp.BOX_SIZE-2*hyp.max.margin)
self.mbr16 = cross_corr.meshgrid_based_rotation(hyp.BOX_SIZE, hyp.BOX_SIZE, hyp.BOX_SIZE)
self.mbr_unpr = cross_corr.meshgrid_based_rotation(32,32,32)
self.hpm = hardPositiveMiner.HardPositiveMiner(self.mbr,self.mbr16,self.mbr_unpr)
elif hyp.do_orientation:
self.mbr16 = cross_corr.meshgrid_based_rotation(hyp.BOX_SIZE, hyp.BOX_SIZE, hyp.BOX_SIZE)
self.mbr_unpr = cross_corr.meshgrid_based_rotation(32,32,32)
self.hpm = None
else:
self.mbr16 = None
self.hpm = None
if hyp.self_improve_iterate:
self.pool_det = DetPool(hyp.det_pool_size)
if hyp.do_eval_recall:
self.eval_dicts = {}
if hyp.dataset_name=="bigbird":
self.recalls = [3, 5, 10]
else:
self.recalls = [10, 20, 30]
if hyp.do_debug or hyp.low_dict_size:
self.pool_size = 11
else:
self.pool_size = hyp.pool_size
self.eval_dicts["eval"] = {}
F = 3
if hyp.eval_recall_o:
self.eval_dicts["eval"]['pool3D_e'] = DoublePool_O(self.pool_size)
self.eval_dicts["eval"]['pool3D_g'] = DoublePool_O(self.pool_size)
else:
self.eval_dicts["eval"]['pool3D_e'] = DoublePool(self.pool_size)
self.eval_dicts["eval"]['pool3D_g'] = DoublePool(self.pool_size)
self.eval_dicts["eval"]['precision3D'] = np.nan*np.array([0.0, 0.0, 0.0], np.float32)
self.eval_dicts["eval"]['neighbors3D'] = np.zeros((F*10, F*11, 3), np.float32)
self.device = torch.device("cuda")
def __init__(self, num_embeddings, embedding_dim, init_embeddings, commitment_cost):
super(VectorQuantizer, self).__init__()
self.embedding_dim = embedding_dim
self.num_embeddings = num_embeddings
if hyp.from_sup:
embeddings_pickled = pickle.load(open(hyp.object_quantize_sup_init,"rb"))
self.num_embeddings = num_embeddings*hyp.num_classes
self.suffix = hyp.suffix
self.labels = hyp.labels
self.embeddings = torch.nn.Embedding(self.num_embeddings,self.embedding_dim)
all_embeds = []
for label in self.labels:
embed = embeddings_pickled[label]
all_embeds.append(torch.cat(embed,dim=0))
all_embeds = torch.cat(all_embeds,dim=0)
self.embeddings.cuda()
self.embeddings.weight.data = all_embeds
else:
self.embeddings = torch.nn.Embedding(self.num_embeddings,
self.embedding_dim)
self.embeddings.cuda()
if init_embeddings is not None:
print("LOADING embeddings FROM NUMPY: ",init_embeddings)
cluster_centers = np.load(init_embeddings)
self.embeddings.weight.data = torch.from_numpy(cluster_centers).to(torch.float32).cuda()
else:
assert False
print("we should never intialize")
limit = 1/self.num_embeddings
self.embeddings.weight.data.uniform_(-limit,+limit)
self.commitment_cost = commitment_cost
self.closest_distances_mean = []
if hyp.vq_rotate:
self.mbr = cross_corr.meshgrid_based_rotation(hyp.BOX_SIZE,hyp.BOX_SIZE,hyp.BOX_SIZE)
def __init__(self, num_embeddings, embedding_dim, init_embeddings, commitment_cost):
super(VectorQuantizer_Supervised, self).__init__()
self.embedding_dim = embedding_dim
self.num_embeddings = num_embeddings
self.embeddings_dict = {}
embeddings_pickled = pickle.load(open(hyp.object_quantize_sup_init,"rb"))
self.suffix = hyp.suffix
self.labels = hyp.labels
self.embeddings = torch.nn.Embedding(len(self.labels),self.num_embeddings*self.embedding_dim)
all_embeds = []
for label in self.labels:
embed = embeddings_pickled[label]
all_embeds.append(torch.cat(embed,dim=0).reshape([-1]))
all_embeds = torch.stack(all_embeds,dim=0)
self.embeddings.cuda()
self.embeddings.weight.data = all_embeds
# self.embedding_cube = torch.nn.Embedding(self.num_embeddings,self.embedding_dim)
# self.embedding_cube.cuda()
# self.embedding_cube.weight.data = torch.cat(embeddings_pickled['cube'],dim=0).to(torch.float32).cuda()
# self.embedding_cylinder = torch.nn.Embedding(self.num_embeddings,self.embedding_dim)
# self.embedding_cylinder.cuda()
# self.embedding_cylinder.weight.data = torch.cat(embeddings_pickled['cylinder'],dim=0).to(torch.float32).cuda()
# self.embedding_sphere = torch.nn.Embedding(self.num_embeddings,self.embedding_dim)
# self.embedding_sphere.cuda()
# self.embedding_sphere.weight.data = torch.cat(embeddings_pickled['sphere'],dim=0).to(torch.float32).cuda()
# for key,val in embeddings_pickled.items():
# embeddings_ind = torch.nn.Embedding(self.num_embeddings,self.embedding_dim)
# embeddings_ind.cuda()
# embeddings_ind.weight.data = torch.cat(val,dim=0).to(torch.float32).cuda()
# self.embeddings[key] = embeddings_ind
# self.embeddings = {}
# self.embeddings['cube'] = self.embeddings
# self.embeddings['cylinder'] = self.embedding_cylinder
# self.embeddings['sphere'] = self.embedding_sphere
self.commitment_cost = commitment_cost
if hyp.vq_rotate:
self.mbr = cross_corr.meshgrid_based_rotation(hyp.BOX_SIZE,hyp.BOX_SIZE,hyp.BOX_SIZE)
def __init__(self, num_embeddings, embedding_dim, init_embeddings, commitment_cost): super(VectorQuantizer_Instance_Vr, self).__init__() self.embedding_dim = embedding_dim self.num_embeddings = num_embeddings self.embeddings = torch.nn.Embedding(self.num_embeddings, self.embedding_dim) self.embeddings.cuda() self.genVar = ResNet3D_NOBN(hyp.feat_dim,hyp.feat_dim) self.genVar.cuda() if init_embeddings is not None: cluster_centers = np.load(init_embeddings) self.embeddings.weight.data = torch.from_numpy(cluster_centers).to(torch.float32).cuda() else: limit = 1/self.num_embeddings self.embeddings.weight.data.uniform_(-limit,+limit) self.commitment_cost = commitment_cost if hyp.vq_rotate: self.mbr = cross_corr.meshgrid_based_rotation(hyp.BOX_SIZE,hyp.BOX_SIZE,hyp.BOX_SIZE)
def __init__(self):
super(VectorQuantizer_Eval, self).__init__()
# init_dir = os.path.join("checkpoints", feat_init)
# ckpt_names = os.listdir(init_dir)
# steps = [int((i.split('-')[1]).split('.')[0]) for i in ckpt_names]
# step = max(steps)
# model_name = 'model-%d.pth'%(step)
# path = os.path.join(init_dir, model_name)
if hyp.use_instances_variation_all:
self.genVar = ResNet3D_NOBN(hyp.feat_dim,hyp.feat_dim)
self.genVar.cuda()
self.embedding_dim = hyp.BOX_SIZE*hyp.BOX_SIZE*hyp.BOX_SIZE*hyp.feat_dim
self.num_embeddings = hyp.object_quantize_dictsize
self.embeddings = torch.nn.Embedding(self.num_embeddings,
self.embedding_dim)
self.embeddings.cuda()
# st()
# loaded_embeddings = torch.load(path)['model_state_dict']['quantizer.embeddings.weight']
# self.embeddings.weight.data = loaded_embeddings
if hyp.vq_rotate:
self.mbr = cross_corr.meshgrid_based_rotation(hyp.BOX_SIZE,hyp.BOX_SIZE,hyp.BOX_SIZE)
def __init__(self, num_embeddings, embedding_dim, init_embeddings, commitment_cost):
super(VectorQuantizer_Supervised_Evaluator, self).__init__()
self.embedding_dim = embedding_dim
self.embeddings = {}
self.suffix = hyp.suffix
self.labels = hyp.labels
self.num_embeddings = len(self.labels)*num_embeddings
# embeddings_pickled = pickle.load(open(hyp.object_quantize_sup_init,"rb"))
# all_embeds = []
# for key,val in embeddings_pickled.items():
# all_embeds.append(torch.cat(val,dim=0))
# self.all_embeds = torch.cat(all_embeds,dim=0)
# for key,val in embeddings_pickled.items():
# embeddings_ind = torch.nn.Embedding(self.num_embeddings,self.embedding_dim)
# embeddings_ind.cuda()
# embeddings_ind.weight.data = torch.cat(val,dim=0).to(torch.float32).cuda()
# self.embeddings[key] = embeddings_ind
# st()
self.commitment_cost = commitment_cost
if hyp.vq_rotate:
self.mbr = cross_corr.meshgrid_based_rotation(hyp.BOX_SIZE,hyp.BOX_SIZE,hyp.BOX_SIZE)
import cross_corr
import torch
import torch.nn.functional as F
import cv2
import ipdb
from scipy.misc import imsave
st = ipdb.set_trace
img = torch.tensor(cv2.imread('new.jpg') / 255.0).permute(
2, 0, 1).unsqueeze(0).cuda()
img_size = 100
img = F.interpolate(img, [img_size, img_size]).to(torch.float32)
orig = img.squeeze(0).permute(1, 2, 0).cpu().numpy()
imsave(f"rotated/original.png", orig)
mbr = cross_corr.meshgrid_based_rotation(img_size,
img_size,
img_size,
angleIncrement=10.0)
rotated_inputs = mbr.rotate2D(img, interpolation="bilinear")
rotated_inputs = rotated_inputs.permute(0, 2, 1, 3, 4)
# for index in range(36):
# rotated_inputs_i = rotated_inputs[0,index]
# rotated_inputs_i_np = rotated_inputs_i.permute(1,2,0).cpu().numpy()
# imsave(f"rotated/{index}.png",rotated_inputs_i_np)
B, angles, C, H, W = list(rotated_inputs.shape)
C = C * H * W
# assert(C==embedding_size)
rot_input = rotated_inputs.reshape(B, angles, -1)
index_to_take = 5
codebook = rot_input[0, index_to_take:index_to_take + 1]
import torch
from torch.autograd import Function
#<<<<<<< HEAD
import ipdb
st = ipdb.set_trace
#=======
import cross_corr
mbr = cross_corr.meshgrid_based_rotation(
16, 16, 16) # TODO: change '16' to actual size.
#>>>>>>> dc4fab3cb9f342258faa7c388190a35e4a29eccd
class VectorQuantization(Function):
@staticmethod
def forward(ctx, inputs, codebook, object_level):
'''
TODO: Making following assumptions for rotation
inputs is of size C x H x W
'''
with torch.no_grad():
# assuming C,H,W. So probably will need to fix that after shape validation
# This flatenning only makes sense when embedding_size is C*H*W.
# Otherwise we'll have to permute the rotated_inputs probably before flatenning.
embedding_size = codebook.size(1)
dictionary_size = codebook.size(0)
codebook_sqr = torch.sum(codebook**2, dim=1)
if object_level:
B, _, _, _ = list(inputs.shape)
if True:
def __init__(self, K, D, object_level):
super().__init__()
self.embedding = nn.Embedding(K, D)
self.object_level = object_level
self.embedding.weight.data.uniform_(-1. / K, 1. / K)
self.mbr = cross_corr.meshgrid_based_rotation(16, 16, 16)