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