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)
for epoch in range(5):
# train for one epoch
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
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)
return saved_for_loss