def _debug(self, message): """ shutdown curses, print message, then exit """ self._shutdown() Tools.pyout(message) sys.exit(0)
def _dimred_fit(self, data, n_components=2): Tools.pyout("FITTING DIMENSIONALITY REDUCTION MODEL") dimred = umap.UMAP() dimred.fit(data) data_ = dimred.transform(data) minmax = max(abs(np.min(data_)), abs(np.max(data_))) self.transform = lambda x: x / minmax Tools.pyout("--------------------------------> DONE") return dimred
def _fwd(self, X): with torch.no_grad(): with autograd.detect_anomaly(): try: X = X.to(self.devices[0]) y = self.tcn(X).detach().cpu().numpy() except Exception as e: os.system('clear') Tools.pyout(e, force=True) sys.exit(0) return y
def _cluster(self, dist_matrix, n_clusters=15): Tools.pyout("CLUSTERING") model = AgglomerativeClustering(affinity='precomputed', linkage='average', n_clusters=n_clusters) model.fit(dist_matrix) colors_ = [] for lbl in model.labels_: colors_.append(Tools.num2hslstr(lbl / max(model.labels_))) Tools.pyout("----> DONE") return colors_
def _reduce(self, dist_matrix, n_components=2): Tools.pyout("FITTING UMAP...") with warnings.catch_warnings(): warnings.simplefilter("ignore") model = manifold.MDS(n_components=n_components, metric='precomputed') reduced = np.array(model.fit_transform(dist_matrix)) Tools.pyout("---------> DONE") reduced -= np.min(reduced, axis=0) reduced /= np.max(reduced, axis=0) reduced = (reduced - 0.5) * 2 return reduced.tolist()
def __getitem__(self, idx): paths = self.batch_paths[idx] X = np.zeros((len(paths),) + self.input_size, dtype=np.float32) for ii, path in enumerate(paths): img = cv2.imread(path) if np.isnan(img).any(): Tools.pyout('nan in', path) img = Transformer.transform(cv2.imread( path), augment=False, BGR=False) X[ii, :, :, :] = self.transform(Image.fromarray(img)).numpy() if np.isnan(X).any(): Tools.pyout('nan in tensor') return torch.FloatTensor(X), paths
def _load(self, path=None): if path is not None and os.path.isfile(path): try: if torch.cuda.is_available(): self.load_state_dict(torch.load(path)) else: self.load_state_dict( torch.load(path, map_location=lambda storage, loc: storage)) Tools.log("Load VAE: Success") return True except Exception as e: Tools.pyout("Load VAE: Fail " + e) Tools.log("Load VAE: Fail " + e) return False return False
def train(self, epoch): self.tcn.switch_mode('train') epoch_loss = 0. cum_loss = 0. cum_ratio_loss = 0. n = 0 for idx in poem(range(len(self.train_set)), "train " + str(epoch)): X, labels, perspectives, paths = self.train_set[idx] if X is None: continue n += 1 X = X.to(self.devices[0]) labels = labels.to(self.devices[1]) perspectives = perspectives.to(self.devices[1]) self.optimizer.zero_grad() y = self.tcn(X) assert not Tools.contains_nan(y) loss = Loss.triplet_semihard_loss(y, labels, perspectives, margin=Config.TCN_MARGIN, device=self.devices[1]) loss.backward() self.optimizer.step() epoch_loss += loss.item() cum_loss += Loss.embedding_accuracy(y, labels, perspectives, device=self.devices[1]).item() cum_ratio_loss += Loss.embedding_accuracy_ratio( y, labels, perspectives) Tools.pyout('Train Epoch: ' + '{} [{}/{} ({:.0f}%)]\tAccuracy: ' '{:.6f}\tRatio: {:.6f}\tLoss: {:.6f}'.format( epoch, idx, len(self.train_set), 100. * idx / len(self.train_set), cum_loss / n, cum_ratio_loss / n, epoch_loss / (n))) Tools.log('Train Epoch: ' + '{} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, n, len(self.train_set), 100., epoch_loss / (n)))
def _write_csv(self, data_root, trials_embeddings, reduce_embd, labels): Tools.pyout("WRITING CSV FILE") SAVE_ROOT = '/media/roblaundry/' + \ data_root.split('/')[-2] + '/results/BT' Tools.makedirs(SAVE_ROOT) with open(os.path.join(SAVE_ROOT, 'data.csv'), 'w+') as csv_file: fieldnames = ['trial_name', 'x', 'y', 'c'] writer = csv.DictWriter(csv_file, fieldnames=fieldnames) writer.writeheader() for ii, embedding in enumerate(trials_embeddings): writer.writerow({ 'trial_name': Tools.fname(embedding[0]), 'x': reduce_embd[ii][0], 'y': reduce_embd[ii][1], 'c': labels[ii].replace(',', ';') }) Tools.pyout("----------> DONE")
def copy(self, root, in_pos, ou_pos): """ extract frames from video files Args: root: string - path to store video files in_pos: string - camera perspective (e.g. left, middle, right) ou_pos: string - name for camera perspective directory in dataset """ for trial in Tools.tqdm_pbar(Tools.list_dirs(root), 'COPYING'): Tools.pyout(trial) fname = '_'.join(Tools.fname(trial).split('_')[:-1]) Tools.makedirs(os.path.join(trial, ou_pos)) vid_folder = self._find(fname) frames = [] for frame in Tools.list_files(os.path.join(trial, in_pos)): frames.append(int(Tools.fname(frame).split('.')[0])) path = os.path.join(vid_folder, 'color-recording-' + ou_pos + '.avi') if not os.path.isfile(path): path = path.replace('.avi', '-x265.mp4') cap = cv2.VideoCapture(path) if not cap.isOpened(): Tools.pyout("ERROR OPENING VideoCapture") raise FileNotFoundError('in copy(): "' + path + '"') ii = 0 with tqdm(total=len(frames)) as pbar: while cap.isOpened(): ret, frame = cap.read() if ret is True: if ii in frames: img = np.copy(frame) cv2.imwrite( os.path.join(trial, ou_pos, str(ii).zfill(5) + '.jpg'), img) pbar.update(1) else: break ii += 1
def test(self, epoch): self.tcn.switch_mode('eval') cum_loss = 0. cum_ratio_loss = 0. n = 0 for idx in poem(range(len(self.val_set)), "eval " + str(epoch)): # if idx == 5: # break with torch.no_grad(): batch = self.val_set[idx] X, labels, perspectives = batch[0], batch[1], batch[2] if X is None: continue n += 1 X = X.to(self.devices[0]) labels = labels.to(self.devices[1]) perspectives = perspectives.to(self.devices[1]) y = self.tcn(X) assert not Tools.contains_nan(y) cum_loss += Loss.embedding_accuracy( y, labels, perspectives, device=self.devices[1]).item() cum_ratio_loss += Loss.embedding_accuracy_ratio( y, labels, perspectives) Tools.pyout('Test Epoch: ' + ( '{} [{}/{} ({:.0f}%)]\tAccuracy: ' '{:.6f}\tRatio: {:.6f}').format(epoch, n, len(self.val_set), 100. * idx / len(self.val_set), cum_loss / (n), cum_ratio_loss / (n))) Tools.log('Test Epoch: ' + ('{} [{}/{} ({:.0f}%)]\tAccuracy: ' '{:.6f}\tRatio: {:.6f}' ).format(epoch, n, len(self.val_set), 100., cum_loss / (n), cum_ratio_loss / (n))) if cum_ratio_loss / (n) < self.best_ratio: self.best_ratio = cum_ratio_loss / (n) self.save_state_dict(self.save_loc) self.last_improvement = epoch
def train(self, epoch): self.vae.train() epoch_loss = 0. for ii in poem(range(min(500, len(self.train_set))), "TRAIN EPOCH " + str(epoch)): X, t, _ = self.train_set[ii] X, t = X.to(self.device), t.to(self.device) y, mu, logvar = self.vae(X) self.optimizer.zero_grad() BCE = F.binary_cross_entropy((1.0 + y.reshape(X.shape[0], -1)) / 2, (1.0 + t.reshape(X.shape[0], -1)) / 2, reduction='sum') logvar = torch.clamp(logvar, -10., 10.) mu = torch.clamp(mu, -10., 10.) KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) BCE /= X.shape[0] * X.shape[1] * X.shape[2] * X.shape[3] KLD /= X.shape[0] * X.shape[1] * X.shape[2] * X.shape[3] loss = BCE + KLD if loss < 12.0: loss.backward() self.optimizer.step() # else: # Tools.pyout("Skip: ") kld = KLD.item() # Tools.pyout(BCE.item(), KLD.item()) if kld != kld: sys.exit(0) epoch_loss += loss.item() Tools.pyout('Train Epoch: ' + '{} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, ii, len(self.train_set), 100. * ii / len(self.train_set), epoch_loss / ii))
def visualize(self, in_folder, ou_folder): Tools.makedirs(ou_folder) data = self._load_embedding_dicts(in_folder) main_frm = np.zeros((480, len(data) * 240, 3), dtype=np.uint8) N = data[0][2].shape[0] # define codec and create VideoWriter object fourcc = cv2.VideoWriter_fourcc(*'mp4v') writer = cv2.VideoWriter( os.path.join(ou_folder, in_folder.split('/')[-1] + '.mp4'), fourcc, 16, (240 * len(data), 480)) Tools.pyout(os.path.join(ou_folder, in_folder.split('/')[-1] + '.mp4')) for ii, pos_data in enumerate(data): data[ii] = self._cluster(*pos_data) for n in poem(range(N), Tools.fname(in_folder)): for ii, pos_data in enumerate(data): fldr, frms, embd, rdcd, lbls, plot = pos_data # load frame image img_pth = os.path.join(fldr, frms[n]) frame = cv2.imread(img_pth) frame = cv2.resize(frame, (240, 240)).astype(np.uint8) # draw cluster color in frame color = plot._prog2color(lbls[n]) cv2.rectangle(frame, (0, 0), (frame.shape[1], frame.shape[0]), color, 10) # draw embeddings plt_img = plot.plot(rdcd[n], color=color, scope=False) # add subfigures to frame main_frm[0:240, ii * 240:(ii + 1) * 240, :] = frame[:, :, :] main_frm[240:480, ii * 240:(ii + 1) * 240, :] = \ plt_img[:, :, :] # Tools.render(main_frm, waitkey=50, name="SEGMENTATION") writer.write(main_frm) writer.release()
def __init__(self, device=None, state_dict_path=None, train_root=None, val_root=None, pos=None, save_loc=None): self.device = device self.state_dict_path = state_dict_path assert pos is not None self.pos = pos self.save_loc = save_loc self.train_set = VAESet(root_dir=train_root, pos=pos, augment=True) self.val_set = VAESet(root_dir=val_root, pos=pos, augment=False) self.vae = VAE().to(device) self.optimizer = optim.Adam(list( filter(lambda p: p.requires_grad, self.vae.parameters())), lr=1e-4) Tools.pyout('VAETrainer loaded')
def eval(self, epoch): self.vae.eval() epoch_loss = 0. with torch.no_grad(): for ii in poem(range(len(self.val_set)), "EVAL EPOCH " + str(epoch)): X, t, _ = self.val_set[ii] X, t = X.to(self.device), t.to(self.device) y, _, _ = self.vae(X) loss = F.binary_cross_entropy( (1.0 + y.reshape(X.shape[0], -1)) / 2, (1.0 + t.reshape(X.shape[0], -1)) / 2, reduction='sum') epoch_loss += loss.item() Tools.pyout('Eval Epoch: ' + '{} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, ii, len(self.val_set), 100. * ii / len(self.val_set), epoch_loss / ii)) if loss < self.best_loss: self.best_loss = loss self.vae.save_state_dict(self.save_loc)
VAE_VIZUA = 1 # visualize vae reconstruction based on alignments VZUAL_REW = 1 # visuzalize reward progression HANDBRAKE = 1 # handbrake-cli to make outputs compatible with browser video if VIZ_EMBED: # can't do this without fitting to find scaling hyperparameters V_RED_FIT = True # select datasets roots = [] roots.append('/media/roblaundry/folding_full') roots.append('/media/roblaundry/folding_single') roots.append('/media/roblaundry/grasping') roots.append('/media/roblaundry/grasping_reverse') roots.append('/media/roblaundry/pouring') os.system('clear') Tools.pyout("STARTING PIPELINE") devices = (torch.device("cuda:0" if torch.cuda.is_available() else "cpu"), torch.device("cuda:0" if torch.cuda.is_available() else "cpu")) for root in roots: if Config.SERVER: save_root = Tools.fname(root) # train TCN on specified dataset if TRAIN_TCN: Tools.pyout("TRAIN_TCN") STATE_DICT_PATHS = (None, None) TRAIN_ROOT = os.path.join(root, 'train') VAL_ROOT = os.path.join(root, 'val') SAVE_LOC = os.path.join(save_root, 'tcn')
def visualize(self, in_folder, ou_folder): try: os.makedirs(ou_folder) except FileExistsError: pass # Tools.pyout(ou_folder) N_perspectives = len(Tools.list_dirs(in_folder)) folder_name = in_folder.split('/')[-1] # define codec and create VideoWriter object fourcc = cv2.VideoWriter_fourcc(*'mp4v') writer = cv2.VideoWriter(os.path.join( ou_folder, in_folder.split('/')[-1] + '.mp4'), fourcc, 16, (240 * N_perspectives, 480)) Tools.pyout(os.path.join( ou_folder, in_folder.split('/')[-1] + '.mp4')) # init frame main_frm = np.zeros((480, 240 * N_perspectives, 3), dtype=np.uint8) # init plots plots = self._init_plots(N_perspectives) # load embedding dicts dicts = self._load_embedding_dicts(in_folder) # loop over all frames frames_exec = ( 'sorted(list(' + ' | '.join('set(dicts[{}])'.format(ii) for ii in range(len(dicts))) + '))') frames = eval(frames_exec) max_frm = int(frames[-1].split('.')[0]) min_frm = int(frames[0].split('.')[0]) for frame in poem(frames, folder_name): for ii, pos in enumerate(sorted(Tools.list_dirs(in_folder))): try: dic = dicts[ii] side_str = pos.split('/')[-1] lpx = 240 * (ii) rpx = 240 * (ii + 1) plot = plots[ii] # add frame from perspective to main frame frame_img = cv2.imread( os.path.join(in_folder, side_str, frame)) frame_img = cv2.resize(frame_img, (240, 240)) main_frm[0:240, lpx:rpx, :] = np.copy( frame_img[:, :, :3]) # get dimred from TCN embedding of frame embedding = dic[frame] with warnings.catch_warnings(): warnings.simplefilter("ignore") projection = tuple( self.transform(self.reduction_model.transform( np.array([embedding])).squeeze())) # plot embedding in plot frame plt_img = plot.plot(projection, prog=(int( frame.split('.')[0]) - min_frm) / (max_frm - min_frm)) # add plot to main frame main_frm[240:480, lpx:rpx, :] = np.copy( plt_img[:, :, :3]) except cv2.error: pass # Tools.render(main_frm, ascii=True) writer.write(main_frm) writer.release()