app is mcp compatible

app.py

@@ -6,16 +6,24 @@ import numpy as np
 from PIL import Image
 from moviepy.editor import *
 
+# Load the OpenPose detector
 openpose = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
 
 def get_frames(video_in):
+    """Extract frames from a video and resize to height 512px.
+    
+    Args:
+        video_in: Path to the input video file.
+        
+    Returns:
+        frames: List of paths to extracted frame images.
+        fps: Frames per second of the original video.
+    """
     frames = []
-    #resize the video
     clip = VideoFileClip(video_in)
     
-    #check fps
     if clip.fps > 30:
-        print("vide rate is over 30, resetting to 30")
+        print("video rate is over 30, resetting to 30")
         clip_resized = clip.resize(height=512)
         clip_resized.write_videofile("video_resized.mp4", fps=30)
     else:
@@ -23,97 +31,103 @@ def get_frames(video_in):
         clip_resized = clip.resize(height=512)
         clip_resized.write_videofile("video_resized.mp4", fps=clip.fps)
     
-    print("video resized to 512 height")
-    
-    # Opens the Video file with CV2
-    cap= cv2.VideoCapture("video_resized.mp4")
-    
+    cap = cv2.VideoCapture("video_resized.mp4")
     fps = cap.get(cv2.CAP_PROP_FPS)
-    print("video fps: " + str(fps))
-    i=0
+    i = 0
     while(cap.isOpened()):
         ret, frame = cap.read()
-        if ret == False:
+        if not ret:
             break
-        cv2.imwrite('kang'+str(i)+'.jpg',frame)
-        frames.append('kang'+str(i)+'.jpg')
-        i+=1
-    
+        frame_path = f'frame_{i}.jpg'
+        cv2.imwrite(frame_path, frame)
+        frames.append(frame_path)
+        i += 1
+
     cap.release()
     cv2.destroyAllWindows()
-    print("broke the video into frames")
-    
     return frames, fps
 
 def get_openpose_filter(i):
-    image = Image.open(i)
+    """Apply OpenPose filter to a single image frame.
     
-    #image = np.array(image)
-
+    Args:
+        i: Path to the image frame.
+        
+    Returns:
+        output_path: Path to the OpenPose-processed image.
+    """
+    image = Image.open(i)
     image = openpose(image)
-    #image = Image.fromarray(image)
-    image.save("openpose_frame_" + str(i) + ".jpeg")
-    return "openpose_frame_" + str(i) + ".jpeg"
+    output_path = f"openpose_frame_{os.path.basename(i)}.jpeg"
+    image.save(output_path)
+    return output_path
 
 def create_video(frames, fps, type):
-    print("building video result")
-    clip = ImageSequenceClip(frames, fps=fps)
-    clip.write_videofile(type + "_result.mp4", fps=fps)
+    """Create a video from a sequence of image frames.
     
-    return type + "_result.mp4"
+    Args:
+        frames: List of image frame paths.
+        fps: Frames per second for the output video.
+        type: A string used as the prefix for naming the result video.
+        
+    Returns:
+        video_path: Path to the resulting video file.
+    """
+    clip = ImageSequenceClip(frames, fps=fps)
+    video_path = f"{type}_result.mp4"
+    clip.write_videofile(video_path, fps=fps)
+    return video_path
 
 def convertG2V(imported_gif):
+    """Convert a GIF file to a standard MP4 video.
+    
+    Args:
+        imported_gif: The input GIF file object.
+        
+    Returns:
+        Path to the converted MP4 video file.
+    """
     clip = VideoFileClip(imported_gif.name)
     clip.write_videofile("my_gif_video.mp4")
     return "my_gif_video.mp4"
-    
+
 def infer(video_in):
+    """Generate an OpenPose-filtered video from an input video.
     
+    This function performs the following steps:
+    1. Extracts frames from the input video and retrieves the original frame rate (FPS).
+    2. Applies the OpenPose model to each frame to extract pose data.
+    3. Reconstructs a new video from the OpenPose-processed frames using the original FPS.
+
+    Args:
+        video_in: The uploaded input video file (MP4 or converted GIF).
     
-    # 1. break video into frames and get FPS
+    Returns:
+        final_vid: The path to the OpenPose-filtered output video.
+        files: A list containing the output video file (for download).
+    """
     break_vid = get_frames(video_in)
-    frames_list= break_vid[0]
+    frames_list = break_vid[0]
     fps = break_vid[1]
-    #n_frame = int(trim_value*fps)
     n_frame = len(frames_list)
     
-    if n_frame >= len(frames_list):
-        print("video is shorter than the cut value")
-        n_frame = len(frames_list)
-    
-    # 2. prepare frames result arrays
     result_frames = []
-    print("set stop frames to: " + str(n_frame))
-    
-    for i in frames_list[0:int(n_frame)]:
+    for i in frames_list[:n_frame]:
         openpose_frame = get_openpose_filter(i)
         result_frames.append(openpose_frame)
-        print("frame " + i + "/" + str(n_frame) + ": done;")
 
-    
     final_vid = create_video(result_frames, fps, "openpose")
-
     files = [final_vid]
 
     return final_vid, files
 
-title="""
+# UI layout
+title = """
 <div style="text-align: center; max-width: 500px; margin: 0 auto;">
-        <div
-        style="
-            display: inline-flex;
-            align-items: center;
-            gap: 0.8rem;
-            font-size: 1.75rem;
-            margin-bottom: 10px;
-        "
-        >
-        <h1 style="font-weight: 600; margin-bottom: 7px;">
-            Video to OpenPose
-        </h1>
-        </div>
-       
+    <div style="display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem; margin-bottom: 10px;">
+        <h1 style="font-weight: 600; margin-bottom: 7px;">Video to OpenPose</h1>
     </div>
+</div>
 """
 
 with gr.Blocks() as demo:
@@ -122,14 +136,13 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 video_input = gr.Video(sources=["upload"])
-                gif_input = gr.File(label="import a GIF instead", file_types=['.gif'])
+                gif_input = gr.File(label="Import a GIF instead", file_types=['.gif'])
                 gif_input.change(fn=convertG2V, inputs=gif_input, outputs=video_input)
                 submit_btn = gr.Button("Submit")
-            
             with gr.Column():
                 video_output = gr.Video()
                 file_output = gr.Files()
 
     submit_btn.click(fn=infer, inputs=[video_input], outputs=[video_output, file_output])
 
-demo.launch()
+demo.launch(mcp_server=True)