def _debug(self, message):
"""
shutdown curses, print message, then exit
"""
self._shutdown()
Tools.pyout(message)
sys.exit(0)
def unpack_video(source_path, oudir, rm_outline=False):
Tools.makedirs(oudir)
# crop = Config.crop_vals(source_path.split('/')[-1])
i = 0
cap = cv2.VideoCapture(source_path)
print(oudir, end='')
while True:
try:
_, frame = cap.read()
# frame = frame[crop['n']:crop['s'], crop['w']:crop['e']]
if frame is None:
break
if rm_outline:
frame = Prep.rm_outline(frame)
# frame = cv2.resize(frame, (368, 368),
# interpolation=cv2.INTER_AREA)
cv2.imwrite(os.path.join(
oudir, str(i).zfill(8) + '.jpg'), frame)
if not Config.SERVER:
cv2.imshow('img', frame)
k = cv2.waitKey(1)
if k == ord('q'):
sys.exit(0)
i += 1
if not i % 100:
print(str(i).zfill(10), end='\r')
except TypeError:
cap.release()
print()
break
def val_with_labels(self, epoch):
"""inference model on validation set, for each frame in each video
find nearest neighbor in each other video and check accuracy on
labels of that frame
"""
self.tcn.switch_mode('eval')
Y = [] # tcn embeddings
V = [] # video names
R = [] # video classes
L = [] # frame features
# compute embeddings for all frames
with torch.no_grad():
for batch_idx in range(len(self.val_set)):
batch = self.val_set[batch_idx]
X, v, r, lbls = batch[0], batch[1], batch[2], batch[3]
X = X.to(self.devices[0])
y = self.tcn(X).cpu().numpy()
Y.extend(list(y))
V.extend(v)
R.extend(r)
L.extend(lbls)
print("Validation Epoch {}: {}/{} ({:.0f}%)".format(
epoch, batch_idx, len(self.val_set),
100. * batch_idx / len(self.val_set)),
' ',
end='\r')
accuracy, correct, N = Loss.labeled_accuracy(Y, V, R, L)
Tools.log('Validation Epoch {}: Accuracy {:.0f}/{} ({:.0f}%)'.format(
epoch, correct, N, 100. * accuracy))
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 __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 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 alphapose(indir=None):
ap = AlphaPose()
for trial_path in Tools.list_dirs(indir):
for pos_path in Tools.list_dirs(trial_path):
ap.process_dir(
pos_path,
os.path.join(pos_path, '3d'))
def _sample_frame(self, pos_folder, used_frames, anchor_idx=None):
valid = False
# frame_pth = None
if anchor_idx is None:
all_frames = Tools.list_files(pos_folder, end='.jpg')
else:
# check that current frame is a TP to anchor frame
all_frames = [
file for file in Tools.list_files(pos_folder, end='.jpg') if
abs(int(file.split('/')[-1].split('.')[0]) - anchor_idx) <
self.pos_range]
random.shuffle(all_frames)
for frame_pth in all_frames:
# check that current frame is not a FP to any used frame
idx = int(frame_pth.split('/')[-1].split('.')[0])
frame_pth = random.choice(all_frames)
if all([abs(idx - u_idx) > self.m * self.pos_range for
u_idx in used_frames]):
valid = True
break
if valid:
frame_idx = int(frame_pth.split('/')[-1].split('.')[0])
return True, frame_pth, frame_idx
else:
return False, None, None
def segment(self, in_root, ou_root):
"""
Extract trials from video
Args:
in_root: string - path to input trial folder
ou_root: string - path to output folder root
"""
# list containing history of anchor points
progression = [0]
folders = []
try:
frames, N = self._load_images(in_root)
n = 0
key = ''
feedback = ''
with tqdm(total=N) as pbar:
while key != ord('+'):
self.main_frm[:, :, :] = frames[n, :, :, ::-1]
if feedback == 'fwd':
# if saved segment in forward dataset:
# display green circle for feedback
self._feedback((0, 255, 0))
if feedback == 'bwd':
# if saved segment in backward dataset:
# display magenta circle for feedback
self._feedback((255, 0, 220))
if feedback == 'skip':
# if new anchor point set without save:
# display blue circle for feedback
self._feedback((255, 0, 0))
if feedback == 'undo':
# if undo:
# display red circle for feedback
self._feedback((0, 0, 255))
# render and await user input
Tools.render(self.main_frm)
feedback = ''
key = self.stdscr.getch()
self.stdscr.addch(20, 25, key)
self.stdscr.refresh()
# process user input
n = self._process_nav(key, n, N, progression)
save = self._process_save(key, in_root, ou_root, folders,
progression[-1], n, progression)
if save:
feedback = save
if self._process_skip(key, folders, n, progression):
feedback = 'skip'
n, undo = self._process_undo(key, progression, folders, n)
if undo:
feedback = 'undo'
pbar.update(n - pbar.n)
except Exception as e:
self._shutdown('Shutting down due to exception')
raise e
def _init_weights(self):
try:
nn.init.xavier_normal_(self.conv1.weight)
nn.init.normal_(self.conv1.bias)
Tools.log("Init TCN on device 0: Success")
except Exception as e:
Tools.log("Init TCN on device 0: Fail... (abort) " + e)
sys.exit(1)
def load(self):
alignment = {}
for trial_root in Tools.tqdm_pbar(Tools.list_dirs(self.root),
"LOADING ALIGNMENT DATA"):
with open(os.path.join(trial_root, 'alignment.json'), 'r') as f:
alignment[trial_root] = json.load(f)
return alignment
def _plot(self, path):
appendage = '/steady' if 'pouring' in self.trial_root else '/left'
frames = Tools.list_files(self.trial_root + appendage)
a_fldr = self.anchor_root + appendage
for ii in range(len(path)):
frame = cv2.imread(frames[ii])
ancho = cv2.imread(os.path.join(a_fldr, path[ii][0]))
# Tools.debug(frame.shape, ancho.shape)
Tools.render(np.concatenate((frame, ancho), axis=1))
def crop_start_end(indir):
D_l = Prep.init_joint_dict(os.path.join(
indir, '0', '3d', 'alphapose-results.json'))
D_r = Prep.init_joint_dict(os.path.join(
indir, '2', '3d', 'alphapose-results.json'))
img_paths_l = sorted(Tools.list_files(os.path.join(indir, '0')))
img_paths_m = sorted(Tools.list_files(os.path.join(indir, '1')))
img_paths_r = sorted(Tools.list_files(os.path.join(indir, '2')))
img_paths_l, img_paths_r = Prep.make_equal_length(
(img_paths_l, img_paths_r))
t_ = None # last positive frame
splits = []
for t, (img_pth_l, img_pth_r) in enumerate(
zip(img_paths_l, img_paths_r)):
img_name = img_pth_l.split('/')[-1]
img_l = cv2.imread(img_pth_l)
img_r = cv2.imread(img_pth_r)
sitting = False
try:
for person in D_l[img_name]:
if Prep.on_couch(person['KP']['LHip'], 0) or \
Prep.on_couch(person['KP']['RHip'], 0):
sitting = True
break
except KeyError:
pass
if not sitting:
try:
for person in D_r[img_name]:
if Prep.on_couch(person['KP']['LHip'], 2) or \
Prep.on_couch(person['KP']['RHip'], 2):
sitting = True
break
except KeyError:
pass
if sitting:
if t_ is None or t > t_ + 5:
splits.append([])
splits[-1].append(img_name)
t_ = t
cv2.circle(img_l, (15, 15), 10, (0, 255, 0), thickness=-1)
else:
cv2.circle(img_l, (15, 15), 10, (0, 0, 255), thickness=-1)
cv2.imshow('foo', np.concatenate((img_l, img_r), axis=1))
if cv2.waitKey(1) & 0xFF == ord('q'):
sys.exit(0)
keep = max(splits, key=lambda x: len(x))
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 _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 _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 _makedirs(self, root):
"""
create folder structure
"""
self._write(root)
try:
Tools.makedirs(root, delete=True)
Tools.makedirs(os.path.join(root, 'left'))
except Exception as e:
print(e)
time.sleep(1000)
raise e
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 _init_weights(self):
try:
# conv1
nn.init.xavier_normal_(self.conv1.weight.data)
nn.init.normal_(self.conv1.bias.data)
# conv2
nn.init.xavier_normal_(self.conv2.weight.data)
nn.init.normal_(self.conv2.bias.data)
Tools.log("Init Encoder: Success")
except Exception as e:
Tools.log("Init Encoder: Fail... (abort) " + e)
sys.exit(1)
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 embedding_accuracy(embeddings, labels, perspectives, device=None):
"""Computes the ratio of positive embeddings in the batch which are
closer together than any negative pair.
Args:
embeddings: (N,D) FloatTensor - TCN embeddings of inputs
labels: (N,) FloatTensor - labels indicating positive pairs
perspectives: (N,) FloatTensor - labels indicating video
perspective
device: torch.cuda.device
Returns:
accuracy Tensor(float)
"""
# Tools.pyout(embeddings)
# compute pairwise distance matrix
pdist_matrix = Tools.pdist(embeddings)
# Tools.pyout(pdist_matrix)
# mask labels equal
leq_mask = Tools.pequal(labels).float()
# mask perspectives equal
peq_mask = Tools.pequal(perspectives).float()
# mask positive pairs: same label, different perspective
pos_mask = leq_mask * (1 - peq_mask)
# mask negative pairs: different label, same perspective
neg_mask = (1 - leq_mask) * peq_mask
# compute max distance between positive pairs for each entry
pos_dists, _ = torch.max(
pdist_matrix * pos_mask, dim=1, keepdim=True)
# compute min distance between negative pairs for each entry
neg_dists = torch.where(
neg_mask > 0,
pdist_matrix,
torch.full(neg_mask.size(), float("inf")).to(device))
neg_dists, _ = torch.min(neg_dists, dim=1, keepdim=True)
# compute ratio where positive distance is smaller than all
# negative distance element wise
dist_diff = (pos_dists < neg_dists).float()
# reduce mean
accuracy = sum(dist_diff) / dist_diff.size()[0]
return accuracy.squeeze()
def _load(self, path):
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 Encoder: Success")
return True
except Exception as e:
Tools.log("Load Encoder: Fail " + e)
return False
return False
def wrist_dict(indir):
datadir = indir.replace('kinect-recordings-3', 'grabbing')
ou_dict = {}
for trial_folder in Tools.list_dirs(indir):
trial_name = trial_folder.split('/')[-1]
D_l = Prep.init_joint_dict(os.path.join(
datadir, trial_name, '0', '3d', 'alphapose-results.json'))
D_m = Prep.init_joint_dict(os.path.join(
datadir, trial_name, '1', '3d', 'alphapose-results.json'))
D_r = Prep.init_joint_dict(os.path.join(
datadir, trial_name, '2', '3d', 'alphapose-results.json'))
ou_dict[os.path.join(trial_folder, 'color-recording-left.avi')] = \
Prep.req_frames(D_l)
ou_dict[os.path.join(
trial_folder, 'color-recording-middle.avi')] = \
Prep.req_frames(D_m)
ou_dict[os.path.join(
trial_folder, 'color-recording-right.avi')] = \
Prep.req_frames(D_r)
with open('/media/roblaundry/kinect-recordings-3/wrist_positions.json',
'w+') as f:
json.dump(ou_dict, f, indent=2)
def _load_embedding_dicts(self, folder):
dicts = []
for json_filename in sorted(Tools.list_files(
folder, substr='embed_', end='.json')):
with open(json_filename, 'r') as f:
dicts.append(json.load(f))
return dicts
def alpha_video(self, folder):
D = Prep.init_joint_dict(
os.path.join(folder, '3d', 'alphapose-results.json'))
for img_pth in sorted(Tools.list_files(folder)):
img_name = img_pth.split('/')[-1]
img = cv2.imread(img_pth)
wait = 1
try:
person = D[img_name][0]
lwrist = (int(person['KP']['LWrist'][0]),
int(person['KP']['LWrist'][1]))
rwrist = (int(person['KP']['RWrist'][0]),
int(person['KP']['RWrist'][1]))
col = (0, 0, 255)
if lwrist[0] < 0 or lwrist[1] < 0 or rwrist[0] < 0 or lwrist[
1] < 0:
print('foo')
wait = 5000
col = (0, 255, 0)
cv2.circle(img, lwrist, 10, col, thickness=-1)
# cv2.circle(img, rwrist, 10, (0, 0, 255), thickness=-1)
# print(person['KP']['RWrist'], person['KP']['LWrist'])
except KeyError:
pass
cv2.imshow('vid', img)
if cv2.waitKey(wait) & 0xFF == ord('q'):
sys.exit(0)
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 __init__(self,
root_dir='./res/datasets/folding/train',
batch_size=Config.TCN_BATCH,
input_size=(3,) + Config.TCN_IMG_SIZE,
pos_range=Config.TCN_POS_RANGE,
negative_multiplier=Config.TCN_NEGATIVE_MULTIPLIER,
transform=None, # unused
augment=False):
self.root = root_dir
self.batch_size = batch_size
self.input_size = input_size
self.pos_range = pos_range
self.m = negative_multiplier
self.augment = augment
self.trial_names = []
self.seeds = [] # use same seed when sampling eval
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.255])])
for f in Tools.list_dirs(root_dir):
if augment:
self.trial_names.append(f)
else:
self.trial_names.append(f)
self.seeds.append(random.randint(0, 9001))
def make_align_dict(self, root, trial_root, lock):
dict_ = {}
for anchor_root in Tools.list_dirs(root):
if not trial_root == anchor_root:
dict_[anchor_root] = self._align(trial_root, anchor_root, lock)
with open(os.path.join(trial_root, 'alignment.json'), 'w+') as f:
json.dump(dict_, f, indent=1)
def _construct_matrix(self, root):
def align(trial_root, anchor_root):
A = Aligner()
path = A.align(trial_root, anchor_root)
return (trial_root, anchor_root, path)
pool = multiprocessing.Pool(Config.N_WORKERS)
lock = multiprocessing.Manager().Lock()
tasks = []
for trial_root in Tools.list_dirs(root):
tasks.append((root, trial_root, lock))
# try:
for _ in Tools.tqdm_pbar(pool.imap_unordered(self._align, tasks),
"ALIGNING (multiprocessing pool)",
total=len(tasks)):
pass
