File size: 6,915 Bytes
b41b87f |
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 |
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import argparse
import csv
import os
import shutil
import sys
from PIL import Image
import torch
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision
import cv2
import numpy as np
import time
sys.path.append('./deep-high-resolution-net.pytorch/lib')
import models
from config import cfg
from config import update_config
from core.function import get_final_preds
from utils.transforms import get_affine_transform
import distutils.core
# os.system('python -m pip install pyyaml==5.3.1')
# dist = distutils.core.run_setup("./detectron2/setup.py")
# temp = ' '.join([f"'{x}'" for x in dist.install_requires])
# cmd = "python -m pip install {0}".format(temp)
# os.system(cmd)
# sys.path.insert(0, os.path.abspath('./detectron2'))
# import detectron2
# # from detectron2.modeling import build_model
# from detectron2 import model_zoo
# from detectron2.engine import DefaultPredictor
# from detectron2.config import get_cfg
# from detectron2.utils.visualizer import Visualizer
# from detectron2.data import MetadataCatalog, DatasetCatalog
# from detectron2.utils.visualizer import Visualizer
# from detectron2.checkpoint import DetectionCheckpointer
# from detectron2.data.datasets import register_coco_instances
# from detectron2.utils.visualizer import ColorMode
from models.detectron2.diver_detector_setup import get_diver_detector
from models.pose_estimator.pose_hrnet import get_pose_net
def box_to_center_scale(box, model_image_width, model_image_height):
"""convert a box to center,scale information required for pose transformation
Parameters
----------
box : list of tuple
list of length 2 with two tuples of floats representing
bottom left and top right corner of a box
model_image_width : int
model_image_height : int
Returns
-------
(numpy array, numpy array)
Two numpy arrays, coordinates for the center of the box and the scale of the box
"""
center = np.zeros((2), dtype=np.float32)
bottom_left_corner = (box[0].data.cpu().item(), box[1].data.cpu().item())
top_right_corner = (box[2].data.cpu().item(), box[3].data.cpu().item())
box_width = top_right_corner[0]-bottom_left_corner[0]
box_height = top_right_corner[1]-bottom_left_corner[1]
bottom_left_x = bottom_left_corner[0]
bottom_left_y = bottom_left_corner[1]
center[0] = bottom_left_x + box_width * 0.5
center[1] = bottom_left_y + box_height * 0.5
aspect_ratio = model_image_width * 1.0 / model_image_height
pixel_std = 200
if box_width > aspect_ratio * box_height:
box_height = box_width * 1.0 / aspect_ratio
elif box_width < aspect_ratio * box_height:
box_width = box_height * aspect_ratio
scale = np.array(
[box_width * 1.0 / pixel_std, box_height * 1.0 / pixel_std],
dtype=np.float32)
if center[0] != -1:
scale = scale * 1.25
return center, scale
def parse_args():
parser = argparse.ArgumentParser(description='Train keypoints network')
# general
parser.add_argument('--cfg', type=str, default='./deep-high-resolution-net.pytorch/experiments/mpii/hrnet/w32_256x256_adam_lr1e-3.yaml')
parser.add_argument('opts',
help='Modify config options using the command-line',
default=None,
nargs=argparse.REMAINDER)
args = parser.parse_args()
# args expected by supporting codebase
args.modelDir = ''
args.logDir = ''
args.dataDir = ''
args.prevModelDir = ''
return args
def get_pose_estimation_prediction(pose_model, image, center, scale):
rotation = 0
trans = get_affine_transform(center, scale, rotation, cfg.MODEL.IMAGE_SIZE)
# trans = cv2.getAffineTransform(srcTri, dstTri)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
model_input = cv2.warpAffine(
image,
trans,
(256, 256),
flags=cv2.INTER_LINEAR)
# pose estimation inference
model_input = transform(model_input).unsqueeze(0)
# switch to evaluate mode
pose_model.eval()
with torch.no_grad():
# compute output heatmap
output = pose_model(model_input)
preds, _ = get_final_preds(
cfg,
output.clone().cpu().numpy(),
np.asarray([center]),
np.asarray([scale]))
return preds
def get_pose_model():
CTX = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
cudnn.benchmark = cfg.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
args = parse_args()
update_config(cfg, args)
pose_model = get_pose_net(cfg, is_train=False)
if cfg.TEST.MODEL_FILE:
print('=> loading model from {}'.format(cfg.TEST.MODEL_FILE))
pose_model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE), strict=False)
else:
print('expected model defined in config at TEST.MODEL_FILE')
pose_model = torch.nn.DataParallel(pose_model, device_ids=cfg.GPUS)
pose_model.to(CTX)
pose_model.eval()
return pose_model
def get_pose_estimation(filepath, image_bgr=None, diver_detector=None, pose_model=None):
if image_bgr is None:
image_bgr = cv2.imread(filepath)
if image_bgr is None:
print("ERROR: image {} does not exist".format(filepath))
return None
if diver_detector is None:
diver_detector = get_diver_detector()
if pose_model is None:
pose_model = get_pose_model()
image = image_bgr[:, :, [2, 1, 0]]
outputs = diver_detector(image_bgr)
scores = outputs['instances'].scores
pred_boxes = []
if len(scores) > 0:
pred_boxes = outputs['instances'].pred_boxes
if len(pred_boxes) >= 1:
for box in pred_boxes:
center, scale = box_to_center_scale(box, cfg.MODEL.IMAGE_SIZE[0], cfg.MODEL.IMAGE_SIZE[1])
image_pose = image.copy() if cfg.DATASET.COLOR_RGB else image_bgr.copy()
box = box.detach().cpu().numpy()
return box, get_pose_estimation_prediction(pose_model, image_pose, center, scale)
# print("pose_preds", pose_preds)
# draw_bbox(box,image_bgr)
# if len(pose_preds)>=1:
# print('drawing preds')
# for kpt in pose_preds:
# draw_pose(kpt,image_bgr) # draw the poses
# break # only want to use the box with the highest confidence score
return None, None
|