Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import shutil | |
| class VideoSequence: | |
| def __init__(self, | |
| base_dir, | |
| key_ind, | |
| input_subdir='videos', | |
| key_subdir='keys0', | |
| tmp_subdir='tmp', | |
| input_format='frame%04d.jpg', | |
| key_format='%04d.jpg', | |
| out_subdir_format='out_%d', | |
| blending_out_subdir='blend', | |
| output_format='%04d.jpg'): | |
| #if (end_frame - beg_frame) % interval != 0: | |
| # end_frame -= (end_frame - beg_frame) % interval | |
| self.__base_dir = base_dir | |
| self.__input_dir = os.path.join(base_dir, input_subdir) | |
| self.__key_dir = os.path.join(base_dir, key_subdir) | |
| self.__tmp_dir = os.path.join(base_dir, tmp_subdir) | |
| self.__input_format = input_format | |
| self.__blending_out_dir = os.path.join(base_dir, blending_out_subdir) | |
| self.__key_format = key_format | |
| self.__out_subdir_format = out_subdir_format | |
| self.__output_format = output_format | |
| self.__key_ind = key_ind | |
| #self.__beg_frame = beg_frame | |
| #self.__end_frame = end_frame | |
| #self.__interval = interval | |
| self.__n_seq = len(key_ind)-1#(end_frame - beg_frame) // interval | |
| self.__make_out_dirs() | |
| os.makedirs(self.__tmp_dir, exist_ok=True) | |
| def beg_frame(self): | |
| return self.__key_ind[0]#self.__beg_frame | |
| def end_frame(self): | |
| return self.__key_ind[-1]#self.__end_frame | |
| def n_seq(self): | |
| return self.__n_seq | |
| def blending_dir(self): | |
| return os.path.abspath(self.__blending_out_dir) | |
| def interval(self, i): | |
| return self.get_sequence_beg_id(i + 1) - self.get_sequence_beg_id(i) | |
| def remove_out_and_tmp(self): | |
| for i in range(self.n_seq + 1): | |
| out_dir = self.__get_out_subdir(i) | |
| shutil.rmtree(out_dir) | |
| shutil.rmtree(self.__tmp_dir) | |
| def get_input_sequence(self, i, is_forward=True): | |
| beg_id = self.get_sequence_beg_id(i) | |
| end_id = self.get_sequence_beg_id(i + 1) | |
| if is_forward: | |
| id_list = list(range(beg_id, end_id)) | |
| else: | |
| id_list = list(range(end_id, beg_id, -1)) | |
| path_dir = [ | |
| os.path.join(self.__input_dir, self.__input_format % id) | |
| for id in id_list | |
| ] | |
| return path_dir | |
| def get_output_sequence(self, i, is_forward=True): | |
| beg_id = self.get_sequence_beg_id(i) | |
| end_id = self.get_sequence_beg_id(i + 1) | |
| if is_forward: | |
| id_list = list(range(beg_id, end_id)) | |
| else: | |
| i += 1 | |
| id_list = list(range(end_id, beg_id, -1)) | |
| out_subdir = self.__get_out_subdir(i) | |
| path_dir = [ | |
| os.path.join(out_subdir, self.__output_format % id) | |
| for id in id_list | |
| ] | |
| return path_dir | |
| def get_temporal_sequence(self, i, is_forward=True): | |
| beg_id = self.get_sequence_beg_id(i) | |
| end_id = self.get_sequence_beg_id(i + 1) | |
| if is_forward: | |
| id_list = list(range(beg_id, end_id)) | |
| else: | |
| i += 1 | |
| id_list = list(range(end_id, beg_id, -1)) | |
| tmp_dir = self.__get_tmp_out_subdir(i) | |
| path_dir = [ | |
| os.path.join(tmp_dir, 'temporal_' + self.__output_format % id) | |
| for id in id_list | |
| ] | |
| return path_dir | |
| def get_edge_sequence(self, i, is_forward=True): | |
| beg_id = self.get_sequence_beg_id(i) | |
| end_id = self.get_sequence_beg_id(i + 1) | |
| if is_forward: | |
| id_list = list(range(beg_id, end_id)) | |
| else: | |
| i += 1 | |
| id_list = list(range(end_id, beg_id, -1)) | |
| tmp_dir = self.__get_tmp_out_subdir(i) | |
| path_dir = [ | |
| os.path.join(tmp_dir, 'edge_' + self.__output_format % id) | |
| for id in id_list | |
| ] | |
| return path_dir | |
| def get_pos_sequence(self, i, is_forward=True): | |
| beg_id = self.get_sequence_beg_id(i) | |
| end_id = self.get_sequence_beg_id(i + 1) | |
| if is_forward: | |
| id_list = list(range(beg_id, end_id)) | |
| else: | |
| i += 1 | |
| id_list = list(range(end_id, beg_id, -1)) | |
| tmp_dir = self.__get_tmp_out_subdir(i) | |
| path_dir = [ | |
| os.path.join(tmp_dir, 'pos_' + self.__output_format % id) | |
| for id in id_list | |
| ] | |
| return path_dir | |
| def get_flow_sequence(self, i, is_forward=True): | |
| beg_id = self.get_sequence_beg_id(i) | |
| end_id = self.get_sequence_beg_id(i + 1) | |
| if is_forward: | |
| id_list = list(range(beg_id, end_id - 1)) | |
| path_dir = [ | |
| os.path.join(self.__tmp_dir, 'flow_f_%04d.npy' % id) | |
| for id in id_list | |
| ] | |
| else: | |
| id_list = list(range(end_id, beg_id + 1, -1)) | |
| path_dir = [ | |
| os.path.join(self.__tmp_dir, 'flow_b_%04d.npy' % id) | |
| for id in id_list | |
| ] | |
| return path_dir | |
| def get_input_img(self, i): | |
| return os.path.join(self.__input_dir, self.__input_format % i) | |
| def get_key_img(self, i): | |
| sequence_beg_id = self.get_sequence_beg_id(i) | |
| return os.path.join(self.__key_dir, | |
| self.__key_format % sequence_beg_id) | |
| def get_blending_img(self, i): | |
| return os.path.join(self.__blending_out_dir, self.__output_format % i) | |
| def get_sequence_beg_id(self, i): | |
| return self.__key_ind[i]#i * self.__interval + self.__beg_frame | |
| def __get_out_subdir(self, i): | |
| dir_id = self.get_sequence_beg_id(i) | |
| out_subdir = os.path.join(self.__base_dir, | |
| self.__out_subdir_format % dir_id) | |
| return out_subdir | |
| def __get_tmp_out_subdir(self, i): | |
| dir_id = self.get_sequence_beg_id(i) | |
| tmp_out_subdir = os.path.join(self.__tmp_dir, | |
| self.__out_subdir_format % dir_id) | |
| return tmp_out_subdir | |
| def __make_out_dirs(self): | |
| os.makedirs(self.__base_dir, exist_ok=True) | |
| os.makedirs(self.__blending_out_dir, exist_ok=True) | |
| for i in range(self.__n_seq + 1): | |
| out_subdir = self.__get_out_subdir(i) | |
| tmp_subdir = self.__get_tmp_out_subdir(i) | |
| os.makedirs(out_subdir, exist_ok=True) | |
| os.makedirs(tmp_subdir, exist_ok=True) | |