import streamlit as st import torch import imageio import numpy as np import time from skimage.transform import resize import warnings import cv2 import subprocess import os from demo import load_checkpoints from demo import make_animation from skimage import img_as_ubyte import shutil def save_image_from_upload(uploaded_file): # Specify the path to save save_path = './uploaded_images' if not os.path.exists(save_path): os.makedirs(save_path) # Open the file in the desired location with write-binary ('wb') mode with open(os.path.join(save_path, uploaded_file.name), "wb") as f: f.write(uploaded_file.getbuffer()) # Write the file to the specified location #st.success(f'Saved file {uploaded_file.name} in {save_path}') try: shutil.copy2(os.path.join(save_path, uploaded_file.name), st.session_state['source_image_path']) print("File copied successfully.") except FileNotFoundError: print("The source file was not found.") except PermissionError: print("Permission denied.") except Exception as e: print(f"Error occurred: {e}") def save_video_from_upload(uploaded_file): # Specify the path to save save_path = './uploaded_videos' if not os.path.exists(save_path): os.makedirs(save_path) # Open the file in the desired location with write-binary ('wb') mode with open(os.path.join(save_path, uploaded_file.name), "wb") as f: f.write(uploaded_file.getbuffer()) # Write the file to the specified location # st.success(f'Saved file {uploaded_file.name} in {save_path}') try: shutil.copy2(os.path.join(save_path, uploaded_file.name), st.session_state['driving_video_path']) print("File copied successfully.") except FileNotFoundError: print("The source file was not found.") except PermissionError: print("Permission denied.") except Exception as e: print(f"Error occurred: {e}") def create_image_video_side_by_side(source, driving, generated=None, output_file='assets/output_video.mp4', fps=20, progress_bar=None): #st.image(source,caption='create_image_video_side_by_side src') total_driving_frames = len(driving) progress_bar.progress(0) #images = l images = [] print("going through video") for i in range(len(driving)): cols = [source] cols.append(driving[i]) if generated is not None: #print("generated data length:"+str(len(generated))) #print("cols type"+str(type(cols[i]))) #print("generated[i] shape" + str(generated[i].shape)) cols.append(generated[i]) #print("len(cols) afer append "+str(len(cols))) # else: # print("generated is None!!!") # Concatenate the images horizontally full_image = np.concatenate(cols, axis=1) # Convert the image array to an RGB image full_image_rgb = np.clip(full_image * 255, 0, 255).astype( np.uint8) if full_image.max() <= 1 else full_image.astype(np.uint8) #print("full_image_rgb shape" + str(full_image_rgb.shape)) # Append to the list of images # if i == 0: # source_rgb = np.clip(source * 255, 0, 255).astype( # np.uint8) if source.max() <= 1 else source.astype(np.uint8) # st.image(source_rgb, caption="source_rgb") # driving_0_rgb = np.clip(driving[i] * 255, 0, 255).astype( # np.uint8) if driving[i].max() <= 1 else driving[i].astype(np.uint8) # st.image(driving_0_rgb, caption="driving_0_rgb") # cols_temp = [source_rgb] # cols_temp.append(driving[i]) # full_image_temp = np.concatenate(cols_temp, axis=1) # st.image(full_image_temp, caption="full_image_temp") images.append(full_image_rgb) progress_percentage = (i + 1) / total_driving_frames progress_bar.progress(progress_percentage) print("going through video done") # Determine the size of the frames height, width, layers = images[0].shape print("images[0].shape"+str(images[0].shape)) size = (width, height) # Define the codec and create VideoWriter object fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 'mp4v' or 'XVID' delete_file_if_exists("temp_gen_video.mp4") out = cv2.VideoWriter("temp_gen_video.mp4", fourcc, fps, size) print("writing video start") for image in images: out.write(cv2.cvtColor(image, cv2.COLOR_RGB2BGR)) out.release() # Release the video writer print("writing video end") print("converting video to H264") delete_file_if_exists(output_file) command = ['ffmpeg', '-i', 'temp_gen_video.mp4', '-c:v', 'libx264', output_file] try: result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True) print("FFMPEG Output:\n", result.stdout) print("Video converted successfully.") except subprocess.CalledProcessError as e: # Output the error in case of failure print("Error during conversion:\n", e.stderr) print("converting video to H264 done") return output_file def delete_file_if_exists(file_path): try: if os.path.exists(file_path): print(f"File {file_path} exists and will be deleted.") os.remove(file_path) print(f"File {file_path} has been deleted.") else: print(f"No file found at {file_path}, nothing to delete.") except OSError as e: print(f"Error deleting file {file_path}: {e}") def read_src_image(): print("reading src imaage") st.session_state['source_image'] = imageio.imread(st.session_state['source_image_path']) st.session_state['source_image'] = resize(st.session_state['source_image'], (st.session_state['pixel'], st.session_state['pixel']))[..., :3] # st.session_state['source_image'] = np.clip(st.session_state['source_image'] * 255, 0, 255).astype( # np.uint8) if st.session_state['source_image'].max() <= 1 else st.session_state['source_image'].astype(np.uint8) def read_driving_video(progress_bar=None): reader = imageio.get_reader(st.session_state['driving_video_path']) st.session_state['fps'] = reader.get_meta_data()['fps'] st.session_state['duration'] = reader.get_meta_data()['duration'] video_width = reader.get_meta_data()['source_size'][0] print("st.session_state['duration']="+str(st.session_state['duration'])) st.session_state['driving_video'] = [] print("reading video") # Calculate the number of frames estimated_frame_count = int(st.session_state['fps'] * st.session_state['duration']) print("estimated_frame_count="+str(estimated_frame_count)) progress_bar.progress(0) try: video_frame_idx=0 for im in reader: # im = np.clip(im * 255, 0, 255).astype( # np.uint8) if im.max() <= 1 else im.astype(np.uint8) st.session_state['driving_video'].append(im) progress_percentage = (video_frame_idx) / (estimated_frame_count+1) # print(f"video_frame_idx = {video_frame_idx} estimated_frame_count={estimated_frame_count}") progress_bar.progress(progress_percentage) video_frame_idx = video_frame_idx+1 except RuntimeError: pass reader.close() print("finished reading video") # st.session_state['driving_video'] = [resize(frame, (st.session_state['pixel'], st.session_state['pixel']))[..., :3] # for frame in st.session_state['driving_video']] print("check resize width ="+str(video_width)) # if video_width != 512: progress_bar.progress(0) # Process each frame, update progress bar along the way if video_width != 512: resized_frames = [] num_frames=len(st.session_state['driving_video']) current_status_placeholder.write("resizing video") for i, frame in enumerate(st.session_state['driving_video']): if i == 0: print("frame.dtype="+str(frame.dtype)) # Resize frame resized_frame = resize(frame, (st.session_state['pixel'], st.session_state['pixel']))[..., :3] if i == 0: print("resized_frame.dtype="+str(resized_frame.dtype)) resized_frames.append(resized_frame) # Update progress bar progress_bar.progress((i + 1) / num_frames) # Update the session state with the resized frames st.session_state['driving_video'] = resized_frames else: for i, frame in enumerate(st.session_state['driving_video']): if frame.dtype != np.float64: # Convert to float64 frame_float64 = frame.astype(np.float64) # Normalize the frame based on its original range if frame.dtype == np.uint8: frame_normalized = frame_float64 / 255.0 elif frame.dtype == np.uint16: frame_normalized = frame_float64 / 65535.0 elif frame.dtype == np.float32: # Assuming float32 range is 0.0 to 1.0, similar normalization might not be needed frame_normalized = frame_float64 st.session_state['driving_video'][i] = frame_normalized def add_animation_to_image(): inference_status_placeholder.write("start inference") print("device=" + str(st.session_state['device'])) predictions = make_animation(st.session_state['source_image'], st.session_state['driving_video'], st.session_state['inpainting'], st.session_state['kp_detector'], st.session_state['dense_motion_network'], st.session_state['avd_network'], device=st.session_state['device'], mode=st.session_state['predict_mode'], progress_bar=create_animation_progress_bar) inference_status_placeholder.write("inference done") # save resulting video st.session_state['output_video_path']='assets/generated_video.mp4' st.session_state['side_by_side_with_generated_video_path']='assets/src_image_driving_video_generated_video_side_by_side.mp4' inference_status_placeholder.write("saving generated video") # for i in range(len(predictions)): # predictions[i] = np.clip(predictions[i] * 255, 0, 255).astype( # np.uint8) if predictions[i].max() <= 1 else predictions[i].astype(np.uint8) #st.image(predictions[0], caption="predictions[0]") imageio.mimsave(st.session_state['output_video_path'], [img_as_ubyte(frame) for frame in predictions], fps=st.session_state['fps']) inference_status_placeholder.write("saving generated video done") print("side_by_side_with_generated_video_path="+st.session_state['side_by_side_with_generated_video_path']) inference_status_placeholder.write("creating side by side video") st.session_state['side_by_side_with_generated_video_path'] = create_image_video_side_by_side(st.session_state['source_image'], st.session_state['driving_video'], generated=predictions, output_file=st.session_state['side_by_side_with_generated_video_path'], fps=st.session_state['fps'],progress_bar=create_animation_progress_bar) inference_status_placeholder.write("creating side by side video done") def is_new_src_image_upload(uploaded_file): if 'last_src_image_uploaded_file' in st.session_state: # Check if the newly uploaded file is different from the last one if (uploaded_file.name != st.session_state.last_src_image_uploaded_file['name'] or uploaded_file.size != st.session_state.last_src_image_uploaded_file['size']): st.session_state.last_src_image_uploaded_file = {'name': uploaded_file.name, 'size': uploaded_file.size} # st.write("A new src image file has been uploaded.") return True else: # st.write("The same src image file has been re-uploaded.") return False else: # st.write("This is the first file upload detected.") st.session_state.last_src_image_uploaded_file = {'name': uploaded_file.name, 'size': uploaded_file.size} return True # Store current file details in session state def is_new_driving_video_upload(uploaded_file): if 'last_driving_video_uploaded_file' in st.session_state: # Check if the newly uploaded file is different from the last one if (uploaded_file.name != st.session_state.last_driving_video_uploaded_file['name'] or uploaded_file.size != st.session_state.last_driving_video_uploaded_file['size']): st.session_state.last_driving_video_uploaded_file = {'name': uploaded_file.name, 'size': uploaded_file.size} # st.write("A new driving video file has been uploaded.") return True else: # st.write("The same driving video file has been re-uploaded.") return False else: # st.write("This is the first file upload detected.") st.session_state.last_driving_video_uploaded_file = {'name': uploaded_file.name, 'size': uploaded_file.size} return True big_text = """

Add motions to still images

""" # Display the styled text st.markdown(big_text, unsafe_allow_html=True) #st.markdown("

Add motions to still images

") current_status_placeholder = st.empty() init_progress_bar = st.progress(0) if 'is_initialized' not in st.session_state: st.session_state['is_initialized'] = True #st.set_option('enableStaticServing ', True) print("init") warnings.filterwarnings("ignore") current_status_placeholder.write("checking CUDA availability") if torch.cuda.is_available(): print("CUDA is available on the following devices:") # Loop through available CUDA devices for i in range(torch.cuda.device_count()): print(f"Device {i}: {torch.cuda.get_device_name(i)}") current_status_placeholder.write(f"Device {i}: {torch.cuda.get_device_name(i)}") else: print("CUDA is not available. Listing CPU only.") print("Device 0: CPU") current_status_placeholder.write("CUDA is not available. Listing CPU only.") st.session_state['device'] = torch.device('cuda:0') st.session_state['dataset_name'] = 'vox' # ['vox', 'taichi', 'ted', 'mgif'] st.session_state['source_image_path'] = 'assets/src_image.png' st.session_state['driving_video_path'] = 'assets/driving_video.mp4' st.session_state['side_by_side_video_path'] = 'assets/src_image_driving_video_side_by_side.mp4' st.session_state['uploaded_src_image_file']=False # st.session_state[ # 'side_by_side_with_generated_video_path'] = 'assets/src_image_driving_video_generated_video_side_by_side.mp4' #side_by_side_with_generated_video_path st.session_state['predict_mode'] = 'relative' # ['standard', 'relative', 'avd'] st.session_state['find_best_frame'] = False # when use the relative mode to animate a face, use 'find_best_frame=True' can get better quality result config_path = 'config/vox-256.yaml' checkpoint_path = 'checkpoints/vox.pth.tar' st.session_state['pixel'] = 512 # for vox, taichi and mgif, the resolution is 256*256 print("start loading model") current_status_placeholder.write("start loading model") st.session_state['inpainting'], st.session_state['kp_detector'], st.session_state['dense_motion_network'], st.session_state['avd_network'] = load_checkpoints(config_path=config_path, checkpoint_path=checkpoint_path, device= st.session_state['device'] ) print("finished loading model") current_status_placeholder.write("finished loading model") current_status_placeholder.write("copying default src image") try: shutil.copy2('assets/default_src_image.png', st.session_state['source_image_path']) print("File copied successfully.") except FileNotFoundError: print("The source file was not found.") except PermissionError: print("Permission denied.") except Exception as e: print(f"Error occurred: {e}") try: current_status_placeholder.write("copying default driving video") shutil.copy2('assets/default_driving_video.mp4', st.session_state['driving_video_path']) print("File copied successfully.") except FileNotFoundError: print("The source file was not found.") except PermissionError: print("Permission denied.") except Exception as e: print(f"Error occurred: {e}") current_status_placeholder.write("reading src image") read_src_image() # st.session_state['thumb_source_image'] = resize(st.session_state['source_image'], (250, 250))[..., :3] current_status_placeholder.write("reading driving video") read_driving_video(init_progress_bar) if os.path.exists('assets/default_src_image_driving_video_side_by_side.mp4'): print("deafult side_by_side_video already exists") try: current_status_placeholder.write("copying side by side video") shutil.copy2('assets/default_src_image_driving_video_side_by_side.mp4', st.session_state['side_by_side_video_path']) print("File copied successfully.") except FileNotFoundError: print("The source file was not found.") except PermissionError: print("Permission denied.") except Exception as e: print(f"Error occurred: {e}") else: current_status_placeholder.write("creating side by side video") st.session_state['side_by_side_video_path'] = create_image_video_side_by_side(st.session_state['source_image'], st.session_state['driving_video'], output_file=st.session_state['side_by_side_video_path'], fps=st.session_state['fps'], progress_bar=init_progress_bar) current_status_placeholder.write("") st.video(st.session_state['side_by_side_video_path']) col1, col2 = st.columns(2) with col1: uploaded_src_image_file = st.file_uploader("Upload a source image... image must be square dimension", type=['jpg', 'jpeg', 'png']) st.markdown(f'Sample 1 download and then upload to above', unsafe_allow_html=True) st.markdown(f'Sample 2 download and then upload to above', unsafe_allow_html=True) with col2: uploaded_driving_video_file = st.file_uploader( "Upload a driving video... video must be square dimension... 512x512 recommended", type=['mp4']) st.markdown( f'Sample 1 download and then upload to above', unsafe_allow_html=True) st.markdown( f'Sample 2 download and then upload to above', unsafe_allow_html=True) if uploaded_src_image_file is not None: if is_new_src_image_upload(uploaded_src_image_file): current_status_placeholder.write("checking uploaded source image") save_path = './uploaded_images' if not os.path.exists(save_path): os.makedirs(save_path) # Open the file in the desired location with write-binary ('wb') mode with open(os.path.join(save_path, "temp_"+uploaded_src_image_file.name), "wb") as f: f.write(uploaded_src_image_file.getbuffer()) # Write the file to the specified location # st.success(f'Saved file temp_{uploaded_src_image_file.name} in {save_path}') image = imageio.imread(os.path.join(save_path, "temp_"+uploaded_src_image_file.name)) height, width = image.shape[:2] # To see details #file_details = {"FileName": uploaded_src_image_file.name, "FileType": uploaded_src_image_file.type, "FileSize": uploaded_src_image_file.size} #st.write(file_details) # Save the file if width == height: current_status_placeholder.write("saving uploaded image") save_image_from_upload(uploaded_src_image_file) current_status_placeholder.write("reading uploaded image") read_src_image() current_status_placeholder.write("creating side by side video") st.session_state['side_by_side_video_path'] = create_image_video_side_by_side(st.session_state['source_image'], st.session_state['driving_video'], output_file=st.session_state[ 'side_by_side_video_path'], fps=st.session_state['fps'], progress_bar=init_progress_bar) print("uploaded_src_image_file Done! ") st.rerun(); else: st.error("Error: Image width and height must be equal.") # if not st.session_state['uploaded_src_image_file']: # st.rerun() # st.session_state['uploaded_src_image_file'] = True # st.video(st.session_state['side_by_side_video_path']) if uploaded_driving_video_file is not None: if is_new_driving_video_upload(uploaded_driving_video_file): # To see details # file_details = {"FileName": uploaded_driving_video_file.name, "FileType": uploaded_driving_video_file.type, "FileSize": uploaded_driving_video_file.size} # st.write(file_details) current_status_placeholder.write("checking uploaded video") save_path = './uploaded_videos' if not os.path.exists(save_path): os.makedirs(save_path) # Open the file in the desired location with write-binary ('wb') mode with open(os.path.join(save_path, "temp_uploaded_video.mp4"), "wb") as f: f.write(uploaded_driving_video_file.getbuffer()) # Write the file to the specified location # st.success(f'Saved file "temp_uploaded_video.mp4" in {save_path}') reader = imageio.get_reader(os.path.join(save_path, "temp_uploaded_video.mp4")) video_width = reader.get_meta_data()['source_size'][0] video_height = reader.get_meta_data()['source_size'][1] # Check if dimensions are not equal if video_width != video_height: st.error("Error: Video width and height must be equal.") else: # Display dimensions # st.write(f"Width: {video_width}px") # st.write(f"Height: {video_height}px") current_status_placeholder.write("saving uploaded video") save_video_from_upload(uploaded_driving_video_file) current_status_placeholder.write("reading uploaded video") read_driving_video(init_progress_bar) current_status_placeholder.write("creating side by side video") st.session_state['side_by_side_video_path'] = create_image_video_side_by_side(st.session_state['source_image'], st.session_state['driving_video'], output_file=st.session_state[ 'side_by_side_video_path'], fps=st.session_state['fps'], progress_bar=init_progress_bar) st.rerun() # st.video(st.session_state['side_by_side_video_path']) # x = st.slider('Select a value') # st.write(x, 'squared is', x * x) # Display the video print("st.session_state['side_by_side_video_path'] =" +st.session_state['side_by_side_video_path']) # Create a button and check if the button is clicked inference_status_placeholder = st.empty() if 'run_button' in st.session_state and st.session_state.run_button == True: st.session_state.running = True else: st.session_state.running = False create_animation_progress_bar = st.progress(0) if st.button('Add motion to Image',disabled=st.session_state.running, key='run_button'): add_animation_to_image() st.session_state['video_generated'] = True st.rerun() # What to do after the button is clicked if 'video_generated' in st.session_state: st.video(st.session_state['side_by_side_with_generated_video_path']) st.video(st.session_state['output_video_path']) del st.session_state['video_generated']