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