|
import gradio as gr |
|
|
|
import os |
|
import cv2 |
|
import numpy as np |
|
from PIL import Image |
|
from moviepy.editor import * |
|
|
|
import sys |
|
sys.path.append('FollowYourPose') |
|
|
|
def get_frames(video_in): |
|
frames = [] |
|
|
|
clip = VideoFileClip(video_in) |
|
start_frame = 0 |
|
end_frame = 50 |
|
|
|
if not os.path.exists('./raw_frames'): |
|
os.makedirs('./raw_frames') |
|
|
|
if not os.path.exists('./mmpose_frames'): |
|
os.makedirs('./mmpose_frames') |
|
|
|
|
|
if clip.fps > 30: |
|
print("vide rate is over 30, resetting to 30") |
|
clip_resized = clip.resize(height=512) |
|
clip_resized = clip_resized.subclip(start_frame / clip_resized.fps, end_frame / clip_resized.fps) |
|
clip_resized.write_videofile("./video_resized.mp4", fps=30) |
|
else: |
|
print("video rate is OK") |
|
clip_resized = clip.resize(height=512) |
|
clip_resized = clip_resized.subclip(start_frame / clip.fps, end_frame / clip.fps) |
|
clip_resized.write_videofile("./video_resized.mp4", fps=clip.fps) |
|
|
|
print("video resized to 512 height") |
|
|
|
|
|
cap= cv2.VideoCapture("./video_resized.mp4") |
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS) |
|
print("video fps: " + str(fps)) |
|
i=0 |
|
while(cap.isOpened()): |
|
ret, frame = cap.read() |
|
if ret == False: |
|
break |
|
cv2.imwrite('./raw_frames/kang'+str(i)+'.jpg',frame) |
|
frames.append('./raw_frames/kang'+str(i)+'.jpg') |
|
i+=1 |
|
|
|
cap.release() |
|
cv2.destroyAllWindows() |
|
print("broke the video into frames") |
|
|
|
return frames, fps |
|
|
|
def get_mmpose_filter(mmpose, i): |
|
|
|
|
|
|
|
image = mmpose(i, fn_index=0)[1] |
|
image = Image.open(image) |
|
|
|
image.save("./mmpose_frames/mmpose_frame_" + str(i).split('/')[-1][:-4] + ".jpeg") |
|
return "./mmpose_frames/mmpose_frame_" + str(i).split('/')[-1][:-4] + ".jpeg" |
|
|
|
def create_video(frames, fps, type): |
|
print("building video result") |
|
clip = ImageSequenceClip(frames, fps=fps) |
|
clip.write_videofile(type + "_result.mp4", fps=fps) |
|
|
|
return type + "_result.mp4" |
|
|
|
|
|
def infer_skeleton(mmpose, video_in): |
|
|
|
|
|
|
|
|
|
break_vid = get_frames(video_in) |
|
frames_list= break_vid[0] |
|
fps = break_vid[1] |
|
|
|
n_frame = len(frames_list) |
|
|
|
if n_frame >= len(frames_list): |
|
print("video is shorter than the cut value") |
|
n_frame = len(frames_list) |
|
|
|
|
|
result_frames = [] |
|
print("set stop frames to: " + str(n_frame)) |
|
|
|
for i in frames_list[0:int(n_frame)]: |
|
mmpose_frame = get_mmpose_filter(mmpose, i) |
|
result_frames.append(mmpose_frame) |
|
print("frame " + i + "/" + str(n_frame) + ": done;") |
|
|
|
|
|
final_vid = create_video(result_frames, fps, "mmpose") |
|
files = [final_vid] |
|
|
|
return final_vid, files |