Example #1
0
            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])
Example #2
0
    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)
Example #3
0
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
Example #4
0
    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)