File size: 10,992 Bytes
face1bf 897d384 face1bf 897d384 face1bf 897d384 face1bf 897d384 face1bf 897d384 face1bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 |
# Pose inferencing
import mmpose
from mmpose.apis import MMPoseInferencer
# Ultralytics
from ultralytics import YOLO
import torch
# Gradio
import gradio as gr
import moviepy.editor as moviepy
# System and files
import os
import glob
import uuid
# Image manipulation
import numpy as np
import cv2
print(torch.__version__)
# Use GPU if available
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
os.system("nvidia-smi")
print("[INFO]: Imported modules!")
human = MMPoseInferencer("human")
hand = MMPoseInferencer("hand")
human3d = MMPoseInferencer(pose3d="human3d")
track_model = YOLO('yolov8n.pt') # Load an official Detect model
print("[INFO]: Downloaded models!")
def check_extension(video):
split_tup = os.path.splitext(video)
# extract the file name and extension
file_name = split_tup[0]
file_extension = split_tup[1]
if file_extension != ".mp4":
print("Converting to mp4")
clip = moviepy.VideoFileClip(video)
video = file_name+".mp4"
clip.write_videofile(video)
return video
def tracking(video, model, boxes=True):
print("[INFO] Is cuda available? ", torch.cuda.is_available())
print(device)
print("[INFO] Loading model...")
# Load an official or custom model
# Perform tracking with the model
print("[INFO] Starting tracking!")
# https://docs.ultralytics.com/modes/predict/
annotated_frame = model(video, boxes=boxes, device=device)
return annotated_frame
def show_tracking(video_content):
# https://docs.ultralytics.com/datasets/detect/coco/
video = cv2.VideoCapture(video_content)
# Track
video_track = tracking(video_content, track_model.track)
# Prepare to save video
#out_file = os.path.join(vis_out_dir, "track.mp4")
out_file = "track.mp4"
print("[INFO]: TRACK", out_file)
fourcc = cv2.VideoWriter_fourcc(*"mp4v") # Codec for MP4 video
fps = video.get(cv2.CAP_PROP_FPS)
height, width, _ = video_track[0][0].orig_img.shape
size = (width,height)
out_track = cv2.VideoWriter(out_file, fourcc, fps, size)
# Go through frames and write them
for frame_track in video_track:
result_track = frame_track[0].plot() # plot a BGR numpy array of predictions
out_track.write(result_track)
print("[INFO] Done with frames")
#print(type(result_pose)) numpy ndarray
out_track.release()
video.release()
cv2.destroyAllWindows() # Closing window
return out_file
def pose3d(video):
video = check_extension(video)
print(device)
# Define new unique folder
add_dir = str(uuid.uuid4())
vis_out_dir = os.path.join("/".join(video.split("/")[:-1]), add_dir)
os.makedirs(vis_out_dir)
result_generator = human3d(video,
vis_out_dir = vis_out_dir,
thickness=4,
radius = 5,
return_vis=True,
kpt_thr=0.3,
rebase_keypoint_height=True,
device=device)
result = [result for result in result_generator] #next(result_generator)
out_file = glob.glob(os.path.join(vis_out_dir, "*.mp4")) #+ glob.glob(os.path.join(vis_out_dir, "*.webm"))
# Reinitialize
return "".join(out_file)
def pose2d(video, kpt_threshold):
video = check_extension(video)
print(device)
# Define new unique folder
add_dir = str(uuid.uuid4())
vis_out_dir = os.path.join("/".join(video.split("/")[:-1]), add_dir)
os.makedirs(add_dir)
result_generator = human(video,
vis_out_dir = add_dir,
return_vis=True,
radius = 5,
thickness=4,
rebase_keypoint_height=True,
kpt_thr=kpt_threshold,
device=device,
pred_out_dir = add_dir
)
result = [result for result in result_generator] #next(result_generator)
out_file = glob.glob(os.path.join(add_dir, "*.mp4")) #+ glob.glob(os.path.join(vis_out_dir, "*.webm"))
kpoints = glob.glob(os.path.join(add_dir, "*.json"))
return "".join(out_file), "".join(kpoints)
def pose2dhand(video, kpt_threshold):
video = check_extension(video)
print(device)
# ultraltics
# Define new unique folder
add_dir = str(uuid.uuid4())
vis_out_dir = os.path.join("/".join(video.split("/")[:-1]), add_dir)
os.makedirs(vis_out_dir)
result_generator = hand(video,
vis_out_dir = vis_out_dir,
return_vis=True,
thickness = 4,
radius = 5,
rebase_keypoint_height=True,
kpt_thr=kpt_threshold,
device=device)
result = [result for result in result_generator] #next(result_generator)
out_file = glob.glob(os.path.join(vis_out_dir, "*.mp4")) #+ glob.glob(os.path.join(vis_out_dir, "*.webm"))
return "".join(out_file)
block = gr.Blocks()
with block:
with gr.Column():
with gr.Tab("Upload video"):
with gr.Column():
with gr.Row():
with gr.Column():
video_input = gr.Video(source="upload", type="filepath", height=612)
# Insert slider with kpt_thr
file_kpthr = gr.Slider(0, 1, value=0.3, label='Keypoint threshold')
with gr.Row():
submit_pose_file = gr.Button("Make 2d pose estimation", variant="primary")
submit_pose3d_file = gr.Button("Make 3d pose estimation", variant="primary")
submit_hand_file = gr.Button("Make 2d hand estimation", variant="primary")
submit_detect_file = gr.Button("Detect and track objects", variant="primary")
with gr.Row():
video_output1 = gr.PlayableVideo(height=512, label = "Estimate human 2d poses", show_label=True)
video_output2 = gr.PlayableVideo(height=512, label = "Estimate human 3d poses", show_label=True)
video_output3 = gr.PlayableVideo(height=512, label = "Estimate human hand poses", show_label=True)
video_output4 = gr.Video(height=512, label = "Detection and tracking", show_label=True, format="mp4")
jsonoutput = gr.Code()
with gr.Tab("General information"):
gr.Markdown("""
\n # Information about the models
\n ## Pose models:
\n All the pose estimation models come from the library [MMpose](https://github.com/open-mmlab/mmpose). It is a library for human pose estimation that provides pre-trained models for 2D and 3D pose estimation.
\n The 2D pose model is used for estimating the 2D coordinates of human body joints from an image or a video frame. The model uses a convolutional neural network (CNN) to predict the joint locations and their confidence scores.
\n The 2D hand model is a specialized version of the 2D pose model that is designed for hand pose estimation. It uses a similar CNN architecture to the 2D pose model but is trained specifically for detecting the joints in the hand.
\n The 3D pose model is used for estimating the 3D coordinates of human body joints from an image or a video frame. The model uses a combination of 2D pose estimation and depth estimation to infer the 3D joint locations.
\n
\n ## Detection and tracking:
\n The tracking method in the Ultralight's YOLOv8 model is used for object tracking in videos. It takes a video file or a camera stream as input and returns the tracked objects in each frame. The method uses the COCO dataset classes for object detection and tracking.
\n The COCO dataset contains 80 classes of objects such as person, car, bicycle, etc. See https://docs.ultralytics.com/datasets/detect/coco/ for all available classes. The tracking method uses the COCO classes to detect and track the objects in the video frames. The tracked objects are represented as bounding boxes with labels indicating the class of the object.""")
gr.Markdown("You can load the keypoints in python in the following way: ")
gr.Code(
value="""def hello_world():
return "Hello, world!"
print(hello_world())""",
language="python",
interactive=True,
show_label=False,
)
# From file
submit_pose_file.click(fn=pose2d,
inputs= [video_input, file_kpthr],
outputs = [video_output1, jsonoutput],
queue=False)
submit_pose3d_file.click(fn=pose3d,
inputs= video_input,
outputs = video_output2,
queue=False)
submit_hand_file.click(fn=pose2dhand,
inputs= [video_input, file_kpthr],
outputs = video_output3,
queue=False)
submit_detect_file.click(fn=show_tracking,
inputs= video_input,
outputs = video_output4,
queue=False)
if __name__ == "__main__":
block.queue(concurrency_count=10, # When you increase the concurrency_count parameter in queue(), max_threads() in launch() is automatically increased as well.
max_size=25, # Maximum number of requests that the queue processes
api_open = False # When creating a Gradio demo, you may want to restrict all traffic to happen through the user interface as opposed to the programmatic API that is automatically created for your Gradio demo.
) # https://www.gradio.app/guides/setting-up-a-demo-for-maximum-performance
block.launch(server_name="0.0.0.0", server_port=7860)
|