def extract_heatmap_info(heatmap_avg, param_thre1=0.1):
all_peaks = []
peak_counter = 0
for part in range(18):
map_ori = heatmap_avg[:, :, part]
map_gau = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map_gau.shape)
map_left[1:, :] = map_gau[:-1, :]
map_right = np.zeros(map_gau.shape)
map_right[:-1, :] = map_gau[1:, :]
map_up = np.zeros(map_gau.shape)
map_up[:, 1:] = map_gau[:, :-1]
map_down = np.zeros(map_gau.shape)
map_down[:, :-1] = map_gau[:, 1:]
peaks_binary = np.logical_and.reduce(
(map_gau >= map_left, map_gau >= map_right, map_gau >= map_up,
map_gau >= map_down, map_gau > param_thre1))
peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse
peaks = list(peaks)
peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
ids = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [peaks_with_score[i] + (ids[i],) for i in range(len(ids))]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
return all_peaks
def validate_epoch(model, loader, dtype):
"""
validation for MultiLabelMarginLoss using f2 score
`model` is a trained subclass of torch.nn.Module
`loader` is a torch.dataset.DataLoader for validation data
`dtype` data type for variables
eg torch.FloatTensor (cpu) or torch.cuda.FloatTensor (gpu)
"""
n_samples = len(loader.sampler)
x, y = loader.dataset[0]
y_array = np.zeros((n_samples))
y_pred_array = np.zeros((n_samples))
bs = loader.batch_size
## Put the model in test mode
model.eval()
for i, (x, y) in enumerate(loader):
x_var = Variable(x.type(dtype), volatile=True)
y_var = Variable(y.type(dtype).long())
y_var = y_var.view(y_var.data.shape[0])
scores = model(x_var)
y_pred = scores.data.max(1)[1].cpu().numpy()
y_array[i * bs:(i + 1) * bs] = y_var.data.cpu().numpy()
y_pred_array[i * bs:(i + 1) * bs] = y_pred
return (y_array == y_pred_array).sum() / float(y_pred_array.shape[0])
def extract_heatmap_info(heatmap_avg, param_thre1=0.1):
all_peaks = []
peak_counter = 0
for part in range(18):
map_ori = heatmap_avg[:, :, part]
map_gau = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map_gau.shape)
map_left[1:, :] = map_gau[:-1, :]
map_right = np.zeros(map_gau.shape)
map_right[:-1, :] = map_gau[1:, :]
map_up = np.zeros(map_gau.shape)
map_up[:, 1:] = map_gau[:, :-1]
map_down = np.zeros(map_gau.shape)
map_down[:, :-1] = map_gau[:, 1:]
peaks_binary = np.logical_and.reduce(
(map_gau >= map_left, map_gau >= map_right, map_gau >= map_up,
map_gau >= map_down, map_gau > param_thre1))
peaks = zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0]) # note reverse
peaks = list(peaks)
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
ids = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (ids[i], ) for i in range(len(ids))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
return all_peaks
def get_scores(model, loader, dtype):
"""
function to get scores and correct y values for a dataset
"""
n_samples = len(loader.sampler)
x, y = loader.dataset[0]
y_array = np.zeros((n_samples, y.size()[0]))
sig_scores_array = np.zeros(y_array.shape)
bs = loader.batch_size
## Put the model in test mode
model.eval()
for i, (x, y) in enumerate(loader):
x_var = Variable(x.type(dtype), volatile=True)
scores = model(x_var)
if dtype == 'torch.cuda.FloatTensor' or dtype is torch.cuda.FloatTensor:
sig_scores = torch.sigmoid(scores).data.cpu().numpy()
y_array[i * bs:(i + 1) * bs, :] = y.cpu().numpy()
sig_scores_array[i * bs:(i + 1) * bs, :] = sig_scores
elif dtype == 'torch.FloatTensor' or dtype is torch.FloatTensor:
sig_scores = torch.sigmoid(scores).data.numpy()
y_array[i * bs:(i + 1) * bs, :] = y.numpy()
sig_scores_array[i * bs:(i + 1) * bs, :] = sig_scores
else:
print('dtype not supported')
return sig_scores_array, y_array
def detect(self, minibatch, humans, scales):
ms = minibatch.shape
tic = time.time()
imageToTest = Variable(T.transpose(
T.transpose((torch.from_numpy(minibatch).float() / 256.0) - 0.5, 2,
3), 1, 2),
volatile=True).cuda()
output1, output2 = self.model(imageToTest)
heatmap_avg = nn.UpsamplingBilinear2d((ms[1], ms[2])).cuda()(output2)
paf_avg = nn.UpsamplingBilinear2d((ms[1], ms[2])).cuda()(output1)
heatmap_avg = T.transpose(heatmap_avg, 1, 2)
heatmap_avg = T.transpose(heatmap_avg, 2, 3)
heatmap_avg = heatmap_avg.cpu().data.numpy()
#print 'heatmap shape : ' , heatmap_avg.shape
paf_avg = paf_avg.cpu().data.numpy()
heatmap_avg = heatmap_avg[:, :, :, :-1]
map = gaussian_filter1d(heatmap_avg, sigma=3, axis=2)
map = gaussian_filter1d(map, sigma=3, axis=1)
map_left = np.zeros(map.shape)
map_left[:, 1:, :, :] = map[:, :-1, :, :]
map_right = np.zeros(map.shape)
map_right[:, :-1, :, :] = map[:, 1:, :, :]
map_up = np.zeros(map.shape)
map_up[:, :, 1:, :] = map[:, :, :-1, :]
map_down = np.zeros(map.shape)
map_down[:, :, :-1, :] = map[:, :, 1:, :]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > self.param_['thre1']))
peaks = np.nonzero(peaks_binary)
for i in range(peaks[0].shape[0]):
human_idx = int(peaks[0][i])
x = int(peaks[1][i] / scales[human_idx])
y = int(peaks[2][i] / scales[human_idx])
#x = peaks[1][i]
#y = peaks[2][i]
joint = peaks[3][i]
humans[human_idx].joints[joint * 2] = x
humans[human_idx].joints[joint * 2 + 1] = y
return humans
def validate(all_model, gender_model, smile_model, loader, dtype):
"""
validation for MultiLabelMarginLoss using f2 score
`model` is a trained subclass of torch.nn.Module
`loader` is a torch.dataset.DataLoader for validation data
`dtype` data type for variables
eg torch.FloatTensor (cpu) or torch.cuda.FloatTensor (gpu)
"""
n_samples = len(loader.sampler)
#x, y,z = loader.dataset[0]
y_array = np.zeros((n_samples))
y_pred_array = np.zeros((n_samples))
z_array = np.zeros((n_samples))
z_pred_array = np.zeros((n_samples))
bs = loader.batch_size
## Put the model in test mode
all_model.eval()
gender_model.eval()
smile_model.eval()
for i, (x, y, z) in enumerate(loader):
x_var = Variable(x.type(dtype), volatile=True)
y_var = Variable(y.type(dtype).long())
y_var = y_var.view(y_var.data.shape[0])
z_var = Variable(z.type(dtype).long())
z_var = z_var.view(z_var.data.shape[0])
noise_x = all_model(x_var)
scores_gender = gender_model(noise_x)
scores_smile = smile_model(noise_x)
y_pred = scores_gender.data.max(1)[1].cpu().numpy()
z_pred = scores_smile.data.max(1)[1].cpu().numpy()
y_array[i * bs:(i + 1) * bs] = y_var.data.cpu().numpy()
y_pred_array[i * bs:(i + 1) * bs] = y_pred
z_array[i * bs:(i + 1) * bs] = z_var.data.cpu().numpy()
z_pred_array[i * bs:(i + 1) * bs] = z_pred
acc_gender = (y_array == y_pred_array).sum() / float(y_pred_array.shape[0])
acc_smile = (z_array == z_pred_array).sum() / float(z_pred_array.shape[0])
y_bool = (y_array == y_pred_array)
z_bool = (z_array == z_pred_array)
acc_all = np.multiply(y_bool, z_bool).sum() / float(y_bool.shape[0])
return acc_all, acc_gender, acc_smile
def get_heatmap_pose(new_key_points):
#print ('Get Heatmap...')
heatmap_pose = np.zeros([256, 256])
for item in new_key_points:
if item != []:
heatmap_pose = fill_in(heatmap_pose, item, 4)
return heatmap_pose
def triple_train_val_balance_dataloaders(datasets,
batch_size,
num_workers,
shuffle=True,
split=0.9,
use_fraction_of_data=1.):
"""
generate three training and three validation dataloaders
to train triple resnet
"""
n_samples = len(datasets[0])
## set up train val split
inds = np.arange(n_samples)
train_inds, val_inds = train_test_split(inds,
test_size=1 - split,
train_size=split)
## logical indexing to use with BalanceSampler
log_train_inds = np.zeros(n_samples)
log_train_inds[train_inds] = 1
log_val_inds = np.zeros(n_samples)
log_val_inds[val_inds] = 1
## reduce the size of your dataset (use for testing only)
if use_fraction_of_data < 1:
train_idx = int(np.ceil(use_fraction_of_data * split * n_samples))
val_idx = int(np.ceil(use_fraction_of_data * (1 - split) * n_samples))
log_train_inds[train_idx:] = 0
log_val_inds[val_idx:] = 0
train_loaders = []
val_loaders = []
for dset in datasets:
train_loaders.append(
BalanceDataLoader(dset,
sampler=BalanceSampler(dset, log_train_inds),
batch_size=batch_size,
shuffle=shuffle,
num_workers=num_workers))
val_loaders.append(
BalanceDataLoader(dset,
sampler=BalanceSampler(dset, log_val_inds),
batch_size=batch_size,
shuffle=shuffle,
num_workers=num_workers))
return train_loaders, val_loaders
def get_18_heatmaps(new_key_points):
heatmaps_18 = []
for item in new_key_points:
heatmap = np.zeros([256, 256])
if (item != []):
heatmap = fill_in(heatmap, item, 4)
heatmaps_18.append(heatmap)
return heatmaps_18
def get_scores(models, loaders, dtype, n_classes=17):
s = np.zeros((3, len(loaders[0].sampler), n_classes))
y_array = np.zeros(s.shape[1:])
bs = loaders[0].batch_size
## Put the model in test mode
models[0].eval()
models[1].eval()
models[2].eval()
## loop over the three models
for i, (model, loader) in enumerate(zip(models, loaders)):
for j, (x, y) in enumerate(loader):
x_var = Variable(x.type(dtype), volatile=True)
if i == 0:
y_array[j * bs:(j + 1) * bs, :] = y.numpy()
score_var = model(x_var)
## store each set of scores
s[i, j * bs:(j + 1) * bs, :] = score_var.data.numpy()
return s, y_array
def test_triple_resnet(models, loaders, mlb, dtype, weights=(1,1,1),
out_file_name="", sigmoid_threshold=None, n_classes=17):
"""
run 3 models on test data and generate a csv file for submission to kaggle
"""
if sigmoid_threshold is None:
sigmoid_threshold = 0.5
## store scores for all three models
s = np.zeros((3, len(loaders[0].sampler), n_classes))
file_names = []
bs = loaders[0].batch_size
## Put the model in test mode
models[0].eval()
models[1].eval()
models[2].eval()
## loop over the three models
for i, (model, loader) in enumerate(zip(models, loaders)):
for j, (x, file_name) in enumerate(loader):
x_var = Variable(x.type(dtype), volatile=True)
if i == 0:
file_names += list(file_name)
score_var = model(x_var)
## store each set of scores
if dtype is torch.FloatTensor:
s[i,j*bs:(j+1)*bs,:] = score_var.data.numpy()
else:
s[i,j*bs:(j+1)*bs,:] = score_var.data.cpu().numpy()
## weighted average of scores from 3 models
scores = (weights[0]*s[0,:,:] + weights[1]*s[1,:,:] + weights[2]*s[2,:,:]) / sum(weights)
## https://discuss.pytorch.org/t/calculating-accuracy-for-a-multi-label-classification-problem/2303
y_pred = torch.sigmoid(torch.from_numpy(scores)).numpy() > sigmoid_threshold
## generate labels from MultiLabelBinarizer
labels = mlb.inverse_transform(y_pred)
## write output file
if out_file_name:
with open(out_file_name, 'w', newline='') as csvfile:
fieldnames = ['image_name', 'tags']
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
for i, labs in enumerate(labels):
str1 = ""
for lab in labs:
## check if there is a label match
str1 += str(lab) + " "
writer.writerow({'image_name': file_names[i], 'tags': str1})
return y_pred
def extract_paf_info(img_raw, paf_avg, all_peaks, param_thre2=0.05, param_thre3=0.5):
connection_all = []
special_k = []
mid_num = 10
for k in range(len(map_ids)):
score_mid = paf_avg[:, :, [x - 19 for x in map_ids[k]]]
candA = all_peaks[limb_seq[k][0] - 1]
candB = all_peaks[limb_seq[k][1] - 1]
nA = len(candA)
nB = len(candB)
if nA != 0 and nB != 0:
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num),
np.linspace(candA[i][1], candB[j][1], num=mid_num))
startend = list(startend)
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0]
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1]
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
score_with_dist_prior = sum(score_midpts) / len(score_midpts)
score_with_dist_prior += min(0.5 * img_raw.shape[0] / norm - 1, 0)
criterion1 = len(np.nonzero(score_midpts > param_thre2)[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append(
[i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if i not in connection[:, 3] and j not in connection[:, 4]:
connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
if len(connection) >= min(nA, nB):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
return special_k, connection_all
def get_pred(model, loader, dtype):
y_pred_array = np.zeros(len(loader.sampler))
bs = loader.batch_size
model.eval()
for i, (x, y) in enumerate(loader):
x_var = Variable(x.type(dtype))
y_var = Variable(y.type(dtype))
y_pred = model(x_var)
y_pred_array[i*bs:(i+1)*bs] = y_pred.data.numpy()[:,0]
return y_pred
def callback_objs(self, msg):
#if self.flag == 0 : return None
tic = time.time()
humans = []
human_imgs = []
scales = []
pubmsg = msg
for item in msg.objects:
if item.class_string == 'person':
item.joints = [-1] * 36
humans.append(item)
img = self.cvbridge.imgmsg_to_cv2(item.cropped, 'bgr8')
imgg = np.zeros(
(self.model_['boxsize'], self.model_['boxsize'], 3))
if img.shape[0] >= img.shape[1]:
scale = float(self.model_['boxsize']) / img.shape[0]
else:
scale = float(self.model_['boxsize']) / img.shape[1]
scales.append(scale)
img2 = cv2.resize(img, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imgg[:img2.shape[0], :img2.shape[1]] = img2.copy()
human_imgs.append(imgg)
if len(humans) == 0:
self.pose_pub.publish(msg)
return None
minibatch = np.stack(human_imgs, axis=0)
print 'minibatch_shape : ', minibatch.shape
num_run = minibatch.shape[0] // 20
humans_updated = []
for i in range(num_run + 1):
s = i * 20
e = min(minibatch.shape[0], (i + 1) * 20)
output = self.detect(minibatch[s:e], humans[s:e], scales[s:e])
humans_updated += output
print 'elapsed time : ', time.time() - tic
tic = time.time()
humans_updated = self.post_processing(humans_updated)
pubmsg.objects = humans_updated
self.pose_pub.publish(pubmsg)
toc = time.time()
print 'elapsed time : ', toc - tic
print ''
self.last_detect = time.time()
self.flag = 0
def __init__(self,
csv_path,
img_path,
img_ext='.jpg',
transform=transforms.Compose([
transforms.Resize(size=(64, 64)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5))
]),
limit_load=None):
df = pd.read_csv(csv_path, nrows=limit_load)
ids_missing_mask = []
for i, row in df.iterrows():
fpath = img_path / (str(int(row['id'])) + img_ext)
ids_missing_mask.append(fpath.exists())
assert all(ids_missing_mask), \
"Some images referenced in the CSV file where not found: {}".format(
df['id'][[not i for i in ids_missing_mask]])
df = df.set_index('id').sort_index()
self.df = df
self.img_path = img_path
self.img_ext = img_ext
self.transform = transform
self.ids = self.df.index
self.lb = LabelBinarizer()
if 'orientation' in self.df.columns.tolist():
self.labels = self.df['orientation'].values - 1
self.labels_binarized = self.lb.fit_transform(
self.df['orientation']).astype(np.float32)
else:
self.labels = np.zeros(self.df.shape[0])
self.labels_binarized = np.zeros(self.df.shape)
def get_triple_resnet_val_scores(models, loaders, dtype,
weights=(1,1,1), n_classes=17):
"""
given three models and training dataloaders,
return sigmoid scores and correct labels
"""
## store scores for all three models
s = np.zeros((3, len(loaders[0].sampler), n_classes))
y_array = np.zeros(s.shape[1:])
bs = loaders[0].batch_size
## Put the model in test mode
models[0].eval()
models[1].eval()
models[2].eval()
## loop over the three models
for i, (model, loader) in enumerate(zip(models, loaders)):
for j, (x, y) in enumerate(loader):
x_var = Variable(x.type(dtype), volatile=True)
if i == 0:
if dtype is torch.cuda.FloatTensor:
y_array[j*bs:(j+1)*bs,:] = y.cpu().numpy()
else:
y_array[j*bs:(j+1)*bs,:] = y.numpy()
score_var = model(x_var)
## store each set of scores
if dtype is torch.cuda.FloatTensor:
s[i,j*bs:(j+1)*bs,:] = score_var.data.cpu().numpy()
elif dtype is torch.FloatTensor:
s[i,j*bs:(j+1)*bs,:] = score_var.data.numpy()
## weighted average of scores from 3 models
scores = (weights[0]*s[0,:,:] + weights[1]*s[1,:,:] + weights[2]*s[2,:,:]) / sum(weights)
sig_scores = torch.sigmoid(torch.from_numpy(scores)).numpy()
return sig_scores, y_array
def test_model(model, loader, mlb, dtype, out_file_name="",
sigmoid_threshold=None, n_classes=17):
"""
run the model on test data and generate a csv file for submission to kaggle
"""
if sigmoid_threshold is None:
sigmoid_threshold = 0.5
y_pred_array = np.zeros((len(loader.sampler), n_classes))
file_names = []
bs = loader.batch_size
## Put the model in test mode
model.eval()
for i, (x, file_name) in enumerate(loader):
x_var = Variable(x.type(dtype), volatile=True)
file_names += list(file_name)
scores = model(x_var)
if dtype is torch.FloatTensor:
y_pred = torch.sigmoid(scores).data.numpy() > sigmoid_threshold
else:
y_pred = torch.sigmoid(scores).data.cpu().numpy() > sigmoid_threshold
y_pred_array[i*bs:(i+1)*bs,:] = y_pred
## generate labels from MultiLabelBinarizer
labels = mlb.inverse_transform(y_pred_array)
## write output file
if out_file_name:
with open(out_file_name, 'w', newline='') as csvfile:
fieldnames = ['image_name', 'tags']
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
for i, labs in enumerate(labels):
str1 = ""
for lab in labs:
## check if there is a label match
str1 += str(lab) + " "
writer.writerow({'image_name': file_names[i], 'tags': str1})
return y_pred_array
paf_avg[m] = paf[0].data
heatmap_avg = T.transpose(
T.transpose(T.squeeze(T.mean(heatmap_avg, 0)), 0, 1), 1, 2).cuda()
heatmap_avg = heatmap_avg.cpu().numpy()
paf_avg = paf_avg.cpu().numpy()
all_peaks = []
peak_counter = 0
#maps =
for part in range(18):
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > param_['thre1']))
# peaks_binary = T.eq(
# peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
peaks = zip(np.nonzero(peaks_binary)[1],
2).cuda()
heatmap_avg = heatmap_avg.cpu().numpy()
paf_avg = paf_avg.cpu().numpy()
toc = time.time()
print 'time is %.5f' % (toc - tic)
tic = time.time()
all_peaks = []
peak_counter = 0
#maps =
for part in range(18):
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > param_['thre1']))
# peaks_binary = T.eq(
# peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
peaks = zip(np.nonzero(peaks_binary)[1],
def get_pose(image, gpu):
T.set_num_threads(T.get_num_threads())
weight_name = './model/pose_model.pth'
(y, x, _) = image.shape
blocks = {}
# find connection in the specified sequence, center 29 is in the position 15
limbSeq = [[2,3], [2,6], [3,4], [4,5], [6,7], [7,8], [2,9], [9,10], \
[10,11], [2,12], [12,13], [13,14], [2,1], [1,15], [15,17], \
[1,16], [16,18], [3,17], [6,18]]
# the middle joints heatmap correpondence
mapIdx = [[31,32], [39,40], [33,34], [35,36], [41,42], [43,44], [19,20], [21,22], \
[23,24], [25,26], [27,28], [29,30], [47,48], [49,50], [53,54], [51,52], \
[55,56], [37,38], [45,46]]
# the joints between corresponding limbs
joints = [
[12, 0],
[12, 1],
[12, 6],
[12, 9],
[12, 13],
[12, 15], #connected to the Nose-Neck
[1, 4],
[4, 5], #connect LShoulder-LElbow-LWrist
[0, 2],
[2, 3], #connect RShoulder-RElbow-RWrist
[6, 7],
[7, 8], #connect RHip-RKnee-RAnkle
[9, 10],
[10, 11], #connect LHip-LKnee-LAnkle
[13, 14], #connect REye-REar
[15, 16] #connect LEye-LEar
]
# visualize
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
block0 = [{
'conv1_1': [3, 64, 3, 1, 1]
}, {
'conv1_2': [64, 64, 3, 1, 1]
}, {
'pool1_stage1': [2, 2, 0]
}, {
'conv2_1': [64, 128, 3, 1, 1]
}, {
'conv2_2': [128, 128, 3, 1, 1]
}, {
'pool2_stage1': [2, 2, 0]
}, {
'conv3_1': [128, 256, 3, 1, 1]
}, {
'conv3_2': [256, 256, 3, 1, 1]
}, {
'conv3_3': [256, 256, 3, 1, 1]
}, {
'conv3_4': [256, 256, 3, 1, 1]
}, {
'pool3_stage1': [2, 2, 0]
}, {
'conv4_1': [256, 512, 3, 1, 1]
}, {
'conv4_2': [512, 512, 3, 1, 1]
}, {
'conv4_3_CPM': [512, 256, 3, 1, 1]
}, {
'conv4_4_CPM': [256, 128, 3, 1, 1]
}]
blocks['block1_1'] = [{
'conv5_1_CPM_L1': [128, 128, 3, 1, 1]
}, {
'conv5_2_CPM_L1': [128, 128, 3, 1, 1]
}, {
'conv5_3_CPM_L1': [128, 128, 3, 1, 1]
}, {
'conv5_4_CPM_L1': [128, 512, 1, 1, 0]
}, {
'conv5_5_CPM_L1': [512, 38, 1, 1, 0]
}]
blocks['block1_2'] = [{
'conv5_1_CPM_L2': [128, 128, 3, 1, 1]
}, {
'conv5_2_CPM_L2': [128, 128, 3, 1, 1]
}, {
'conv5_3_CPM_L2': [128, 128, 3, 1, 1]
}, {
'conv5_4_CPM_L2': [128, 512, 1, 1, 0]
}, {
'conv5_5_CPM_L2': [512, 19, 1, 1, 0]
}]
for i in range(2, 7):
blocks['block%d_1' % i] = [{
'Mconv1_stage%d_L1' % i: [185, 128, 7, 1, 3]
}, {
'Mconv2_stage%d_L1' % i: [128, 128, 7, 1, 3]
}, {
'Mconv3_stage%d_L1' % i: [128, 128, 7, 1, 3]
}, {
'Mconv4_stage%d_L1' % i: [128, 128, 7, 1, 3]
}, {
'Mconv5_stage%d_L1' % i: [128, 128, 7, 1, 3]
}, {
'Mconv6_stage%d_L1' % i: [128, 128, 1, 1, 0]
}, {
'Mconv7_stage%d_L1' % i: [128, 38, 1, 1, 0]
}]
blocks['block%d_2' % i] = [{
'Mconv1_stage%d_L2' % i: [185, 128, 7, 1, 3]
}, {
'Mconv2_stage%d_L2' % i: [128, 128, 7, 1, 3]
}, {
'Mconv3_stage%d_L2' % i: [128, 128, 7, 1, 3]
}, {
'Mconv4_stage%d_L2' % i: [128, 128, 7, 1, 3]
}, {
'Mconv5_stage%d_L2' % i: [128, 128, 7, 1, 3]
}, {
'Mconv6_stage%d_L2' % i: [128, 128, 1, 1, 0]
}, {
'Mconv7_stage%d_L2' % i: [128, 19, 1, 1, 0]
}]
def make_layers(cfg_dict):
layers = []
for i in range(len(cfg_dict) - 1):
one_ = cfg_dict[i]
for k, v in one_.items():
if 'pool' in k:
layers += [
nn.MaxPool2d(kernel_size=v[0],
stride=v[1],
padding=v[2])
]
else:
conv2d = nn.Conv2d(in_channels=v[0],
out_channels=v[1],
kernel_size=v[2],
stride=v[3],
padding=v[4])
layers += [conv2d, nn.ReLU(inplace=True)]
one_ = cfg_dict[-1].keys()
k = list(one_)[0]
v = cfg_dict[-1][k]
conv2d = nn.Conv2d(in_channels=v[0],
out_channels=v[1],
kernel_size=v[2],
stride=v[3],
padding=v[4])
layers += [conv2d]
return nn.Sequential(*layers)
layers = []
for i in range(len(block0)):
one_ = block0[i]
#for k,v in one_.iteritems():
for k, v in one_.items():
if 'pool' in k:
layers += [
nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2])
]
else:
conv2d = nn.Conv2d(in_channels=v[0],
out_channels=v[1],
kernel_size=v[2],
stride=v[3],
padding=v[4])
layers += [conv2d, nn.ReLU(inplace=True)]
models = {}
models['block0'] = nn.Sequential(*layers)
#for k,v in blocks.iteritems():
for k, v in blocks.items():
models[k] = make_layers(v)
model = pose_model(models)
model.load_state_dict(torch.load(weight_name))
model.cuda(gpu)
model.float()
model.eval()
param_, model_ = config_reader()
#torch.nn.functional.pad(img pad, mode='constant', value=model_['padValue'])
tic = time.time()
oriImg = image
imageToTest = Variable(T.transpose(
T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(), 0), 2, 3), 1,
2),
volatile=True).cuda(gpu)
multiplier = [
x * model_['boxsize'] / oriImg.shape[0] for x in param_['scale_search']
]
heatmap_avg = torch.zeros(
(len(multiplier), 19, oriImg.shape[0], oriImg.shape[1])).cuda(gpu)
paf_avg = torch.zeros(
(len(multiplier), 38, oriImg.shape[0], oriImg.shape[1])).cuda(gpu)
#print(heatmap_avg.size())
toc = time.time()
#print('time to load model is %.5f'%(toc-tic))
tic = time.time()
for m in range(len(multiplier)):
scale = multiplier[m]
h = int(oriImg.shape[0] * scale)
w = int(oriImg.shape[1] * scale)
pad_h = 0 if (h % model_['stride']
== 0) else model_['stride'] - (h % model_['stride'])
pad_w = 0 if (w % model_['stride']
== 0) else model_['stride'] - (w % model_['stride'])
new_h = h + pad_h
new_w = w + pad_w
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_['stride'], model_['padValue'])
imageToTest_padded = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
feed = Variable(T.from_numpy(imageToTest_padded)).cuda(gpu)
output1, output2 = model(feed)
heatmap = nn.UpsamplingBilinear2d(
(oriImg.shape[0], oriImg.shape[1])).cuda(gpu)(output2)
paf = nn.UpsamplingBilinear2d(
(oriImg.shape[0], oriImg.shape[1])).cuda(gpu)(output1)
heatmap_avg[m] = heatmap[0].data
paf_avg[m] = paf[0].data
toc = time.time()
#print('time to forward pass is %.5f'%(toc-tic))
tic = time.time()
heatmap_avg = T.transpose(
T.transpose(T.squeeze(T.mean(heatmap_avg, 0)), 0, 1), 1, 2)
paf_avg = T.transpose(T.transpose(T.squeeze(T.mean(paf_avg, 0)), 0, 1), 1,
2)
heatmap_avg = heatmap_avg.cpu().numpy()
paf_avg = paf_avg.cpu().numpy()
toc = time.time()
#print('time to take averages is %.5f'%(toc-tic))
tic = time.time()
all_peaks = []
peak_counter = 0
#maps =
for part in range(18):
map_ori = heatmap_avg[:, :, part]
map_ak = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map_ak.shape)
map_left[1:, :] = map_ak[:-1, :]
map_right = np.zeros(map_ak.shape)
map_right[:-1, :] = map_ak[1:, :]
map_up = np.zeros(map_ak.shape)
map_up[:, 1:] = map_ak[:, :-1]
map_down = np.zeros(map_ak.shape)
map_down[:, :-1] = map_ak[:, 1:]
peaks_binary = np.logical_and.reduce(
(map_ak >= map_left, map_ak >= map_right, map_ak >= map_up,
map_ak >= map_down, map_ak > param_['thre1']))
# peaks_binary = T.eq(
# peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
peaks = list(
zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0])) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
some_id = range(peak_counter, peak_counter + len(list(peaks)))
peaks_with_score_and_id = [
peaks_with_score[i] + (some_id[i], ) for i in range(len(some_id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
if norm != 0:
vec = np.divide(vec, norm)
startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num)))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
vec_y, vec[1])
if norm != 0:
score_with_dist_prior = sum(score_midpts) / len(
score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
else:
score_with_dist_prior = sum(score_midpts) / len(
score_midpts) + min(0.5 * oriImg.shape[0] - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > param_['thre2'])
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] + candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=lambda x: x[2],
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3] and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][
indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if (subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int),
2] + connection_all[k][i][2]
if (subset[j][indexA] != partAs[i]):
subset[j][indexA] = partAs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partAs[i].astype(int),
2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
#print("found = 2")
membership = ((subset[j1] >= 0).astype(int) +
(subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[
partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(
candidate[connection_all[k][i, :2].astype(int),
2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
#canvas = cv2.imread(test_image) # B,G,R order
# for i in range(18):
# for j in range(len(all_peaks[i])):
# cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
results = []
stickwidth = 5
for n in range(len(subset)):
limbs = {}
angles = []
position = []
for i in range(17):
index = subset[n][np.array(limbSeq[i]) - 1]
if -1 in index:
continue
# cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
limb = ((X[0] - X[1], Y[0] - Y[1]))
limbs[tuple(limbSeq[i])] = limb
#if i == 12:
#position.append(np.mean(X)/x)
#position.append(np.mean(Y)/y)
# mX = np.mean(X)
# mY = np.mean(Y)
# length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
# angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
# polygon = cv2.ellipse2Poly((int(mY),int(mX)), (int(length/2), stickwidth), int(angle), 0, 360, 1)
# cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
# canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
results.append(extract.get_joints(limbs, joints, limbSeq))
#Parallel(n_jobs=1)(delayed(handle_one)(i) for i in range(18))
toc = time.time()
#torch.cuda.empty_cache()
#print('time to feature extract is %.5f'%(toc-tic))
return (results)
def extract_paf_info(img_raw,
paf_avg,
all_peaks,
param_thre2=0.05,
param_thre3=0.5):
connection_all = []
special_k = []
mid_num = 10
for k in range(len(map_ids)):
score_mid = paf_avg[:, :, [x - 19 for x in map_ids[k]]]
candA = all_peaks[limb_seq[k][0] - 1]
candB = all_peaks[limb_seq[k][1] - 1]
nA = len(candA)
nB = len(candB)
if nA != 0 and nB != 0:
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
vec = np.divide(vec, norm)
startend = zip(
np.linspace(candA[i][0], candB[j][0], num=mid_num),
np.linspace(candA[i][1], candB[j][1], num=mid_num))
startend = list(startend)
vec_x = np.array([
score_mid[int(round(startend[I][1])),
int(round(startend[I][0])), 0]
for I in range(len(startend))
])
vec_y = np.array([
score_mid[int(round(startend[I][1])),
int(round(startend[I][0])), 1]
for I in range(len(startend))
])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
vec_y, vec[1])
score_with_dist_prior = sum(score_midpts) / len(
score_midpts)
score_with_dist_prior += min(
0.5 * img_raw.shape[0] / norm - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > param_thre2)
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] + candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=lambda x: x[2],
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if i not in connection[:, 3] and j not in connection[:, 4]:
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if len(connection) >= min(nA, nB):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
return special_k, connection_all
print('warming up')
writer_open_pose = skvideo.io.FFmpegWriter("data/output.mp4",
inputdict={'-r': str(20)},
outputdict={'-r': str(20)})
writer_heat_map = skvideo.io.FFmpegWriter("data/heat_map.mp4",
inputdict={'-r': str(20)},
outputdict={'-r': str(20)})
video_capture = cv2.VideoCapture(
"/local_home/project/Tracking-with-darkflow/data/AVG-TownCentre.mp4")
# fourcc = cv2.VideoWriter_fourcc(*'XVID')
# out_file = cv2.VideoWriter('',fourcc, 20.0, (360,480))
counter = 22
ret, frame = video_capture.read()
# while True:
heat_map = np.zeros((frame.shape[0], frame.shape[1], 1), np.uint8)
# _ = handle_one(np.ones((540,960,3)), heat_map)
# to be removed
ret, frame = video_capture.read()
canvas, heat_map, track_list = handle_one(frame, heat_map)
while video_capture.isOpened():
counter -= 1
start = time.clock()
print("before track")
for track in track_list:
track.update(frame)
print("in track")
print("after track")
def handle_one(oriImg):
# for visualize
canvas = np.copy(oriImg)
imageToTest = Variable(T.transpose(
T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(), 0), 2, 3), 1,
2),
volatile=True).cuda()
print oriImg.shape
scale = model_['boxsize'] / float(oriImg.shape[0])
print scale
h = int(oriImg.shape[0] * scale)
w = int(oriImg.shape[1] * scale)
pad_h = 0 if (h % model_['stride']
== 0) else model_['stride'] - (h % model_['stride'])
pad_w = 0 if (w % model_['stride']
== 0) else model_['stride'] - (w % model_['stride'])
new_h = h + pad_h
new_w = w + pad_w
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest,
model_['stride'],
model_['padValue'])
imageToTest_padded = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
feed = Variable(T.from_numpy(imageToTest_padded)).cuda()
output1, output2 = model(feed)
heatmap = nn.UpsamplingBilinear2d(
(oriImg.shape[0], oriImg.shape[1])).cuda()(output2)
paf = nn.UpsamplingBilinear2d(
(oriImg.shape[0], oriImg.shape[1])).cuda()(output1)
print heatmap.size()
print paf.size()
print type(heatmap)
heatmap_avg = T.transpose(T.transpose(heatmap[0], 0, 1), 1,
2).data.cpu().numpy()
paf_avg = T.transpose(T.transpose(paf[0], 0, 1), 1, 2).data.cpu().numpy()
all_peaks = []
peak_counter = 0
#maps =
for part in range(18):
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > param_['thre1']))
# peaks_binary = T.eq(
# peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
peaks = zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
vec_y, vec[1])
score_with_dist_prior = sum(
score_midpts) / len(score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > param_['thre2'])
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] + candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=lambda x: x[2],
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3] and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][
indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if (subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int),
2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
print "found = 2"
membership = ((subset[j1] >= 0).astype(int) +
(subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[
partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(
candidate[connection_all[k][i, :2].astype(int),
2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# canvas = cv2.imread(test_image) # B,G,R order
for i in range(18):
for j in range(len(all_peaks[i])):
cv2.circle(canvas,
all_peaks[i][j][0:2],
4,
colors[i],
thickness=-1)
stickwidth = 4
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i]) - 1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1])**2 + (Y[0] - Y[1])**2)**0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv2.ellipse2Poly(
(int(mY), int(mX)), (int(length / 2), stickwidth), int(angle),
0, 360, 1)
cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
return canvas
def post_process(oriImg, canvas, paf_avg, all_peak_idxs, nonzero_vals,
dont_draw):
all_peaks = [[] for x in range(N_JOINTS)]
for i in xrange(all_peak_idxs.shape[0]):
all_peaks[all_peak_idxs[i, 0]].append(
(all_peak_idxs[i, 2], all_peak_idxs[i, 1], nonzero_vals[i], i))
PEAKT = time.time()
connection_all = []
special_k = []
mid_num = 10
KTTT = time.time()
# don't really know how this works
for k in xrange(len(mapIdx)):
idxs = []
for x in mapIdx[k]:
idxs.append(x - 19)
score_mid = paf_avg[idxs, :, :]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
if norm == 0:
continue
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num))
vec_x = np.zeros(len(startend))
vec_y = np.zeros(len(startend))
for I in xrange(len(startend)):
vec_x[I] = score_mid[0,
int(round(startend[I][1])),
int(round(startend[I][0]))]
vec_y[I] = score_mid[1,
int(round(startend[I][1])),
int(round(startend[I][0]))]
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
vec_y, vec[1])
score_with_dist_prior = np.sum(
score_midpts) / len(score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > param_['thre2'])
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] + candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=get_third_item,
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3] and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
cand_tmp = []
for sublist in all_peaks:
for item in sublist:
cand_tmp.append(item)
candidate = np.asarray(cand_tmp)
# don't really know how this works either
for k in xrange(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][
indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if (subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int),
2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
#print "found = 2"
membership = ((subset[j1] >= 0).astype(int) +
(subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[
partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = np.sum(
candidate[connection_all[k][i, :2].astype(int),
2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
#print subset.shape
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
N_BONES = N_JOINTS - 1
if dont_draw:
res = {'found_ppl': []}
for n in range(len(subset)):
skel = {}
for i in range(N_BONES):
index = subset[n][np.array(limbSeq[i]) - 1]
if -1 in index:
continue
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
bone_name = bone_map[i]
skel[bone_name] = ((X[0], Y[0]), (X[1], Y[1]))
res['found_ppl'].append(skel)
return res
p = time.time()
canvas = canvas.copy()
for i in range(N_JOINTS):
for j in range(len(all_peaks[i])):
cv2.circle(canvas,
all_peaks[i][j][0:2],
4,
colors[i],
thickness=-1)
stickwidth = 4
N_BONES = N_JOINTS - 1
for i in range(N_BONES):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i]) - 1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1])**2 + (Y[0] - Y[1])**2)**0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv2.ellipse2Poly(
(int(mY), int(mX)), (int(length / 2), stickwidth), int(angle),
0, 360, 1)
cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
cv2.putText(canvas, str(i), (int(mY), int(mX)),
cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 3)
cv2.putText(canvas, str(i), (int(mY), int(mX)),
cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 0), 2)
canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
print('drawing took: ', time.time() - p)
return canvas
sampler=BalanceSampler(dset, log_val_inds),
batch_size=batch_size,
shuffle=shuffle,
num_workers=num_workers))
return train_loaders, val_loaders
if __name__ == '__main__':
from balance_batch_dataloader import *
csv_path = 'data/train_v2.csv'
img_path = 'data/train-jpg'
img_ext = '.jpg'
dtype = torch.FloatTensor
training_dataset = ResnetOptimizeDataset(csv_path, img_path, dtype)
inds = np.arange(10000)
logical_inds = np.zeros(len(training_dataset))
logical_inds[inds] = 1
bbs = BalanceSampler(training_dataset, logical_inds)
print(len(bbs))
train_loader = BalanceDataLoader(training_dataset,
sampler=bbs,
batch_size=32,
num_workers=1)
for t, (x, y) in enumerate(train_loader):
col_sum = y.sum(dim=0).numpy().flatten()
print(col_sum > 0)
break
def process_image(oriImg, model, model_params, jjac_info):
# for visualize
t0 = time.time()
canvas = np.copy(oriImg)
#imageToTest = Variable(T.transpose(T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(), 0), 2, 3), 1, 2),
# volatile=True).cuda()
#print oriImg.shape
tic = time.time()
scale = model_params['model_']['boxsize'] / float(oriImg.shape[0])
#print scale
h = int(oriImg.shape[0] * scale)
w = int(oriImg.shape[1] * scale)
pad_h = 0 if (h % model_params['model_']['stride']
== 0) else model_params['model_']['stride'] - (
h % model_params['model_']['stride'])
pad_w = 0 if (w % model_params['model_']['stride']
== 0) else model_params['model_']['stride'] - (
w % model_params['model_']['stride'])
new_h = h + pad_h
new_w = w + pad_w
#print 'scaled width and height ({}, {})'.format(h, w)
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_params['model_']['stride'],
model_params['model_']['padValue'])
imageToTest_padded = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
feed = Variable(T.from_numpy(imageToTest_padded)).cuda()
output1, output2 = model(feed)
heatmap = nn.UpsamplingBilinear2d(
(oriImg.shape[0], oriImg.shape[1])).cuda()(output2)
paf = nn.UpsamplingBilinear2d(
(oriImg.shape[0], oriImg.shape[1])).cuda()(output1)
toc = time.time()
model_running_time = toc - tic
#print heatmap.size() # (360, 640, 3)
#print paf.size()
#print type(heatmap)
tic = time.time()
heatmap_avg = T.transpose(T.transpose(heatmap[0], 0, 1), 1,
2).data.cpu().numpy()
#paf_avg = T.transpose(T.transpose(paf[0], 0, 1), 1, 2).data.cpu().numpy()
all_peaks = []
peak_counter = 0
# 13-17 are head
# 10, 11, 12, 13 are legs
# hand_parts = [5, 6, 7, 8, 15, 16]
# 4 is right hand
# todo is it actually left elbow below because inverse? confirm
# 5 was left elbow
# 6 was left shoulder
# 7 is left hand. Right hand only?
# 15 is head. right part of head
body_part_index_to_name = {
'4': 'Right Hand',
'5': 'Left Elbow',
'6': 'Left Shoulder',
'7': 'Left Hand',
'15': 'Left Eye',
'14': '1',
'16': '2',
'17': '3'
}
hand_parts = [4, 7, 15, 16, 17]
# hand_parts = list(range(18))
# for part in range(18):
for part in hand_parts:
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > model_params['param_']['thre1']))
# peaks_binary = T.eq(
# peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
peaks = zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
# print peaks_with_score_and_id
peak_counter += len(peaks)
#print 'heat map peak stuff time is %.5f' % (time.time() - tic)
# connection_all = []
# special_k = []
# mid_num = 10
# for k in range(len(mapIdx)):
# score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
# candA = all_peaks[limbSeq[k][0] - 1]
# candB = all_peaks[limbSeq[k][1] - 1]
# nA = len(candA)
# nB = len(candB)
# indexA, indexB = limbSeq[k]
# if (nA != 0 and nB != 0):
# connection_candidate = []
# for i in range(nA):
# for j in range(nB):
# vec = np.subtract(candB[j][:2], candA[i][:2])
# norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
# vec = np.divide(vec, norm)
#
# startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
# np.linspace(candA[i][1], candB[j][1], num=mid_num))
#
# vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
# for I in range(len(startend))])
# vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
# for I in range(len(startend))])
#
# score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
# score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
# 0.5 * oriImg.shape[0] / norm - 1, 0)
# criterion1 = len(np.nonzero(score_midpts > param_['thre2'])[0]) > 0.8 * len(score_midpts)
# criterion2 = score_with_dist_prior > 0
# if criterion1 and criterion2:
# connection_candidate.append(
# [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
#
# connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
# connection = np.zeros((0, 5))
# for c in range(len(connection_candidate)):
# i, j, s = connection_candidate[c][0:3]
# if (i not in connection[:, 3] and j not in connection[:, 4]):
# connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
# if (len(connection) >= min(nA, nB)):
# break
#
# connection_all.append(connection)
# else:
# special_k.append(k)
# connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
# subset = -1 * np.ones((0, 20))
# candidate = np.array([item for sublist in all_peaks for item in sublist])
#
# for k in range(len(mapIdx)):
# if k not in special_k:
# partAs = connection_all[k][:, 0]
# partBs = connection_all[k][:, 1]
# indexA, indexB = np.array(limbSeq[k]) - 1
#
# for i in range(len(connection_all[k])): # = 1:size(temp,1)
# found = 0
# subset_idx = [-1, -1]
# for j in range(len(subset)): # 1:size(subset,1):
# if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
# subset_idx[found] = j
# found += 1
#
# if found == 1:
# j = subset_idx[0]
# if (subset[j][indexB] != partBs[i]):
# subset[j][indexB] = partBs[i]
# subset[j][-1] += 1
# subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
# elif found == 2: # if found 2 and disjoint, merge them
# j1, j2 = subset_idx
# print "found = 2"
# membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
# if len(np.nonzero(membership == 2)[0]) == 0: # merge
# subset[j1][:-2] += (subset[j2][:-2] + 1)
# subset[j1][-2:] += subset[j2][-2:]
# subset[j1][-2] += connection_all[k][i][2]
# subset = np.delete(subset, j2, 0)
# else: # as like found == 1
# subset[j1][indexB] = partBs[i]
# subset[j1][-1] += 1
# subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
#
# # if find no partA in the subset, create a new subset
# elif not found and k < 17:
# row = -1 * np.ones(20)
# row[indexA] = partAs[i]
# row[indexB] = partBs[i]
# row[-1] = 2
# row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
# subset = np.vstack([subset, row])
#
# # delete some rows of subset which has few parts occur
# deleteIdx = [];
# for i in range(len(subset)):
# if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
# deleteIdx.append(i)
# subset = np.delete(subset, deleteIdx, axis=0)
# canvas = cv2.imread(test_image) # B,G,R order
found_hand = False
found_head = False
found_2_hands = False
tic = time.time()
# for i in range(18):
for i in range(len(hand_parts)):
for j in range(len(all_peaks[i])):
#cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
# Right hand is first in loop. so if it fails we have backup
if body_part_index_to_name[str(hand_parts[i])] == 'Right Hand': #
# BGR order. So blue is below
# cv2.circle(canvas, all_peaks[i][j][0:2], 4, (255, 0, 0), thickness=-1)
#print 'Hand position:', all_peaks[i][j][0:2]
found_hand = True
hand_position = all_peaks[i][j][0:2]
elif body_part_index_to_name[str(hand_parts[i])] == 'Left Hand':
# cv2.circle(canvas, all_peaks[i][j][0:2], 4, (0, 255, 0), thickness=-1)
if not found_hand:
print 'Did not find Right hand but found left hand!'
found_hand = True
hand_position = all_peaks[i][j][0:2]
else:
print 'Found two hands'
found_2_hands = True
second_hand_position = all_peaks[i][j][0:2]
#hand_position = all_peaks[i][j][0:2]
elif body_part_index_to_name[str(hand_parts[i])] == 'Left Eye':
# BGR order. So red is below
# cv2.circle(canvas, all_peaks[i][j][0:2], 4, (0, 0, 255), thickness=-1)
#print 'Head position:', all_peaks[i][j][0:2]
found_head = True
head_position = all_peaks[i][j][0:2]
jjac_info['last_x_head_pos'].append(head_position[1])
if len(jjac_info['last_x_head_pos']) > 10:
jjac_info['last_x_head_pos'].pop()
biggest_diff = max(jjac_info['last_x_head_pos']) - min(
jjac_info['last_x_head_pos'])
if biggest_diff > jjac_info['biggest_diff']:
jjac_info['biggest_diff'] = biggest_diff
# espeak_command = 'Largest y range from last_x_head_pos: {}'.format(biggest_diff)
espeak_command = 'Largest range from last positions: {}'.format(
biggest_diff)
print espeak_command
# os.system(espeak_command)
# jjac_info['head_y_range']
elif body_part_index_to_name[str(hand_parts[i])] == '1':
cv2.circle(canvas,
all_peaks[i][j][0:2],
4, (0, 0, 255),
thickness=-1)
elif body_part_index_to_name[str(hand_parts[i])] == '2':
cv2.circle(canvas,
all_peaks[i][j][0:2],
4, (0, 255, 255),
thickness=-1)
elif body_part_index_to_name[str(hand_parts[i])] == '3':
cv2.circle(canvas,
all_peaks[i][j][0:2],
4, (255, 255, 255),
thickness=-1)
else:
# cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
cv2.circle(canvas,
all_peaks[i][j][0:2],
4, (128, 128, 128),
thickness=-1)
#global how_many_times_hands_went_over_head, hands_over_head
if found_hand and found_head:
# both hands over head or 1 hand over head
if (found_2_hands and hand_position[1] < head_position[1]
and second_hand_position[1] < head_position[1]
) or hand_position[1] < head_position[1]:
print 'Hand is higher than head {} < {}'.format(
hand_position[1], head_position[1])
# if first time hand over head, then increase jumping jack count
if not jjac_info['hands_over_head']:
jjac_info['hands_over_head'] = True
jjac_info['num_jumping_jacks'] += 1
print 'Number of jumping jacks:', jjac_info[
'num_jumping_jacks']
# speak speech command example
# espeak '1 Jumping Jack'
espeak_command = "espeak {}".format(
'\'{} Jumping jack{}\''.format(
jjac_info['num_jumping_jacks'],
's' if jjac_info['num_jumping_jacks'] > 1 else ''))
# os.system(espeak_command) # slows it down?
else:
jjac_info['hands_over_head'] = False
print 'hand is lower than head {} > {}'.format(
hand_position[1], head_position[1])
#print 'Drawing circles time is %.5f' % (time.time() - tic)
cv2.putText(
canvas,
"No. Jumping Jacks: {} ".format(jjac_info['num_jumping_jacks']),
(0, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, 255)
cv2.putText(
canvas, "Hands over head: {}".format(
'Yes' if jjac_info['hands_over_head'] else 'No'), (0, 45),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, 255)
# stickwidth = 4
# for i in range(17):
# for n in range(len(subset)):
# index = subset[n][np.array(limbSeq[i]) - 1]
# if -1 in index:
# continue
# cur_canvas = canvas.copy()
# Y = candidate[index.astype(int), 0]
# X = candidate[index.astype(int), 1]
# mX = np.mean(X)
# mY = np.mean(Y)
# length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
# angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
# polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
# cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
# canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
print 'cv and print: %.5f, Model time: %.5f, Full processing time is %.5f.' % (
time.time() - tic, model_running_time, time.time() - t0)
return canvas
def handle_one(oriImg):
# for visualize
canvas = np.copy(oriImg)
imageToTest = Variable(T.transpose(T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(),0),2,3),1,2),volatile=True).cuda()
print oriImg.shape
scale = model_['boxsize'] / float(oriImg.shape[0])
print scale
h = int(oriImg.shape[0]*scale)
w = int(oriImg.shape[1]*scale)
pad_h = 0 if (h%model_['stride']==0) else model_['stride'] - (h % model_['stride'])
pad_w = 0 if (w%model_['stride']==0) else model_['stride'] - (w % model_['stride'])
new_h = h+pad_h
new_w = w+pad_w
imageToTest = cv2.resize(oriImg, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_['stride'], model_['padValue'])
imageToTest_padded = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,2,0,1))/256 - 0.5
feed = Variable(T.from_numpy(imageToTest_padded)).cuda()
output1,output2 = model(feed)
heatmap = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output2)
paf = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output1)
print heatmap.size()
print paf.size()
print type(heatmap)
heatmap_avg = T.transpose(T.transpose(heatmap[0],0,1),1,2).data.cpu().numpy()
paf_avg = T.transpose(T.transpose(paf[0],0,1),1,2).data.cpu().numpy()
all_peaks = []
peak_counter = 0
#maps =
for part in range(18):
map_ori = heatmap_avg[:,:,part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:,:] = map[:-1,:]
map_right = np.zeros(map.shape)
map_right[:-1,:] = map[1:,:]
map_up = np.zeros(map.shape)
map_up[:,1:] = map[:,:-1]
map_down = np.zeros(map.shape)
map_down[:,:-1] = map[:,1:]
peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param_['thre1']))
# peaks_binary = T.eq(
# peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:,:,[x-19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0]-1]
candB = all_peaks[limbSeq[k][1]-1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if(nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0]*vec[0] + vec[1]*vec[1])
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
score_with_dist_prior = sum(score_midpts)/len(score_midpts) + min(0.5*oriImg.shape[0]/norm-1, 0)
criterion1 = len(np.nonzero(score_midpts > param_['thre2'])[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([i, j, score_with_dist_prior, score_with_dist_prior+candA[i][2]+candB[j][2]])
connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
connection = np.zeros((0,5))
for c in range(len(connection_candidate)):
i,j,s = connection_candidate[c][0:3]
if(i not in connection[:,3] and j not in connection[:,4]):
connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
if(len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:,0]
partBs = connection_all[k][:,1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if(subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
print "found = 2"
membership = ((subset[j1]>=0).astype(int) + (subset[j2]>=0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(candidate[connection_all[k][i,:2].astype(int), 2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = [];
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2]/subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# canvas = cv2.imread(test_image) # B,G,R order
for i in range(18):
for j in range(len(all_peaks[i])):
cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
stickwidth = 4
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i])-1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv2.ellipse2Poly((int(mY),int(mX)), (int(length/2), stickwidth), int(angle), 0, 360, 1)
cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
return canvas
paf_avg = T.transpose(T.transpose(T.squeeze(T.mean(paf_avg, 0)),0,1),1,2).cuda()
heatmap_avg=heatmap_avg.cpu().numpy()
paf_avg = paf_avg.cpu().numpy()
toc =time.time()
print 'time is %.5f'%(toc-tic)
tic = time.time()
all_peaks = []
peak_counter = 0
#maps =
for part in range(18):
map_ori = heatmap_avg[:,:,part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:,:] = map[:-1,:]
map_right = np.zeros(map.shape)
map_right[:-1,:] = map[1:,:]
map_up = np.zeros(map.shape)
map_up[:,1:] = map[:,:-1]
map_down = np.zeros(map.shape)
map_down[:,:-1] = map[:,1:]
peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param_['thre1']))
# peaks_binary = T.eq(
# peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]
