def embed(self, force=False): self.tcn.switch_mode('eval') for root in poem(sorted(self.roots, reverse=True), "EMBEDDING"): for trial in poem(Tools.list_dirs(root), Tools.fname(root)): for pos_path in poem(Tools.list_dirs(trial), Tools.fname(trial)): pos = pos_path.split('/')[-1] dataset = EmbedderSet(root_dir=pos_path) embeddings = {} for X, paths in poem(dataset, Tools.fname(pos_path)): if len(paths) > 0: y = self._fwd(X) if not np.isfinite(y).all(): pyout('before', paths) sys.exit(0) for ii, path in enumerate(paths): embeddings[path.split('/')[-1]] = np.copy( y[ii, :]).tolist() for key in embeddings: if np.isnan(embeddings[key]).any(): pyout('after', pos_path, key) sys.exit(0) with open(os.path.join(trial, 'embed_' + pos + '.json'), 'w+') as f: json.dump(embeddings, f, indent=1) del dataset del embeddings
def __init__(self, root_dir=None, pos=None, batch_size=Config.TCN_BATCH, input_size=(3, 224, 224), output_size=(3, 224, 224), augment=False): self.root = root_dir assert pos is not None self.batch_size = batch_size self.input_size = input_size self.output_size = output_size self.augment = augment frames = [] for trial_folder in Tools.list_dirs(root_dir): for frame_pth in Tools.list_files( os.path.join(trial_folder, pos), end='.jpg'): frames.append(frame_pth) if augment: random.shuffle(frames) self.batches = [[]] for frame in poem(frames, 'LOADING ' + Tools.fname(root_dir)): if len(self.batches[-1]) >= batch_size: self.batches.append([frame]) else: self.batches[-1].append(frame)
def _load_embedding_dicts(self, folder): try: out = [] for json_filename in sorted( Tools.list_files(folder, end='.json', substr='embed_')): pos = Tools.fname(json_filename).replace('embed_', '').replace( '.json', '') root_folder = os.path.join(folder, pos) with open(json_filename, 'r') as f: D = json.load(f) X = np.zeros((len(D), len(D[list(D)[0]])), dtype=float) frames = sorted(list(D)) for ii, frame in enumerate(frames): X[ii, :] = np.array(D[frame]) out.append((root_folder, frames, X, self._reduce(X))) if self.averaged: X_avg = np.stack([pos_data[2] for pos_data in out], axis=0) X_avg = np.mean(X_avg, axis=0) for ii, (fldr, frms, _, rdcd) in enumerate(out): out[ii] = (fldr, frms, X_avg, rdcd) return out except Exception as e: Tools.debug(e) Tools.debug(folder, ex=0)
def _load_images(self, path): """ read images into working memoryk Args: path: string - input path Returns: imgs: list(np.array) - list of frames from 'left' perspective N: int - number of frames """ N = len(Tools.list_files(os.path.join(path, 'left'))) imgs = np.zeros((N, 240, 360, 3), dtype=np.uint8) load_path = os.path.join(path, 'left') for fii, frm_pth in Tools.tqdm_pbar(enumerate( Tools.list_files(load_path)), path, total=N): if self.zfill_n is None: self.zfill_n = len(Tools.fname(frm_pth).split('.')[0]) img = cv2.imread(frm_pth) img = cv2.resize(img, (360, 240)) imgs[fii, :, :, :] = np.copy(img[:, :, :]) return imgs, N
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 visualize(self, in_folder, ou_folder): Tools.makedirs(ou_folder) nn = self._nearest_neighbor(in_folder) fourcc = cv2.VideoWriter_fourcc(*'mp4v') writer = cv2.VideoWriter( os.path.join(ou_folder, in_folder.split('/')[-1] + '.mp4'), fourcc, 16, (480, 480)) frames = [ Tools.fname(f) for f in Tools.list_files(os.path.join(in_folder, self.POS), end='.jpg') ] cv_plot = CVLine((240, 480), minx=0, maxx=len(frames), miny=0, maxy=1) main_frm = np.zeros((480, 480, 3), dtype=np.uint8) for fii, frame in Tools.tqdm_pbar(enumerate(frames), Tools.fname(in_folder), total=len(frames)): frame = cv2.imread(os.path.join(in_folder, self.POS, frame)) try: ancho = cv2.imread( os.path.join(nn, self.POS, self.alignments[in_folder][nn][fii][0])) except Exception as e: Tools.debug(self.alignments[in_folder][nn]) Tools.debug(e, ex=0) frame = cv2.resize(frame, (240, 240)) ancho = cv2.resize(ancho, (240, 240)) main_frm[:240, :240, :] = frame[:, :, :] main_frm[:240, 240:, :] = ancho[:, :, :] R = self._reward_score(in_folder, fii) main_frm[240:, :, :] = cv_plot.plot((fii, R)) writer.write(main_frm) writer.release()
def _reward_score(self, trial, index): cum_gain = 0. cum_weight = 0. for match_trial in self.alignments[trial]: match_frm, weight = self.alignments[trial][match_trial][index] match_nmr = int(''.join(filter(str.isdigit, match_frm))) try: match_min = min( int(''.join(filter(str.isdigit, Tools.fname(f)))) for f in Tools.list_files( os.path.join(match_trial, self.POS), end='.jpg')) match_max = max( int(''.join(filter(str.isdigit, Tools.fname(f)))) for f in Tools.list_files( os.path.join(match_trial, self.POS), end='.jpg')) except Exception as e: Tools.debug(os.path.join(Tools.pathstr(match_trial), self.POS)) Tools.debug(e, ex=0) cum_gain += (match_nmr - match_min) / (match_max - match_min) cum_weight += 1 return cum_gain / cum_weight
def __init__(self, root=None): self.POS = tuple([ Tools.fname(f) for f in Tools.list_dirs(Tools.list_dirs(root)[0]) ]) Tools.debug(self.POS) if 'steady' in self.POS: self.POS = 'steady' elif 'middle' in self.POS: self.POS = 'middle' else: self.POS = self.POS[0] am = AlignMatrix(root) self.alignments = am.load()
def walk_dirs(self, ou_path): """ find trials and store Args: ou_path: string - path to output file """ recording_paths = [] for f0 in Tools.tqdm_pbar(Tools.list_dirs(self.root), "GATHERING"): fname = Tools.fname(f0) if fname == '2017-11-09-to-encode': for f1 in Tools.list_dirs(f0): for f2 in Tools.list_dirs(f1): if self._valid_folder(f2): if self._length(f2, encoded=False) >= 60: recording_paths.append(f2) elif fname == '2017-11-06-still some parts to compress': for f1 in Tools.list_dirs(f0): if self._valid_folder(f1): if self._length(f1, encoded=False) >= 60: recording_paths.append(f1) elif fname == 'toencode': for f1 in Tools.list_dirs(f0): fname1 = Tools.fname(f1) if "copy" not in fname1 and not fname1 == "encoded": for f2 in Tools.list_dirs(f1): if self._valid_folder(f2): if self._length(f2, encoded=False) >= 60: recording_paths.append(f2) else: for f1 in Tools.list_dirs(f0): if self._valid_folder(f1): if self._length(f1) >= 60: recording_paths.append(f1) with open(ou_path, 'w+') as f: f.write('\n'.join(recording_paths))
def __init__(self, device=None, state_dict_root=None, root=None): self.device = device self.POS = tuple([ Tools.fname(f) for f in Tools.list_dirs(Tools.list_dirs(root)[0]) ]) self.VAE_DICT, self.EMB_SIZE = self._load_vae_dicts(root) am = AlignMatrix(root) self.alignments = am.load() self.VAE = [ VAE(state_dict_path=os.path.join(state_dict_root, pos, 'vae_mdl.pth')).to(device) for pos in poem(self.POS, "LOADING VAE MODELS") ] self.cv_hist, self.labels = self._init_hist((240, 240 * len(self.POS)), root) self.lbl_dict = self._init_labels(root)
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 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 embed(self): dict_ = None folder = None with torch.no_grad(): for ii in poem(range(len(self.dataset)), "STORING EMBEDDINGS " + self.pos): X, _, paths = self.dataset[ii] X = X.to(self.device) _, mu, _ = self.vae(X) mu = mu.cpu().numpy().tolist() for pth, emb in zip(paths, mu): if not folder == os.path.join(*pth.split('/')[:-1]): if dict_ is not None: with open('/' + folder + '_vae.json', 'w+') as f: json.dump(dict_, f, indent=1) folder = os.path.join(*pth.split('/')[:-1]) dict_ = {} dict_[Tools.fname(pth)] = emb with open('/' + folder + '_vae.json', 'w+') as f: json.dump(dict_, f, indent=1)
import cv2 import os import shutil from src.utils.tools import Tools from src.segmenter_grasping import SegmenterGrasping IN_ROOT = '' OU_ROOT = '' RGB = False # whether videos are in RGB (true) format or BGR (false) format # MANUALLY SEGMENT RECORDINGS s = SegmenterGrasping() for trial in Tools.tqdm_pbar(Tools.list_dirs(IN_ROOT), description="SEGMENTING"): s.segment(trial, os.path.join(OU_ROOT, Tools.fname(trial))) # COPY ACTUAL FRAMES TO OU_ROOT for trial in Tools.tqdm_pbar(Tools.list_dirs(OU_ROOT), description="MOVING FRAMES"): for pos in Tools.list_dirs(trial): with open(os.path.join(pos, 'frames.txt'), 'r') as f: for path in f: shutil.copy( path.replace('\n', ''), os.path.join(pos, Tools.fname(path.replace('\n', '')))) os.remove(os.path.join(pos, 'frames.txt')) # BGR->RGB, RESIZE, AND CROP IMAGES for trial in Tools.tqdm_pbar(Tools.list_dirs(IN_ROOT), description="CROPPING"): for pos in Tools.list_dirs(trial):
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') EPOCHS = 100 trainer = TCNTrainer(devices=devices, state_dict_paths=STATE_DICT_PATHS, train_root=TRAIN_ROOT, val_root=VAL_ROOT, save_loc=SAVE_LOC) trainer.train_loop(0, EPOCHS)
def visualize(self, in_folder, ou_folder): Tools.makedirs(ou_folder) self.len = [ len(Tools.list_files(os.path.join(in_folder, pos))) for pos in self.POS ] frames = [ Tools.fname(f) for f in Tools.list_files( os.path.join(in_folder, min(self.POS, key=lambda x: len(x)))) ] n_zfill = len(frames[0].split('.')[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(self.POS), 480 if self.lbl_dict is None else 720)) # init frame # if os.path.exists(in_folder.replace('_view', '') + '.json'): if self.lbl_dict is not None: main_frm = np.zeros((720, len(self.POS) * 240, 3), dtype=np.uint8) else: main_frm = np.zeros((480, len(self.POS) * 240, 3), dtype=np.uint8) for fii, frame in poem(enumerate(frames), Tools.fname(in_folder), total=len(frames)): try: vae_embds = self._imagine(in_folder, fii, n_zfill) for pii, pos in enumerate(self.POS): orig_img = cv2.imread(os.path.join(in_folder, pos, frame)) orig_img = cv2.resize(orig_img, (240, 240)) main_frm[0:240, pii * 240:(pii + 1) * 240, :] = \ np.copy(orig_img[:, :, :3]) with torch.no_grad(): X = torch.FloatTensor(vae_embds[pos]).to( self.device).unsqueeze(0) y = self.VAE[pii].decode(X) imag_img = y.cpu().numpy().squeeze() imag_img = (imag_img + 1.) * 127.5 imag_img = imag_img.transpose(1, 2, 0)[:, :, ::-1].astype( np.uint8) imag_img = cv2.resize(imag_img, (240, 240)) main_frm[240:480, pii * 240:(pii + 1) * 240, :] = \ np.copy(imag_img) if self.lbl_dict is not None: est_lbls = self._estimate_labels(in_folder, fii, n_zfill) vals = tuple([(float(self.lbl_dict[in_folder][frame][lbl]), el) for lbl, el in zip(self.labels, est_lbls)]) main_frm[480:720, :, :] = self.cv_hist.plot(values=vals) except IndexError: pass writer.write(main_frm) writer.release()
def segment_video(self, in_folder, ou_folder): """ Extract trials from video Args: in_folder: string - path to intput video folder ou_fodler: string - path to output folder root """ try: l_color = self._load_videos(in_folder)[0] img = np.copy(l_color[0]) img = cv2.resize(img, (240, 160)) Tools.render(img[:, :, ::-1]) except Exception as e: print(in_folder) print(e) time.sleep(1000) raise e progression = [0] created_folders = [] ii = 0 key = '' try: with tqdm(total=len(l_color)) as pbar: while key != ord('+'): feedback = '' key = self.stdscr.getch() self.stdscr.addch(20, 25, key) self.stdscr.refresh() if key == ord('d'): # move forwards slow ii = min(ii + 1, len(l_color) - 1) elif key == ord('e'): # move forwards medium ii = min(ii + 5, len(l_color) - 1) elif key == ord('3'): # move forwards fast ii = min(ii + 25, len(l_color) - 1) elif key == ord('a'): # move backwards slow ii = max(progression[-1], ii - 1) elif key == ord('q'): # move backwards medium ii = max(progression[-1], ii - 5) elif key == ord('1'): # move backwards fast ii = max(progression[-1], ii - 25) elif key == ord(']'): # STORE SEGMENT # make new folder store_folder_number = len( [f for f in created_folders if f is not None]) savepath = os.path.join( ou_folder, Tools.fname(in_folder) + '_' + str(store_folder_number).zfill(3)) self._makedirs(savepath) # store images self._store_imgs(l_color, savepath + '/left', progression[-1], ii) # update progression data progression.append(ii) created_folders.append(savepath) feedback = 'store' elif key == ord('u'): # skip progression.append(ii) created_folders.append(None) feedback = 'skip' elif key == ord('.'): # undo if len(progression) > 1: if created_folders[-1] is not None: shutil.rmtree(created_folders[-1]) progression.pop(-1) created_folders.pop(-1) ii = progression[-1] feedback = 'back' pbar.update(ii - pbar.n) img = np.copy(l_color[ii]) img = cv2.resize(img, (240, 160)) if feedback == 'store': cv2.circle(img, (10, 10), 10, (0, 255, 0), -1) if feedback == 'skip': cv2.circle(img, (10, 10), 10, (255, 0, 0), -1) if feedback == 'back': cv2.circle(img, (10, 10), 10, (0, 0, 255), -1) Tools.render(img[:, :, ::-1]) except Exception as e: print(e) time.sleep(1000) raise e with open('./res/finished_files.txt', 'a+') as f: f.write(in_folder + '\n')