Spaces:
Sleeping
Sleeping
| import shutil, cv2, os | |
| import numpy as np | |
| # Map {image_number: [row_num, col_num]} | |
| default_concat_mapping = {1: [1, 1], | |
| 2: [1, 2], | |
| 3: [1, 3], | |
| 4: [1, 4], | |
| 5: [1, 5], | |
| 6: [2, 3], | |
| 7: [2, 4], | |
| 8: [2, 4], | |
| 9: [3, 3], | |
| 10:[2, 5], | |
| 11:[3, 4], | |
| 12:[3, 4], | |
| 13:[4, 4], | |
| 14:[4, 4], | |
| 15:[4, 4], | |
| 16:[4, 4]} | |
| default_model_postfix_mapping = {"pix2pixHD": "_fake_B.png", | |
| "cut" : "_fake_B.png", | |
| "ugatit" : "_fake_A2B.png", | |
| "5G" : "_fake_B.png", | |
| "5.7G" : "_fake_B.png", | |
| "900Mreorg": "_fake_B.png"} | |
| class ConcatImages: | |
| def __init__(self, ): | |
| self.cv2_img_read_type = cv2.IMREAD_COLOR | |
| self.cv2_img_read_size = (512, 512) | |
| self.blank_img_size = (512, 512, 3) | |
| self.paired_image_list = [] | |
| self.save_dir = None | |
| self.save_name = None | |
| self.blank_img = np.zeros(self.blank_img_size).astype(np.uint8) | |
| # Text config | |
| self.text_pos = (10, 30) | |
| self.text_font = cv2.FONT_HERSHEY_COMPLEX | |
| self.text_color = (0, 0, 255) # Red | |
| self.text_size = 1.0 | |
| self.text_overstriking = 2 | |
| def run_concat_images(self, save_dir, paired_image_path_list, paired_image_text_list=None, row_num=0, col_num=0): | |
| self.save_dir = save_dir | |
| os.makedirs(self.save_dir, exist_ok=True) | |
| total_img_number = len(paired_image_path_list) | |
| assert total_img_number <= row_num * col_num, "Error image number {} > (row: {} x col: {})".format("%d"%total_img_number, "%d"%row_num, "%d"%col_num) | |
| flag_add_text = False | |
| if paired_image_text_list == None: | |
| flag_add_text = False | |
| else: | |
| assert len(paired_image_path_list) == len(paired_image_text_list), "Not match len(paired_image_path_list) %d and len(paired_image_text_list) %d".format(len(paired_image_path_list), len(paired_image_text_list)) | |
| flag_add_text = True | |
| self.read_images_list(paired_image_path_list, paired_image_text_list, flag_add_text) | |
| self.concat_format_grids_byorder(row_num, col_num) | |
| def read_images_list(self, paired_image_path_list, paired_image_text_list, flag_add_text): | |
| self.paired_image_list = [] | |
| if flag_add_text == True: | |
| for img_idx in range(len(paired_image_path_list)): | |
| self.paired_image_list += [self.add_text(self.read_image(paired_image_path_list[img_idx]), | |
| paired_image_text_list[img_idx])] | |
| else: | |
| for img_idx in range(len(paired_image_path_list)): | |
| self.paired_image_list += [self.read_image(paired_image_path_list[img_idx])] | |
| self.save_name = paired_image_path_list[0].split("/")[-1] | |
| def concat_format_grids_byorder(self, row_num, col_num): | |
| # Stack format | |
| res = None | |
| total_img_number = len(self.paired_image_list) | |
| added_blank_img_number = row_num * col_num - total_img_number | |
| self.paired_image_list += [ self.blank_img for added_blank_img_idx in range(added_blank_img_number)] | |
| img_idx = 0 | |
| row_res = None | |
| col_res = None | |
| for row_idx in range(row_num): | |
| for col_idx in range(col_num): | |
| if col_idx == 0: | |
| col_res = self.paired_image_list[img_idx] | |
| else: | |
| col_res = np.hstack([col_res, self.paired_image_list[img_idx]]) | |
| img_idx += 1 | |
| if row_idx == 0: | |
| row_res = col_res | |
| else: | |
| row_res = np.vstack([row_res, col_res]) | |
| cv2.imwrite(os.path.join(self.save_dir, self.save_name), row_res) | |
| def concat_format_custom(self, ): | |
| pass | |
| def read_image(self, img_path): | |
| return cv2.resize(cv2.imread(img_path, self.cv2_img_read_type), self.cv2_img_read_size) | |
| def add_text(self, img, text): | |
| return cv2.putText(img, text, self.text_pos, self.text_font, self.text_size, self.text_color, self.text_overstriking) | |
| class ConcatVideos: | |
| def __init__(self, ): | |
| self.save_dir = None | |
| self.save_name = None | |
| self.video_path_list = [] | |
| self.total_video_number = 0 | |
| self.ffmpeg_cmd_str = None | |
| self.alphabet_mapping_idx = 97 | |
| self.flag_readd_video_text = True | |
| # Text config | |
| self.text_pos_x = 10 | |
| self.text_pos_y = 30 | |
| self.text_color = "red" | |
| self.text_size = 22 | |
| self.text_font = self.set_text_font("/usr/share/fonts/") | |
| def run_concat_videos(self, save_dir, paired_video_path_list, paired_video_text_list=None, row_num=0, col_num=0, flag_readd_video_name=True): | |
| self.save_dir = save_dir | |
| os.makedirs(self.save_dir, exist_ok=True) | |
| flag_add_text = False | |
| if paired_video_text_list == None: | |
| flag_add_text = False | |
| else: | |
| assert len(paired_video_path_list) == len(paired_video_text_list), "Not match len(paired_video_path_list) %d and len(paired_video_text_list) %d".format(len(paired_video_path_list), len(paired_video_path_list)) | |
| flag_add_text = True | |
| if flag_add_text == True: | |
| self.video_path_list = self.add_video_text(paired_video_path_list, paired_video_text_list) | |
| assert len(self.video_path_list) == len(paired_video_path_list) | |
| else: | |
| self.video_path_list = paired_video_path_list | |
| self.flag_readd_video_text = flag_readd_video_name | |
| if self.flag_readd_video_text == True: | |
| assert paired_video_text_list is not None, "Must have paired_video_text_list." | |
| self.total_video_number = len(paired_video_path_list) | |
| assert self.total_video_number <= row_num * col_num, "Error video number {} > (row: {} x col: {})".format("%d"%self.total_video_number, "%d"%row_num, "%d"%col_num) | |
| self.concat_format_grids_byorder(row_num, col_num) | |
| def concat_format_grids_byorder(self, row_num, col_num): | |
| self.ffmpeg_cmd_str = "ffmpeg " | |
| for video_fn_path in self.video_path_list: | |
| self.ffmpeg_cmd_str += "-i " + video_fn_path + " " | |
| # Video concat format | |
| self.ffmpeg_cmd_str += "-filter_complex '" | |
| video_cnt = 0 | |
| for t_row in range(row_num): | |
| for t_col in range(col_num): | |
| if video_cnt == 0: | |
| self.ffmpeg_cmd_str += "[0:v]pad=iw*{}:ih*{}".format(col_num, row_num) | |
| else: | |
| self.ffmpeg_cmd_str += "[{}];[{}][{}:v]overlay=w*{}:h*{}".format(chr(self.alphabet_mapping_idx + video_cnt), chr(self.alphabet_mapping_idx + video_cnt), video_cnt, t_col, t_row) | |
| video_cnt += 1 | |
| if video_cnt >= self.total_video_number: | |
| break | |
| self.ffmpeg_cmd_str += "'" | |
| # Video concat config | |
| self.ffmpeg_cmd_str += " -c:v libx264 -crf 22 -preset veryfast -y " | |
| self.ffmpeg_cmd_str += os.path.join(self.save_dir, "output.mp4") | |
| print("self.ffmpeg_cmd_str", self.ffmpeg_cmd_str) | |
| os.system(self.ffmpeg_cmd_str) | |
| def add_video_text(self, video_path_list, video_name_list): | |
| video_w_name_path_list = [] | |
| for video_path, video_name in zip(video_path_list, video_name_list): | |
| assert os.path.isfile(video_path), "When adding video text, Not a video file path {}.".format(video_path) | |
| video_fn = video_path.split("/")[-1] | |
| video_w_name_fn = os.path.join(os.path.dirname(video_path), video_fn.split(".")[0] + "_w_name." + video_fn.split(".")[-1]) | |
| if self.flag_readd_video_text == True: | |
| print("ffmpeg -i {} -vf drawtext=fontfile={}:fontcolor={}:fontsize={}:text='{}':x={}:y={} -y {}".format(video_path, self.text_font, self.text_color, self.text_size, video_name, self.text_pos_x, self.text_pos_y, video_w_name_fn)) | |
| os.system("ffmpeg -i {} -vf drawtext=fontfile={}:fontcolor={}:fontsize={}:text='{}':x={}:y={} -y {}".format(video_path, self.text_font, self.text_color, self.text_size, video_name, self.text_pos_x, self.text_pos_y, video_w_name_fn)) | |
| video_w_name_path_list += [video_w_name_fn] | |
| return video_w_name_path_list | |
| def set_text_font(self, font_path): | |
| for root, dirs, files in os.walk(font_path): | |
| if "DejaVuSans.ttf" in files: | |
| print("Successfully set text font in ffmpeg: {}".format("%s"%os.path.join(root, "DejaVuSans.ttf"))) | |
| return os.path.join(root, "DejaVuSans.ttf") | |
| assert False, "Not exist DejaVuSans.ttf in {}".format("%s"%font_path) | |