Spaces:
Running
Running
File size: 6,868 Bytes
95f8bbc 2ee3801 95f8bbc 2ee3801 95f8bbc aa34300 95f8bbc 2ee3801 95f8bbc aa34300 95f8bbc |
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 |
import ntpath
import os
import shutil
import numpy as np
import torch.utils.data
from tqdm import tqdm
from SPPE.src.main_fast_inference import *
from common.utils import calculate_area
from dataloader import DetectionLoader, DetectionProcessor, DataWriter, Mscoco, VideoLoader
from fn import getTime
from opt import opt
from pPose_nms import write_json
args = opt
args.dataset = 'coco'
args.fast_inference = False
args.save_img = True
args.sp = True
if not args.sp:
torch.multiprocessing.set_start_method('forkserver', force=True)
torch.multiprocessing.set_sharing_strategy('file_system')
def model_load():
model = None
return model
def image_interface(model, image):
pass
def generate_kpts(video_file, progress):
final_result, video_name = handle_video(video_file, progress)
# ============ Changing ++++++++++
kpts = []
no_person = []
for i in range(len(final_result)):
if not final_result[i]['result']: # No people
no_person.append(i)
kpts.append(None)
continue
kpt = max(final_result[i]['result'],
key=lambda x: x['proposal_score'].data[0] * calculate_area(x['keypoints']), )['keypoints']
kpts.append(kpt.data.numpy())
for n in no_person:
kpts[n] = kpts[-1]
no_person.clear()
for n in no_person:
kpts[n] = kpts[-1] if kpts[-1] else kpts[n-1]
# ============ Changing End ++++++++++
name = f'{args.outputpath}/{video_name}.npz'
kpts = np.array(kpts).astype(np.float32)
print('kpts npz save in ', name)
np.savez_compressed(name, kpts=kpts)
return kpts
def handle_video(video_file, progress):
# =========== common ===============
args.video = video_file
base_name = os.path.basename(args.video)
video_name = base_name[:base_name.rfind('.')]
# =========== end common ===============
# =========== image ===============
# img_path = f'outputs/alpha_pose_{video_name}/split_image/'
# args.inputpath = img_path
# args.outputpath = f'outputs/alpha_pose_{video_name}'
# if os.path.exists(args.outputpath):
# shutil.rmtree(f'{args.outputpath}/vis', ignore_errors=True)
# else:
# os.mkdir(args.outputpath)
#
# # if not len(video_file):
# # raise IOError('Error: must contain --video')
#
# if len(img_path) and img_path != '/':
# for root, dirs, files in os.walk(img_path):
# im_names = sorted([f for f in files if 'png' in f or 'jpg' in f])
# else:
# raise IOError('Error: must contain either --indir/--list')
#
# # Load input images
# data_loader = ImageLoader(im_names, batchSize=args.detbatch, format='yolo').start()
# print(f'Totally {data_loader.datalen} images')
# =========== end image ===============
# =========== video ===============
args.outputpath = f'output_alphapose/{video_name}'
if os.path.exists(args.outputpath):
shutil.rmtree(f'{args.outputpath}/vis', ignore_errors=True)
else:
os.mkdir(args.outputpath)
videofile = args.video
mode = args.mode
if not len(videofile):
raise IOError('Error: must contain --video')
# Load input video
data_loader = VideoLoader(videofile, batchSize=args.detbatch).start()
(fourcc, fps, frameSize) = data_loader.videoinfo()
print('the video is {} f/s'.format(fps))
# =========== end video ===============
# Load detection loader
print('Loading YOLO model..')
sys.stdout.flush()
det_loader = DetectionLoader(data_loader, batchSize=args.detbatch).start()
# start a thread to read frames from the file video stream
det_processor = DetectionProcessor(det_loader).start()
# Load pose model
pose_dataset = Mscoco()
if args.fast_inference:
pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
else:
pose_model = InferenNet(4 * 1 + 1, pose_dataset)
pose_model
pose_model.eval()
runtime_profile = {
'dt': [],
'pt': [],
'pn': []
}
# Data writer
save_path = os.path.join(args.outputpath, 'AlphaPose_' + ntpath.basename(video_file).split('.')[0] + '.avi')
# writer = DataWriter(args.save_video, save_path, cv2.VideoWriter_fourcc(*'XVID'), fps, frameSize).start()
writer = DataWriter(args.save_video).start()
print('Start pose estimation...')
im_names_desc = tqdm(range(data_loader.length()))
batchSize = args.posebatch
for i in progress.tqdm(range(data_loader.length()), desc="Step 1: 2D Detecting"):
start_time = getTime()
with torch.no_grad():
(inps, orig_img, im_name, boxes, scores, pt1, pt2) = det_processor.read()
if orig_img is None:
print(f'{i}-th image read None: handle_video')
break
if boxes is None or boxes.nelement() == 0:
writer.save(None, None, None, None, None, orig_img, im_name.split('/')[-1])
continue
ckpt_time, det_time = getTime(start_time)
runtime_profile['dt'].append(det_time)
# Pose Estimation
datalen = inps.size(0)
leftover = 0
if datalen % batchSize:
leftover = 1
num_batches = datalen // batchSize + leftover
hm = []
for j in range(num_batches):
inps_j = inps[j * batchSize:min((j + 1) * batchSize, datalen)]
hm_j = pose_model(inps_j)
hm.append(hm_j)
hm = torch.cat(hm)
ckpt_time, pose_time = getTime(ckpt_time)
runtime_profile['pt'].append(pose_time)
hm = hm.cpu().data
writer.save(boxes, scores, hm, pt1, pt2, orig_img, im_name.split('/')[-1])
ckpt_time, post_time = getTime(ckpt_time)
runtime_profile['pn'].append(post_time)
if args.profile:
# TQDM
im_names_desc.set_description(
'det time: {dt:.4f} | pose time: {pt:.4f} | post processing: {pn:.4f}'.format(
dt=np.mean(runtime_profile['dt']), pt=np.mean(runtime_profile['pt']), pn=np.mean(runtime_profile['pn']))
)
if (args.save_img or args.save_video) and not args.vis_fast:
print('===========================> Rendering remaining images in the queue...')
print('===========================> If this step takes too long, you can enable the --vis_fast flag to use fast rendering (real-time).')
while writer.running():
pass
writer.stop()
final_result = writer.results()
write_json(final_result, args.outputpath)
return final_result, video_name
if __name__ == "__main__":
os.chdir('../..')
print(os.getcwd())
# handle_video(img_path='outputs/image/kobe')
generate_kpts('outputs/dance.mp4')
|