center = (int(body_part.x * image_w + 0.5), int(body_part.y * image_h + 0.5))
centers[i] = center
cv2.circle(npimg, center, 3, CocoColors[i], thickness=3, lineType=8, shift=0)
# draw line
for pair_order, pair in enumerate(CocoPairsRender):
if pair[0] not in human.body_parts.keys() or pair[1] not in human.body_parts.keys():
continue
# npimg = cv2.line(npimg, centers[pair[0]], centers[pair[1]], common.CocoColors[pair_order], 3)
cv2.line(npimg, centers[pair[0]], centers[pair[1]], CocoColors[pair_order], 3)
return npimg
weight_name = './network/weight/pose_model.pth'
model = get_model('vgg19')
model.load_state_dict(torch.load(weight_name))
model.cuda()
model.float()
model.eval()
if __name__ == "__main__":
video_capture = cv2.VideoCapture(0)
while True:
# Capture frame-by-frame
ret, oriImg = video_capture.read()
shape_dst = np.min(oriImg.shape[0:2])
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
# model
model = get_model(trunk='vgg19')
model = torch.nn.DataParallel(model).cuda()
# load pretrained
use_vgg(model)
# Fix the VGG weights first, and then the weights will be released
for i in range(20):
for param in model.module.model0[i].parameters():
param.requires_grad = False
trainable_vars = [param for param in model.parameters() if param.requires_grad]
optimizer = torch.optim.SGD(trainable_vars,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
def main(args):
sys.path.append(
args.openpose_dir) # In case calling from an external script
from lib.network.rtpose_vgg import get_model
from lib.network.rtpose_vgg import use_vgg
from lib.network import im_transform
from evaluate.coco_eval import get_outputs, handle_paf_and_heat
from lib.utils.common import Human, BodyPart, CocoPart, CocoColors, CocoPairsRender, draw_humans
from lib.utils.paf_to_pose import paf_to_pose_cpp
from lib.config import cfg, update_config
update_config(cfg, args)
model = get_model('vgg19')
model = torch.nn.DataParallel(model).cuda()
use_vgg(model)
# model.load_state_dict(torch.load(args.weight))
checkpoint = torch.load(args.weight)
epoch = checkpoint['epoch']
best_loss = checkpoint['best_loss']
state_dict = checkpoint['state_dict']
# state_dict = {key.replace("module.",""):value for key, value in state_dict.items()} # Remove "module." from vgg keys
model.load_state_dict(state_dict)
# optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(args.weight, epoch))
model.float()
model.eval()
image_folders = args.image_folders.split(',')
for i, image_folder in enumerate(image_folders):
print(
f"\nProcessing {i} of {len(image_folders)}: {' '.join(image_folder.split('/')[-4:-2])}"
)
if args.all_frames: # Split video and run inference on all frames
output_dir = os.path.join(os.path.dirname(image_folder),
'predictions', 'pose2d',
'openpose_pytorch_ft_all')
os.makedirs(output_dir, exist_ok=True)
video_path = os.path.join(
image_folder,
'scan_video.avi') # break up video and run on all frames
temp_folder = image_folder.split('/')[-3] + '_openpose'
image_folder = os.path.join(
'/tmp', f'{temp_folder}') # Overwrite image_folder
os.makedirs(image_folder, exist_ok=True)
split_video(video_path, image_folder)
else: # Just use GT-annotated frames
output_dir = os.path.join(os.path.dirname(image_folder),
'predictions', 'pose2d',
'openpose_pytorch_ft')
os.makedirs(output_dir, exist_ok=True)
img_mask = os.path.join(image_folder, '??????.png')
img_names = glob(img_mask)
for img_name in img_names:
image_file_path = img_name
oriImg = cv2.imread(image_file_path) # B,G,R order
shape_dst = np.min(oriImg.shape[0:2])
with torch.no_grad():
paf, heatmap, im_scale = get_outputs(oriImg, model, 'rtpose')
humans = paf_to_pose_cpp(heatmap, paf, cfg)
# Save joints in OpenPose format
image_h, image_w = oriImg.shape[:2]
people = []
for i, human in enumerate(humans):
keypoints = []
for j in range(18):
if j == 8:
keypoints.extend([
0, 0, 0
]) # Add extra joint (midhip) to correspond to body_25
if j not in human.body_parts.keys():
keypoints.extend([0, 0, 0])
else:
body_part = human.body_parts[j]
keypoints.extend([
body_part.x * image_w, body_part.y * image_h,
body_part.score
])
person = {"person_id": [i - 1], "pose_keypoints_2d": keypoints}
people.append(person)
people_dict = {"people": people}
_, filename = os.path.split(image_file_path)
name, _ = os.path.splitext(filename)
frame_id = int(name)
with open(
os.path.join(output_dir,
f"scan_video_{frame_id:012}_keypoints.json"),
'w') as outfile:
json.dump(people_dict, outfile)
if args.all_frames:
shutil.rmtree(image_folder) # Delete image_folder
return rtpose_model
args = parse_args()
update_config(cfg, args)
print("Loading dataset...")
# load train data
preprocess = transforms.Compose([
transforms.Normalize(),
transforms.RandomApply(transforms.HFlip(), 0.5),
transforms.RescaleRelative(scale_range=(cfg.DATASET.SCALE_MIN, cfg.DATASET.SCALE_MAX)),
transforms.Crop(cfg.DATASET.IMAGE_SIZE),
transforms.CenterPad(cfg.DATASET.IMAGE_SIZE),
])
# model
rtpose_vgg = get_model(trunk='vgg19')
# load pretrained
use_vgg(rtpose_vgg)
class rtpose_lightning(pl.LightningModule):
def __init__(self, preprocess, target_transforms, model, optimizer):
super(rtpose_lightning, self).__init__()
self.preprocess = preprocess
self.model = model
self.opt = optimizer
self.target_transforms = target_transforms
def forward(self, x):
_, saved_for_loss = self.model.forward(x)
