Update app.py

app.py

@@ -19,15 +19,11 @@ for model_id in model_ids:
 
 
 
-def get_frame_count_in_duration(filepath):
-    video = cv2.VideoCapture(filepath)
-    fps = video.get(cv2.CAP_PROP_FPS)
+def get_frame_count(filepath):
+    video = cv2.VideoCapture(filepath) 
     frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
-    duration = frame_count / fps
-    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
     video.release()
-    return gr.update(visible=False), gr.update(visible=True), gr.update(maximum=frame_count)
+    return gr.update(maximum=frame_count)
 
 def get_video_dimension(filepath):
     video = cv2.VideoCapture(filepath)
@@ -38,40 +34,17 @@ def get_video_dimension(filepath):
     video.release()
     return width, height, fps, frame_count
 
-def adjust_to_multiple_of_12(number):
-    remainder = number % 12
-    if remainder != 0:
-        adjustment = 12 - remainder
-        number += adjustment
-    return number
-    
-def resize_video(input_file):
-    # Load the video clip
-    clip = VideoFileClip(input_file)
-    print(f"WIDTH TARGET: 512")
-    # Calculate the aspect ratio
-    ratio = 512 / clip.size[0]
-    new_height = int(clip.size[1] * ratio)
-    new_height_adjusted = adjust_to_multiple_of_12(new_height)
-    new_width_adjusted = adjust_to_multiple_of_12(512)
-    print(f"OLD H: {new_height} | NEW H: {new_height_adjusted}")
-    print(f"OLD W: 512 | NEW W: {new_width_adjusted}")
-
-    # Close the video clip
-    clip.close()
-
+def resize_video(input_vid, output_vid, width, height, fps):
     # Open the input video file
-    video = cv2.VideoCapture(input_file)
+    video = cv2.VideoCapture(input_vid)
+
+    # Get the original video's width and height
+    original_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
+    original_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
     # Create a VideoWriter object to write the resized video
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for the output video
-   
-    # Check if the file already exists
-    if os.path.exists('video_resized.mp4'):
-        # Delete the existing file
-        os.remove('video_resized.mp4')
-        
-    output_video = cv2.VideoWriter('video_resized.mp4', fourcc, 8.0, (512, 512))
+    output_video = cv2.VideoWriter(output_vid, fourcc, fps, (width, height))
 
     while True:
         # Read a frame from the input video
@@ -80,7 +53,7 @@ def resize_video(input_file):
             break
 
         # Resize the frame to the desired dimensions
-        resized_frame = cv2.resize(frame, (512, 512))
+        resized_frame = cv2.resize(frame, (width, height))
 
         # Write the resized frame to the output video file
         output_video.write(resized_frame)
@@ -89,56 +62,62 @@ def resize_video(input_file):
     video.release()
     output_video.release()
 
-    
-    
-    #final_video_resized = os.path.join(temp_output_path, 'video_resized.mp4')
-    test_w, test_h, fps, frame_count = get_video_dimension('video_resized.mp4')
-    print(f"resized clip dims : {test_w}, {test_h}, {fps}")
-    return gr.update(visible=False), gr.update(value='video_resized.mp4', visible=True), gr.update(maximum=frame_count)
-        
-def run_inference(prompt, video_path, condition, video_length):
+    return output_vid
 
-    output_path = 'output/'
-    os.makedirs(output_path, exist_ok=True)
+def chunkify(video_path, fps, nb_frames):
+    chunks_array = []
 
-    # Construct the final video path
-    video_path_output = os.path.join(output_path, f"{prompt}.mp4")
+    video_capture = cv2.VideoCapture(video_path)
+    chunk_start_frame = 0
+    frames_per_chunk = 12
 
-    # Check if the file already exists
-    if os.path.exists(video_path_output):
-        # Delete the existing file
-        os.remove(video_path_output)
+    while chunk_start_frame < nb_frames:
+        chunk_end_frame = min(chunk_start_frame + frames_per_chunk, total_frames)
 
-    if video_length > 12:
-        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{video_path}' --output_path '{output_path}' --width 512 --height 512 --fps 8 --video_length {video_length} --is_long_video"
-    else:
-        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{video_path}' --output_path '{output_path}' --width 512 --height 512 --fps 8 --video_length {video_length}"
-    subprocess.run(command, shell=True)
+        video_capture.set(cv2.CAP_PROP_POS_FRAMES, chunk_start_frame)
+        success, frame = video_capture.read()
+        if not success:
+            break
 
-    # Construct the video path
-    video_path_output = os.path.join(output_path, f"{prompt}.mp4")
+        chunk_name = f"chunk_{chunk_start_frame}-{chunk_end_frame}.mp4"
+        chunk_video = cv2.VideoWriter(chunk_name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (frame.shape[1], frame.shape[0]))
+
+        for frame_number in range(chunk_start_frame, chunk_end_frame):
+            video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
+            success, frame = video_capture.read()
+            if not success:
+                break
+
+            chunk_video.write(frame)
+
+        chunk_video.release()
+        chunks_array.append(chunk_name)
 
+        chunk_start_frame += frames_per_chunk
+
+    video_capture.release()
+    print(f"CHUNKS: {chunks_array}")
+    return chunks_array
+    
         
+def run_inference(prompt, video_path, condition, video_length):
 
-    return "done", video_path_output 
+    # Get FPS of original video input
+    target_fps = get_video_dimension(video_path)[2] 
+    print(f"INPUT FPS: {target_fps}")
+    
+    # Count total frames according to fps
+    total_frames = get_video_dimension(video_path)[3]
 
-def run_inference_chunks(prompt, video_path, condition, video_length):
+    # Resize the video
+    resized = resize_video(video_path, 'resized.mp4', 512, 512, target_fps)
 
-    # Specify the input and output paths
-    input_vid = video_path
-    resized_vid = 'resized.mp4'
+    # Chunkify the video into 12 frames chunks
+    chunks = chunkify(resized, target_fps, total_frames)
 
-    # Call the function to resize the video
-    video_path = resize_video(input_vid, resized_vid, width=512)
-    width, height, fps = get_video_dimension(video_path)
 
-    print(f"{width} x {height} | {fps}")
 
-    # Split the video into chunks mp4 of 12 frames at video fps
-    # Store chunks as mp4 paths in an array
 
-    # For each mp4 chunks in chunks arrays, run command
-    # store video result in processed chunks array
     
     output_path = 'output/'
     os.makedirs(output_path, exist_ok=True)
@@ -152,18 +131,18 @@ def run_inference_chunks(prompt, video_path, condition, video_length):
         os.remove(video_path_output)
 
     if video_length > 12:
-        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{video_path}' --output_path '{output_path}' --width {width} --height {height} --fps {fps} --video_length {video_length} --is_long_video"
+        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{video_path}' --output_path '{output_path}' --width 512 --height 512 --fps 8 --video_length {video_length} --is_long_video"
     else:
-        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{video_path}' --output_path '{output_path}' --width {width} --height {height} --fps {fps} --video_length {video_length}"
+        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{video_path}' --output_path '{output_path}' --width 512 --height 512 --fps 8 --video_length {video_length}"
     subprocess.run(command, shell=True)
 
     # Construct the video path
     video_path_output = os.path.join(output_path, f"{prompt}.mp4")
 
-        
-
     return "done", video_path_output 
 
+ 
+
 css="""
 #col-container {max-width: 810px; margin-left: auto; margin-right: auto;}
 """
@@ -174,8 +153,8 @@ with gr.Blocks(css=css) as demo:
         """)
         with gr.Row():
             with gr.Column():
-                video_in = gr.Video(source="upload", type="filepath", visible=True)
-                video_path = gr.Video(source="upload", type="filepath", visible=False)
+                #video_in = gr.Video(source="upload", type="filepath", visible=True)
+                video_path = gr.Video(source="upload", type="filepath", visible=True)
                 prompt = gr.Textbox(label="prompt")
                 with gr.Row():
                     condition = gr.Dropdown(label="Condition", choices=["depth", "canny", "pose"], value="depth")
@@ -185,9 +164,9 @@ with gr.Blocks(css=css) as demo:
             with gr.Column():
                 video_res = gr.Video(label="result")
                 status = gr.Textbox(label="result")
-    video_in.change(fn=resize_video,
-                      inputs=[video_in],
-                      outputs=[video_in, video_path, video_length]
+    video_path.change(fn=get_frame_count,
+                      inputs=[video_path],
+                      outputs=[video_length]
                      )
     submit_btn.click(fn=run_inference, 
                      inputs=[prompt,