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 _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 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 _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 _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 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,
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 unpack_videos(indir, oudir, devices):
found_positive = False
for iter_path in sorted(Tools.list_dirs(indir)):
if True or all(Prep.video_length(v) >= 60
for v in Tools.list_files(iter_path)) and \
len(Tools.list_files(iter_path)) > 0:
found_positive = True
print('Unpacking', iter_path)
iter_targ_path = os.path.join(
oudir, iter_path.split('/')[-1])
Tools.makedirs(iter_targ_path, delete=True)
for vid_path in Tools.list_files(iter_path):
if 'color' in vid_path:
if 'left' in vid_path:
Prep.unpack_video(
vid_path, os.path.join(iter_targ_path, '0'))
elif 'middle' in vid_path:
Prep.unpack_video(
vid_path, os.path.join(iter_targ_path, '1'))
elif 'right' in vid_path:
Prep.unpack_video(
vid_path, os.path.join(iter_targ_path, '2'))
else:
print('Error unpacking', vid_path)
elif 'depth' in vid_path:
if 'left' in vid_path:
Prep.unpack_video(
vid_path,
os.path.join(iter_targ_path, '0/3d'),
rm_outline=True)
elif 'middle' in vid_path:
Prep.unpack_video(
vid_path,
os.path.join(iter_targ_path, '1/3d'),
rm_outline=True)
elif 'right' in vid_path:
Prep.unpack_video(
vid_path,
os.path.join(iter_targ_path, '2/3d'),
rm_outline=True)
else:
print('Error unpacking', vid_path)
else:
print('Ignore', iter_path)
return found_positive
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 __init__(self,
root_dir=None,
batch_size=Config.TCN_BATCH,
input_size=(3,) + Config.TCN_IMG_SIZE):
self.root = root_dir
self.input_size = input_size
self.batch_size = batch_size
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.255])])
self.frame_names = Tools.list_files(self.root)
self.batch_paths = [[]]
for frame_name in sorted(Tools.list_files(self.root, end='.jpg')):
if len(self.batch_paths[-1]) >= self.batch_size:
# start new batch
self.batch_paths.append([frame_name])
else:
# append to last batch
self.batch_paths[-1].append(frame_name)
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_labels(self, root):
lbl_dict = {}
for label_file in Tools.list_files(root, end='.json'):
with open(label_file, 'r') as f:
if os.path.exists(label_file.replace('.json', '_view')):
lbl_dict[label_file.replace('.json', '_view')] = \
json.load(f)
else:
lbl_dict[label_file.replace('.json', '')] = json.load(f)
if len(lbl_dict) == 0:
return None
else:
return lbl_dict
def _heatmap(self, matrix, path=None):
"""
draw a heatmap of distance matrix or accumulated cost matrix with
selected path
"""
img = np.copy(matrix)
img -= np.min(img)
img /= np.max(img)
if path is None:
img = cv2.resize(img, (360, 360))
return img[::-1, :]
else:
img = (np.stack((img, ) * 3, axis=-1) * 255).astype(np.uint8)
appendage = '/steady' if 'pouring' in self.trial_root else '/left'
frames = Tools.list_files(self.trial_root + appendage)
anchors = Tools.list_files(self.anchor_root + appendage)
# print(len(frames), len(path))
for ii in range(len(path) - 2, -1, -1):
p0 = path[ii]
p1 = path[ii + 1]
img = cv2.line(img, p0, p1, (0, 0, 255), 2)
frame = cv2.imread(frames[p0[0]])
ancho = cv2.imread(anchors[p0[1]])
frame = cv2.resize(frame, (240, 240))
ancho = cv2.resize(ancho, (240, 240))
if not ii % 3:
img_ = np.copy(img)
img_ = cv2.resize(img_, (240, 240))
Tools.render(np.concatenate(
(img_[::-1, :, :], frame[:, :, ::1], ancho), axis=1),
waitkey=50)
return img[::-1, :, :]
def _load_dicts(self, indir):
# load dictionaries of camera positions
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'))
# load dictionaries of real positions relative to camera
lvid_path = os.path.join(
indir.replace('/grabbing/', '/kinect-recordings-3/'),
'color-recording-left.avi')
rvid_path = os.path.join(
indir.replace('/grabbing/', '/kinect-recordings-3/'),
'color-recording-right.avi')
xyz = None
# load image paths of video and make equal length
img_paths_l = sorted(Tools.list_files(os.path.join(indir, '0')))
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))
return D_l, D_r, lvid_path, rvid_path, xyz, img_paths_l, img_paths_r
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 _init_embeddings(self, root, lock):
embd_files = Tools.list_files(root, end='.json', substr='embed_')
# lock.acquire()
with open(embd_files[0], 'r') as f:
sample_dict = json.load(f)
frames = list(sample_dict)
nr_frames = len(sample_dict)
emb_N = len(sample_dict[list(sample_dict)[0]])
embeddings = np.zeros((nr_frames, emb_N * len(embd_files)),
dtype=float)
# lock.release()
for ii, emb_F in enumerate(embd_files):
# lock.acquire()
with open(emb_F, 'r') as f:
dict_ = json.load(f)
# lock.release()
for jj, frame in enumerate(sorted(list(dict_))):
embeddings[jj, ii * emb_N:(ii + 1) * emb_N] = \
np.array(dict_[frame])
return frames, embeddings
def _init_hist(self, shape, root):
labels = []
minvals = []
maxvals = []
for label_file in Tools.list_files(root, end='.json'):
# Tools.pyout(label_file)
with open(label_file, 'r') as f:
d = json.load(f)
for key in d:
# Tools.pyout(key, d[key]['contact'])
for lbl in d[key]:
if lbl not in labels:
labels.append(lbl)
minvals.append(float(d[key][lbl]))
maxvals.append(float(d[key][lbl]))
else:
minvals[labels.index(lbl)] = min(
minvals[labels.index(lbl)], float(d[key][lbl]))
maxvals[labels.index(lbl)] = max(
maxvals[labels.index(lbl)], float(d[key][lbl]))
if len(labels) > 0:
return CVHist(shape, labels, minvals, maxvals, n_bars=2), labels
else:
return None, None
def detect_wrists(indir):
ct = CameraTransform()
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'))
while True:
img_paths_l = sorted(Tools.list_files(os.path.join(indir, '0')))
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))
trajectory = {}
for img_pth_l, img_pth_r in \
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)
img_ld = cv2.imread(img_pth_l.replace(
img_name, '3d/' + img_name))
img_rd = cv2.imread(img_pth_r.replace(
img_name, '3d/' + img_name))
posture = np.zeros((2, 12))
weights = np.zeros((2, 12))
on_couch = False
for img, img_d, D, pos2real, pos in [
(img_l, img_ld, D_l, ct.left2real, 0),
(img_r, img_rd, D_r, ct.right2real, 2)]:
try:
contenders = []
for person in D[img_name]:
if Prep.on_couch(person['KP']['LHip'], pos) or \
Prep.on_couch(person['KP']['RHip'], pos):
contenders.append(person)
if len(contenders) > 0:
on_couch = True
person = max(contenders, key=lambda x: x['score'])
xprw = int(person['KP']['RWrist'][0])
yprw = int(person['KP']['RWrist'][1])
crw = person['KP']['RWrist'][2]
zprw = ct.gray_val(img_d, xprw, yprw)
xpre = int(person['KP']['RElbow'][0])
ypre = int(person['KP']['RElbow'][1])
cre = person['KP']['RElbow'][2]
zpre = ct.gray_val(img_d, xpre, ypre)
xplw = int(person['KP']['LWrist'][0])
yplw = int(person['KP']['LWrist'][1])
clw = person['KP']['LWrist'][2]
zplw = ct.gray_val(img_d, xplw, yplw)
xple = int(person['KP']['LElbow'][0])
yple = int(person['KP']['LElbow'][1])
cle = person['KP']['LElbow'][2]
zple = ct.gray_val(img_d, xple, yple)
xrrw, yrrw, zrrw = pos2real(xprw, yprw, zprw)
xrre, yrre, zrre = pos2real(xpre, ypre, zpre)
xrlw, yrlw, zrlw = pos2real(xplw, yplw, zplw)
xrle, yrle, zrle = pos2real(xple, yple, zple)
arx = m.atan(float(xrrw - xrre) /
float(zrrw - zrre + 1e-42))
ary = m.atan(float(yrrw - yrre) /
float(zrrw - zrre + 1e-42))
arz = m.atan(float(xrrw - xrre) /
float(yrrw - yrre + 1e-42)) / 2
alx = m.atan(float(xrlw - xrle) /
float(zrlw - zrle + 1e-42))
aly = m.atan(float(yrlw - yrle) /
float(zrlw - zrle + 1e-42))
alz = m.atan(float(xrlw - xrle) /
float(yrlw - yrle + 1e-42)) / 2
posture[0 if pos == 0 else 1, :] = np.array([xrlw, yrlw, zrlw,
alx, aly, alz,
xrrw, yrrw, zrrw,
arx, ary, arz])
weights[0 if pos == 0 else 1, :] = np.array([clw, clw, clw,
clw + cle, clw + cle, clw + cle,
crw, crw, crw,
crw + cre, crw + cre, crw + cre])
for IMG in [img, img_d]:
for xpw, ypw, zpw, xrw, yrw, zrw, \
xpe, ype, zpe, xre, yre, zre in (
(xprw, yprw, zprw, xrrw, yrrw, zrrw,
xpre, ypre, zpre, xrre, yrre, zrre),
(xplw, yplw, zplw, xrlw, yrlw, zrlw,
xple, yple, zple, xrle, yrle, zrle)):
measure = zrw # ct.real2left(xr, yr, zr)
color = Prep.scr2rgb(-measure / 140. + 0.5)
cv2.line(IMG, (xpw, ypw),
(xpe, ype), color)
cv2.circle(IMG, (xpw, ypw), 3,
color, thickness=-1)
cv2.circle(IMG, (xpe, ype), 3,
color, thickness=-1)
except KeyError:
pass
if on_couch:
posture = np.average(posture, axis=0, weights=weights)
trajectory[img_name] = posture.copy().tolist()
disp = np.concatenate((img_l, img_r), axis=1)
if not Config.SERVER:
cv2.imshow('img', disp)
if cv2.waitKey(30) & 0xFF == ord('q'):
sys.exit(0)
# with open(os.path.join(indir, 'trajectory.json'), 'w+') as f:
# json.dump(trajectory, f, indent=2)
break
def extract_trajectory(indir):
trajectory = {}
# load dictionaries of camera projections
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'))
# load dictionaries of real camera positions
lvid_path = os.path.join(indir.replace(
'/grabbing/', '/kinect-recordings-3/'), 'color-recording-left.avi')
rvid_path = os.path.join(indir.replace(
'/grabbing/', '/kinect-recordings-3/'), 'color-recording-right.avi')
xyz = None
# load image paths of video
img_paths_l = sorted(Tools.list_files(os.path.join(indir, '0')))
img_paths_r = sorted(Tools.list_files(os.path.join(indir, '2')))
# if videos are not equal lenght, cut from start and end to make them
# equal.
img_paths_l, img_paths_r = Prep.make_equal_length(
(img_paths_l, img_paths_r))
# loop over all frames
for img_pth_l, img_pth_r in zip(img_paths_l, img_paths_r):
# left and right frame have same name (but in different folder)
img_name = img_pth_l.split('/')[-1]
frame = str(int(img_name.split('.')[0]))
# for testing: display images
img_l = cv2.imread(img_pth_l)
img_r = cv2.imread(img_pth_r)
# get person on couch in left view (if any)
try:
peoples = [person for person in D_l[img_name] if
Prep.on_couch(person['KP']['LHip'], 0) or
Prep.on_couch(person['KP']['RHip'], 0)]
if len(peoples) == 0:
lperson = None
else:
lperson_idx = max(range(len(peoples)),
key=lambda x: peoples[x]['score'])
lperson = peoples[lperson_idx]
except KeyError:
pass
# get person on couch in right view (if any)
try:
peoples = [person for person in D_r[img_name] if
Prep.on_couch(person['KP']['LHip'], 2) or
Prep.on_couch(person['KP']['RHip'], 2)]
if len(peoples) == 0:
rperson = None
else:
rperson_idx = max(range(len(peoples)),
key=lambda x: peoples[x]['score'])
rperson = peoples[rperson_idx]
except KeyError:
pass
# do not add frame to trajectory if both are None
if lperson is not None or rperson is not None:
# left wrist
if lperson is None: # so rperson is not None
# get left wrist from right frame
lwrist = Prep.affine_transform(
xyz[rvid_path][frame]['xyz'][rperson_idx * 2 + 1], 'R')
cv2.circle(
img_r,
(int(rperson['KP']['LWrist'][0]),
int(rperson['KP']['LWrist'][1])),
10, (0, 255, 0), -1)
# get right wrist from right frame
rwrist = Prep.affine_transform(
xyz[rvid_path][frame]['xyz'][rperson_idx * 2], 'R')
cv2.circle(
img_r,
(int(rperson['KP']['RWrist'][0]),
int(rperson['KP']['RWrist'][1])),
10, (0, 255, 0), -1)
elif rperson is None: # so lperson is not None
# get left wrist from left frame
lwrist = Prep.affine_transform(
xyz[lvid_path][frame]['xyz'][lperson_idx * 2 + 1], 'L')
cv2.circle(
img_l,
(int(lperson['KP']['LWrist'][0]),
int(lperson['KP']['LWrist'][1])),
10, (0, 255, 0), -1)
# get right wrist from left frame
rwrist = Prep.affine_transform(
xyz[lvid_path][frame]['xyz'][lperson_idx * 2], 'L')
cv2.circle(
img_l,
(int(lperson['KP']['RWrist'][0]),
int(lperson['KP']['RWrist'][1])),
10, (0, 255, 0), -1)
else:
if rperson['KP']['LWrist'][2] > lperson['KP']['LWrist'][2]:
# get left wrist from right frame
lwrist = Prep.affine_transform(
xyz[rvid_path][frame]['xyz'][rperson_idx * 2 + 1],
'R')
cv2.circle(
img_l,
(int(lperson['KP']['LWrist'][0]),
int(lperson['KP']['LWrist'][1])),
10, (0, 0, 255), -1)
cv2.circle(
img_r,
(int(rperson['KP']['LWrist'][0]),
int(rperson['KP']['LWrist'][1])),
10, (0, 255, 0), -1)
else:
# get left wrist from left frame
lwrist = Prep.affine_transform(
xyz[lvid_path][frame]['xyz'][lperson_idx * 2 + 1],
'L')
cv2.circle(
img_l,
(int(lperson['KP']['LWrist'][0]),
int(lperson['KP']['LWrist'][1])),
10, (0, 255, 0), -1)
cv2.circle(
img_r,
(int(rperson['KP']['LWrist'][0]),
int(rperson['KP']['LWrist'][1])),
10, (0, 0, 255), -1)
if rperson['KP']['RWrist'][2] > lperson['KP']['RWrist'][2]:
# get right wrist from right frame
rwrist = Prep.affine_transform(
xyz[rvid_path][frame]['xyz'][rperson_idx * 2], 'R')
cv2.circle(
img_l,
(int(lperson['KP']['RWrist'][0]),
int(lperson['KP']['RWrist'][1])),
10, (0, 0, 255), -1)
cv2.circle(
img_r,
(int(rperson['KP']['RWrist'][0]),
int(rperson['KP']['RWrist'][1])),
10, (0, 255, 0), -1)
else:
# get right wrist from left frame
rwrist = Prep.affine_transform(
xyz[lvid_path][frame]['xyz'][lperson_idx * 2], 'L')
cv2.circle(
img_l,
(int(lperson['KP']['RWrist'][0]),
int(lperson['KP']['RWrist'][1])),
10, (0, 255, 0), -1)
cv2.circle(
img_r,
(int(rperson['KP']['RWrist'][0]),
int(rperson['KP']['RWrist'][1])),
10, (0, 0, 255), -1)
trajectory[img_name] = [lwrist[0], lwrist[1], lwrist[2],
rwrist[0], rwrist[1], rwrist[2]]
# resize for easier rendering
img_l = cv2.resize(
img_l, (img_l.shape[1] // 2, img_l.shape[0] // 2))
img_r = cv2.resize(
img_r, (img_r.shape[1] // 2, img_r.shape[0] // 2))
# display video
Tools.render('DEBUG', np.concatenate((img_l, img_r), axis=1))
# 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):
for frame_pth in Tools.list_files(pos, end='.jpg'):
if RGB:
img = cv2.imread(frame_pth)
else:
img = cv2.imread(frame_pth)[:, :, ::-1]
img = cv2.resize(img, (552, 368))
if 'left' in pos:
D = -378
img = img[:, D:D + 368, :]
if 'middle' in pos:
D = 100
img = img[:, D:D + 368, :]
if 'right' in pos:
D = 60
img = img[:, D:D + 368, :]
cv2.imwrite(frame_pth, img)
def _paste(self):
root_dir = self.dataset.root.replace('datasets', 'plots')
for trial_pth in Tools.list_dirs(root_dir):
l_pths = sorted(Tools.list_files(os.path.join(trial_pth, '0')))
m_pths = sorted(Tools.list_files(os.path.join(trial_pth, '1')))
r_pths = sorted(Tools.list_files(os.path.join(trial_pth, '2')))
minlen = min(len(l_pths), len(m_pths), len(r_pths))
while len(l_pths) > minlen:
l_pths.pop(-1)
while len(m_pths) > minlen:
m_pths.pop(-1)
while len(r_pths) > minlen:
r_pths.pop(-1)
for L, M, R in zip(l_pths, m_pths, r_pths):
plot_l = cv2.imread(L, 1)
plot_m = cv2.imread(M, 1)
plot_r = cv2.imread(R, 1)
fram_l = cv2.imread(L.replace('plots', 'datasets'), 1)
fram_m = cv2.imread(M.replace('plots', 'datasets'), 1)
fram_r = cv2.imread(R.replace('plots', 'datasets'), 1)
bc = cv2.BORDER_CONSTANT
white = [255, 255, 255]
fram_l = cv2.copyMakeBorder(fram_l,
82,
83,
162,
163,
bc,
value=white)
fram_m = cv2.copyMakeBorder(fram_m,
82,
83,
162,
163,
bc,
value=white)
fram_r = cv2.copyMakeBorder(fram_r,
82,
83,
162,
163,
bc,
value=white)
out = cv2.copyMakeBorder(fram_l,
0,
480,
0,
1280,
bc,
value=white)
for c in range(3):
out[0:480, 640:1280, c] = fram_m[:, :, c]
out[0:480, 1280:1920, c] = fram_r[:, :, c]
out[480:960, 0:640, c] = plot_l[:, :, c]
out[480:960, 640:1280, c] = plot_m[:, :, c]
out[480:960, 1280:1920, c] = plot_r[:, :, c]
cv2.imwrite(L.replace(os.path.join(trial_pth, '0'), trial_pth),
out)
shutil.rmtree(os.path.join(trial_pth, '0'))
shutil.rmtree(os.path.join(trial_pth, '1'))
shutil.rmtree(os.path.join(trial_pth, '2'))
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()