def image_preprocess(image_path: str, prev_image: str, next_image: str, center,
scale):
trans_matrix = get_affine_transform(center, scale, 0, image_size)
image_data = read_image(image_path)
image_data = cv2.warpAffine(image_data,
trans_matrix,
(int(image_size[0]), int(image_size[1])),
flags=cv2.INTER_LINEAR)
image_data = image_transforms(image_data)
if prev_image is None or next_image is None:
return image_data
else:
prev_image_data = read_image(prev_image)
prev_image_data = cv2.warpAffine(
prev_image_data,
trans_matrix, (int(image_size[0]), int(image_size[1])),
flags=cv2.INTER_LINEAR)
prev_image_data = image_transforms(prev_image_data)
next_image_data = read_image(next_image)
next_image_data = cv2.warpAffine(
next_image_data,
trans_matrix, (int(image_size[0]), int(image_size[1])),
flags=cv2.INTER_LINEAR)
next_image_data = image_transforms(next_image_data)
return image_data, prev_image_data, next_image_data
def draw_skeleton_in_origin_image(batch_image_list, batch_joints_list, batch_bbox_list, save_dir, vis_skeleton=True, vis_bbox=True):
"""
:param batch_image_list: batch image path
:param batch_joints_list: joints coordinates in image Coordinate reference system
:batch_bbox_list: xyxy
:param save_dir:
:return: No return
"""
skeleton_image_save_folder = osp.join(save_dir, "skeleton")
bbox_image_save_folder = osp.join(save_dir, "bbox")
together_save_folder = osp.join(save_dir, "SkeletonAndBbox")
if vis_skeleton and vis_bbox:
save_folder = together_save_folder
else:
save_folder = skeleton_image_save_folder
if vis_bbox:
save_folder = bbox_image_save_folder
batch_final_coords = batch_joints_list
for index, image_path in enumerate(batch_image_list):
final_coords = batch_final_coords[index]
final_coords = coco2posetrack_ord_infer(final_coords)
bbox = batch_bbox_list[index]
image_name = image_path[image_path.index("images") + len("images") + 1:]
# image_name = image_path[image_path.index("frames") + len("frames") + 1:]
vis_image_save_path = osp.join(save_folder, image_name)
if osp.exists(vis_image_save_path):
processed_image = read_image(vis_image_save_path)
processed_image = add_poseTrack_joint_connection_to_image(processed_image, final_coords, sure_threshold=0.2,
flag_only_draw_sure=True) if vis_skeleton else processed_image
processed_image = add_bbox_in_image(processed_image, bbox) if vis_bbox else processed_image
save_image(vis_image_save_path, processed_image)
else:
image_data = read_image(image_path)
processed_image = image_data.copy()
processed_image = add_poseTrack_joint_connection_to_image(processed_image, final_coords, sure_threshold=0.2,
flag_only_draw_sure=True) if vis_skeleton else processed_image
processed_image = add_bbox_in_image(processed_image, bbox) if vis_bbox else processed_image
save_image(vis_image_save_path, processed_image)
def _get_spatiotemporal_window(self, data_item):
filename = data_item['filename']
img_num = data_item['imgnum']
image_file_path = data_item['image']
num_frames = data_item['nframes']
data_numpy = read_image(image_file_path)
zero_fill = len(osp.basename(image_file_path).replace('.jpg', ''))
if zero_fill == 6:
is_posetrack18 = True
else:
is_posetrack18 = False
current_idx = int(osp.basename(image_file_path).replace('.jpg', ''))
if self.distance_whole_otherwise_segment:
farthest_distance = self.distance
prev_delta_range = range(
1,
min((current_idx + 1) if is_posetrack18 else current_idx,
farthest_distance))
next_delta_range = range(
1,
min((num_frames - current_idx) if is_posetrack18 else
(num_frames - current_idx + 1), farthest_distance))
else:
prev_delta_range = range(
1,
min(current_idx + 1 if is_posetrack18 else current_idx,
self.previous_distance))
next_delta_range = range(
1,
min((num_frames - current_idx) if is_posetrack18 else
(num_frames - current_idx + 1), self.next_distance))
prev_delta_range = list(prev_delta_range)
next_delta_range = list(next_delta_range)
# setting deltas
if len(prev_delta_range) == 0:
prev_delta = 0
margin_left = 1
else:
if self.random_aux_frame:
prev_delta = random.choice(prev_delta_range)
else:
prev_delta = prev_delta_range[-1]
margin_left = prev_delta
if len(next_delta_range) == 0:
next_delta = 0
margin_right = 1
else:
if self.random_aux_frame:
next_delta = random.choice(next_delta_range)
else:
next_delta = next_delta_range[-1]
margin_right = next_delta
prev_idx = current_idx - prev_delta
next_idx = current_idx + next_delta
prev_image_file = osp.join(osp.dirname(image_file_path),
str(prev_idx).zfill(zero_fill) + '.jpg')
next_image_file = osp.join(osp.dirname(image_file_path),
str(next_idx).zfill(zero_fill) + '.jpg')
# checking for files existence
if not osp.exists(prev_image_file):
error_msg = "Can not find image :{}".format(prev_image_file)
self.logger.error(error_msg)
raise Exception(error_msg)
if not osp.exists(next_image_file):
error_msg = "Can not find image :{}".format(next_image_file)
self.logger.error(error_msg)
raise Exception(error_msg)
data_numpy_prev = read_image(prev_image_file)
data_numpy_next = read_image(next_image_file)
if self.color_rgb:
# cv2 read_image color channel is BGR
data_numpy = cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB)
data_numpy_prev = cv2.cvtColor(data_numpy_prev, cv2.COLOR_BGR2RGB)
data_numpy_next = cv2.cvtColor(data_numpy_next, cv2.COLOR_BGR2RGB)
if data_numpy is None:
self.logger.error('=> fail to read {}'.format(image_file_path))
raise ValueError('Fail to read {}'.format(image_file_path))
if data_numpy_prev is None:
self.logger.error(
'=> PREV SUP: fail to read {}'.format(prev_image_file))
raise ValueError(
'PREV SUP: Fail to read {}'.format(prev_image_file))
if data_numpy_next is None:
self.logger.error(
'=> NEXT SUP: fail to read {}'.format(next_image_file))
raise ValueError(
'NEXT SUP: Fail to read {}'.format(next_image_file))
joints = data_item['joints_3d']
joints_vis = data_item['joints_3d_vis']
center = data_item["center"]
scale = data_item["scale"]
score = data_item['score'] if 'score' in data_item else 1
r = 0
if self.is_train:
if (np.sum(joints_vis[:, 0]) > self.num_joints_half_body
and np.random.rand() < self.prob_half_body):
c_half_body, s_half_body = half_body_transform(
joints, joints_vis, self.num_joints, self.upper_body_ids,
self.aspect_ratio, self.pixel_std)
center, scale = c_half_body, s_half_body
scale_factor = self.scale_factor
rotation_factor = self.rotation_factor
# scale = scale * np.random.uniform(1 - scale_factor[0], 1 + scale_factor[1])
if isinstance(scale_factor, list) or isinstance(
scale_factor, tuple):
scale_factor = scale_factor[0]
scale = scale * np.clip(np.random.randn() * scale_factor + 1,
1 - scale_factor, 1 + scale_factor)
r = np.clip(np.random.randn() * rotation_factor, -rotation_factor * 2, rotation_factor * 2) \
if random.random() <= 0.6 else 0
if self.flip and random.random() <= 0.5:
data_numpy = data_numpy[:, ::-1, :]
data_numpy_prev = data_numpy_prev[:, ::-1, :]
data_numpy_next = data_numpy_next[:, ::-1, :]
joints, joints_vis = fliplr_joints(joints, joints_vis,
data_numpy.shape[1],
self.flip_pairs)
center[0] = data_numpy.shape[1] - center[0] - 1
# calculate transform matrix
trans = get_affine_transform(center, scale, r, self.image_size)
input_x = cv2.warpAffine(
data_numpy,
trans, (int(self.image_size[0]), int(self.image_size[1])),
flags=cv2.INTER_LINEAR)
input_prev = cv2.warpAffine(
data_numpy_prev,
trans, (int(self.image_size[0]), int(self.image_size[1])),
flags=cv2.INTER_LINEAR)
input_next = cv2.warpAffine(
data_numpy_next,
trans, (int(self.image_size[0]), int(self.image_size[1])),
flags=cv2.INTER_LINEAR)
if self.transform:
input_x = self.transform(input_x)
input_prev = self.transform(input_prev)
input_next = self.transform(input_next)
# joint transform like image
for i in range(self.num_joints):
if joints_vis[
i, 0] > 0.0: # joints_vis num_joints x 3 (x_vis,y_vis)
joints[i, 0:2] = exec_affine_transform(joints[i, 0:2], trans)
# H W
for index, joint in enumerate(joints):
x, y, _ = joint
if x < 0 or y < 0 or x > self.image_size[0] or y > self.image_size[
1]:
joints_vis[index] = [0, 0, 0]
# target_heatmaps, target_heatmaps_weight = self._generate_target(joints, joints_vis, self.heatmap_size, self.num_joints)
target_heatmaps, target_heatmaps_weight = generate_heatmaps(
joints,
joints_vis,
self.sigma,
self.image_size,
self.heatmap_size,
self.num_joints,
use_different_joints_weight=self.use_different_joints_weight,
joints_weight=self.joints_weight)
target_heatmaps = torch.from_numpy(target_heatmaps) # H W
target_heatmaps_weight = torch.from_numpy(target_heatmaps_weight)
meta = {
'image': image_file_path,
'prev_sup_image': prev_image_file,
'next_sup_image': next_image_file,
'filename': filename,
'imgnum': img_num,
'joints': joints,
'joints_vis': joints_vis,
'center': center,
'scale': scale,
'rotation': r,
'score': score,
'margin_left': margin_left,
'margin_right': margin_right,
}
return input_x, input_prev, input_next, target_heatmaps, target_heatmaps_weight, meta
def video():
logger.info("Start")
base_video_path = "./input"
base_img_vis_save_dirs = './output/vis_img'
json_save_base_dirs = './output/json'
create_folder(json_save_base_dirs)
video_list = list_immediate_childfile_paths(base_video_path,
ext=['mp3', 'mp4'])
input_image_save_dirs = []
SAVE_JSON = True
SAVE_VIS_VIDEO = True
SAVE_VIS_IMAGE = True
SAVE_BOX_IMAGE = True
base_img_vis_box_save_dirs = './output/vis_img_box'
# 1.Split the video into images
for video_path in tqdm(video_list):
video_name = osp.basename(video_path)
temp = video_name.split(".")[0]
image_save_path = os.path.join(base_video_path, temp)
image_vis_save_path = os.path.join(base_img_vis_save_dirs, temp)
image_vis_box_save_path = os.path.join(base_img_vis_box_save_dirs,
temp)
input_image_save_dirs.append(image_save_path)
create_folder(image_save_path)
create_folder(image_vis_save_path)
create_folder(image_vis_box_save_path)
video2images(video_path, image_save_path) # jpg
# 2. Person Instance detection
logger.info("Person Instance detection in progress ...")
video_candidates = {}
for index, images_dir in tqdm(enumerate(input_image_save_dirs)):
# if index >= 1:
# continue
video_name = osp.basename(images_dir)
image_list = list_immediate_childfile_paths(images_dir, ext='jpg')
video_candidates_list = []
for image_path in tqdm(image_list):
candidate_bbox = inference_yolov3(image_path)
for bbox in candidate_bbox:
# bbox - x, y, w, h
video_candidates_list.append({
"image_path": image_path,
"bbox": bbox,
"keypoints": None
})
video_candidates[video_name] = {
"candidates_list": video_candidates_list,
"length": len(image_list)
}
logger.info("Person Instance detection finish")
# 3. Singe Person Pose Estimation
logger.info("Single person pose estimation in progress ...")
for video_name, video_info in video_candidates.items():
video_candidates_list = video_info["candidates_list"]
video_length = video_info["length"]
prev_image_id = None
for person_info in tqdm(video_candidates_list):
image_path = person_info["image_path"]
xywh_box = person_info["bbox"]
print(os.path.basename(image_path))
image_idx = int(os.path.basename(image_path).replace(".jpg", ""))
# from
prev_idx, next_id = image_idx - 1, image_idx + 1
if prev_idx < 0:
prev_idx = 0
if image_idx >= video_length - 1:
next_id = video_length - 1
prev_image_path = os.path.join(
os.path.dirname(image_path),
"{}.jpg".format(str(prev_idx).zfill(zero_fill)))
next_image_path = os.path.join(
os.path.dirname(image_path),
"{}.jpg".format(str(next_id).zfill(zero_fill)))
# current_image = read_image(image_path)
# prev_image = read_image(prev_image_path)
# next_image = read_image(next_image_path)
bbox = xywh_box
keypoints = inference_PE(image_path, prev_image_path,
next_image_path, bbox)
person_info["keypoints"] = keypoints.tolist()[0]
# posetrack points
new_coord = coco2posetrack_ord_infer(keypoints[0])
# pose
if SAVE_VIS_IMAGE:
image_save_path = os.path.join(
os.path.join(base_img_vis_save_dirs, video_name),
image_path.split("/")[-1])
if osp.exists(image_save_path):
current_image = read_image(image_save_path)
else:
current_image = read_image(image_path)
pose_img = add_poseTrack_joint_connection_to_image(
current_image,
new_coord,
sure_threshold=0.3,
flag_only_draw_sure=True)
save_image(image_save_path, pose_img)
if SAVE_BOX_IMAGE:
image_save_path = os.path.join(
os.path.join(base_img_vis_box_save_dirs, video_name),
image_path.split("/")[-1])
if osp.exists(image_save_path):
current_image = read_image(image_save_path)
else:
current_image = read_image(image_path)
xyxy_box = bbox[0], bbox[
1], bbox[0] + bbox[2], bbox[1] + bbox[3]
box_image = add_bbox_in_image(current_image, xyxy_box)
save_image(image_save_path, box_image)
if SAVE_JSON:
joints_info = {"Info": video_candidates_list}
temp = "result_" + video_name + ".json"
write_json_to_file(joints_info,
os.path.join(json_save_base_dirs, temp))
print("------->json Info save Complete!")
print("------->Visual Video Compose Start")
if SAVE_VIS_VIDEO:
image2video(os.path.join(base_img_vis_save_dirs, video_name),
video_name)
print("------->Complete!")