def get_loc(
self,
p=np.array([0, 0, 0]),
o=np.array([0, 0, 0, 1]),
source='CameraTop_frame',
target='map'): #pose = np.array([x,y,z]) : position w.r.t. robot
pp = PoseStamped()
pp.pose.position.x = p[0]
pp.pose.position.y = p[1]
pp.pose.position.z = p[2]
pp.pose.orientation.x = o[0]
pp.pose.orientation.y = o[1]
pp.pose.orientation.z = o[2]
pp.pose.orientation.w = o[3]
#pp.header.stamp = rospy.get_rostime()
pp.header.frame_id = source #'CameraDepth_frame'
#print rospy.Time()
self.transform.waitForTransform(target,
source,
time=rospy.Time(),
timeout=rospy.Duration(3.0))
asdf = self.transform.getLatestCommonTime(target, source)
pp.header.stamp = asdf
result = self.transform.transformPose(target, pp)
result_p = np.array([
result.pose.position.x, result.pose.position.y,
result.pose.position.z
])
result_o = np.array([
result.pose.orientation.x, result.pose.orientation.y,
result.pose.orientation.z, result.pose.orientation.w
])
return result_p, result_o
def get_pos_wrt_robot(self,
cropped_cloud,
x,
y,
size=10,
scan_len=50,
scan='point'):
#scan : point(around), vertical(line)
h = cropped_cloud.shape[0]
w = cropped_cloud.shape[1]
if scan == 'point':
x1 = min(h, max(0, x - size // 2))
x2 = min(h, max(0, x + size // 2))
y1 = min(w, max(0, y - size // 2))
y2 = min(w, max(0, y + size // 2))
roi = cropped_cloud[x1:x2, y1:y2]
mask = roi[:, :, 0] > 0
masked = roi[mask]
if masked.size == 0: return np.array([0, 0, 0])
mask = masked[:, 0] == masked[:, 0].min()
masked = masked[mask]
return masked[0] #self.point_clouds[x,y]
else:
xx1 = min(h, max(0, x - scan_len))
xx2 = min(h, max(0, x + scan_len))
roi = cropped_cloud[xx1:xx2, y - 2:y + 2, :]
mask = roi[:, :, 0] > 0
masked = roi[mask]
if masked.size == 0: return np.array([0, 0, 0])
mask = masked[:, 0] == masked[:, 0].min()
masked = masked[mask]
return masked[0] #self.point_clouds[x,y]
def get_subsets(connection_all, special_k, all_peaks):
# 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(map_ids)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limb_seq[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])
return subset, candidate
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_subsets(connection_all, special_k, all_peaks):
# 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(map_ids)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limb_seq[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])
return subset, candidate
def get_im_names(im_dir, pattern='*.jpg', return_np=True, return_path=False):
import glob
"""Get the image names in a dir. Optional to return numpy array, paths."""
im_paths = glob.glob(osp.join(im_dir, pattern))
im_names = [osp.basename(path) for path in im_paths]
ret = im_paths if return_path else im_names
if return_np:
ret = np.array(ret)
return ret
def __getitem__(self, index):
if (self.mode == "train"):
label = torch.from_numpy(self.labels_train[index]).type(self.dtype)
img_name = self.img_names_train[index]
img = np.array(Image.open(self.csv_path + img_name[0])).T
img = torch.from_numpy(img).type(self.dtype)
return img, label
if (self.mode == "val"):
label = torch.from_numpy(self.labels_val[index]).type(self.dtype)
img_name = self.img_names_val[index]
img = np.array(Image.open(self.csv_path + img_name[0])).T
img = torch.from_numpy(img).type(self.dtype)
return img, label
if (self.mode == "test"):
label = torch.from_numpy(self.labels_test[index]).type(self.dtype)
img_name = self.img_names_test[index]
img = np.array(Image.open(self.csv_path + img_name)).T
img = torch.from_numpy(img).type(self.dtype)
return img, label
def __init__(self, csv_path, file_name, dtype, mode):
train_data = pd.read_csv(csv_path + file_name, header=0, skiprows=[0])
self.dtype = dtype
self.mode = mode
self.csv_path = csv_path
if (mode == "train" or mode == "val"):
labels = train_data.ix[:, 5:7]
img_names = train_data.ix[:, 0:1]
img_names_train, img_names_val, labels_train, labels_val = train_test_split(
img_names,
labels,
random_state=0,
train_size=0.7,
test_size=0.3)
self.N = img_names_train.shape[0]
self.V = img_names_val.shape[0]
self.img_names_train = np.array(img_names_train).reshape(
[self.N, 1])
self.labels_train = np.array(labels_train).reshape([self.N, 2])
self.labels_val = np.array(labels_val).reshape([self.V, 2])
self.img_names_val = np.array(img_names_val).reshape([self.V, 1])
if (mode == "test"):
test_data = pd.read_csv(csv_path + file_name,
header=0,
skiprows=[0])
self.T = test_data.shape[0]
self.img_names_test = np.array(test_data.ix[:, 0:1]).reshape(
[self.T, 1])
self.labels_test = np.array(test_data.ix[:,
5:7]).reshape([self.T, 2])
def generate_label_index_dict(dataset, logical_inds=None):
"""
generate a dict mapping class names to image indices containing that label
logical inds is a set of ones and zeros of the length of your data set
used for train-val split
"""
mlb_matrix = np.array(dataset.y_train)
if logical_inds is None:
logical_inds = np.ones(mlb_matrix.shape[0])
test_matrix = np.eye(17)
labels = dataset.mlb.inverse_transform(test_matrix)
labels = [label[0] for label in labels]
returndict = {}
for label in labels:
returndict[label] = np.array([])
for col_index, label in enumerate(labels):
col = mlb_matrix[:, col_index]
## elementwise multiplication of two arrays of ones and zeros
filtr = (col > 0) * logical_inds
returndict[label] = np.where(filtr)[0]
return returndict
def define_bbox(res, boundaries):
joints_for_bbox = np.where(res[:, 0] != 0)[0]
bbox = np.array([
np.min(res[joints_for_bbox, 0]),
np.max(res[joints_for_bbox, 0]),
np.min(res[joints_for_bbox, 1]),
np.max(res[joints_for_bbox, 1])
]).astype(int)
width = bbox[1] - bbox[0]
for i in range(len(bbox)):
bbox[i] -= width
width *= (-1)
if (bbox[i] < boundaries[i] and width < 0) or (bbox[i] > boundaries[i]
and width > 0):
bbox[i] = boundaries[i]
return bbox
def __init__(self, csv_path, img_path, img_ext, transform=None):
tmp_df = pd.read_csv(csv_path, header=None)
targetFile = img_path + tmp_df[0][0] + img_ext
#self.mlb = MultiLabelBinarizer()
self.img_path = img_path
self.img_ext = img_ext
if transform is None:
self.transform = transforms.Compose([
transforms.Grayscale(), # Default output channels is 1
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
else:
self.transform = transform
self.words = tmp_df[0]
#self.X_train = np.empty([64, 64, 4, len(tmp_df[0])]) # This feels so janky # hardcoding ok because it is custom after all
self.X_train = np.empty(
[64, 64, len(tmp_df[0])]
) # This feels so janky # hardcoding ok because it is custom after all
#self.X_train = tmp_df[0] # X should not be words, shoudl be images
print(self.X_train.shape)
#for word in tmp_df[0]:
for wi in range(len(tmp_df[0])):
#for wi in range(len(tmp_df[0])-1, -1, -1):
#image = imageio.imread(img_path + tmp_df[0][wi] + img_ext)
image = Image.open(img_path + tmp_df[0][wi] + img_ext).convert('L')
#try:
#self.X_train = np.append(self.X_train, np.expand_dims(np.array(image), 3))
#self.X_train[wi] = np.expand_dims(np.array(image), 3)
self.X_train[:, :, wi] = np.array(image)
#print('added ' + str(wi) + ' ' + tmp_df[0][wi])
#except Exception as e:
# print (e)
#self.X_train = self.X_train[:,:,:,1:] # Take everything except the first which was just for starting
vectors = tmp_df.loc[:, 1:300]
self.y_train = vectors.as_matrix()
# To make it more like MNIST
self.imgs = self.X_train
def balanceData(self, x, y):
from count import scan
from count import upsample
from count import downsample
from sklearn.utils import resample
# dic,avg,maxdis,mindis = scan(y)
# print dic
most = ['No Finding']
most2 = ['Infiltration']
maj = ['Effusion', 'Atelectasis']
med = [
'Pneumothorax', 'Mass', 'Pleural_Thickening', 'Consolidation',
'Nodule'
]
mino = [
'Edema', 'Hernia', 'Pneumonia', 'Emphysema', 'Fibrosis',
'Cardiomegaly'
]
rest = [s for s in most + maj + med + mino + most2 if s != 'Hernia']
zt = [
idx for idx in range(len(y))
if 'No Finding' not in y[idx].split('|')
]
# and 'Infiltration' not in y[idx].split('|')]# and 'Effusion' not in y[idx].split('|')
# and 'Atelectasis' not in y[idx].split('|')]
# zt = zt+ downsample(9426.66666667, 10000, most2, maj+med+most, y)
# zt = zt + upsample(9426.66666667,3000, ['Hernia'], maj+most2, y)
# zt = zt + upsample(9426.66666667,10000, mino, maj+most2+med, y)
zt = zt + downsample(9426.66666667, 9426, ['No Finding'], ['-'], y)
# zt = zt + downsample(9426.66666667, 7000, maj,most+most2, y)
# zt = zt + upsample(1,3000, 'Pneumonia', y)
# zt = zt + upsample(1,3000, 'Fibrosis', y)
# zt = zt + upsample(1,3000, 'Edema', y) + upsample(1,3000, 'Emphysema', y)
#print scan([y[i] for i in zt])
from random import shuffle
#print max(zt)
shuffle(zt)
return np.array(zt)
def link_key_point(img_canvas, candidate, subset, stickwidth=4):
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limb_seq[i]) - 1]
if -1 in index:
continue
cur_canvas = img_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])
img_canvas = cv2.addWeighted(img_canvas, 0.4, cur_canvas, 0.6, 0)
return img_canvas
def __init__(self, folder, transform=None):
reader = pd.read_csv(folder + 'Data_Entry_2017.csv')
self.mlb = MultiLabelBinarizer()
imgs = reader['Image Index']
lb = reader['Finding Labels']
IDX = self.balanceData(imgs, lb)
self.xt = pd.Series([imgs[i] for i in IDX])
print(type(imgs))
new_lb = pd.Series(np.array([lb[i] for i in IDX]))
self.transform = transform
# self.yt = [[0.00]*14] * len(self.xt)
# self.yt = np.array(self.yt)
# self.yt = np.reshape(self.yt,(len(self.yt),14))
#one hot
self.yt = self.mlb.fit_transform(lb.str.split('|')).astype(np.float32)
def __getitem__(self, index):
"""
return X_train image and y_train index
"""
img_str = self.X_train[index] + self.img_ext
load_path = os.path.join(self.img_path, img_str)
if self.img_ext == '.jpg':
img = Image.open(load_path)
img = img.convert('RGB')
img = self.transforms(img)
elif self.img_ext == '.npy':
img = np.load(load_path)
img = np.array(img, dtype=np.int32)
img = img.reshape((7, 256, 256))
label = torch.from_numpy(self.y_train[index]).type(self.dtype)
return img, label
def __getitem__(self, index):
"""
return X_train image and y_train index
"""
img_str = self.X_train[index] + self.img_ext
load_path = os.path.join(self.img_path, img_str)
## branching for different backends
if self.img_ext == '.jpg':
img = Image.open(load_path)
img = img.convert('RGB')
img = self.transforms(img)
elif self.img_ext == '.npy':
img = np.load(load_path)
img = np.array(img, dtype=np.int32)
img = img.reshape((7, 256, 256))
return img, self.X_train[index]
def link_key_point(img_canvas, candidate, subset, stickwidth=4):
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limb_seq[i]) - 1]
if -1 in index:
continue
cur_canvas = img_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])
img_canvas = cv2.addWeighted(img_canvas, 0.4, cur_canvas, 0.6, 0)
return img_canvas
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_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)
for i in open(subdirectory + fi + ".dat").readlines():
datContent = ast.literal_eval(i)
except IOError:
print("dat file not found", fi)
continue
if pla == "pitcher":
player = "Pitcher"
else:
player = "Batter"
bottom_b = datContent[player]['bottom']
left_b = datContent[player]['left']
right_b = datContent[player]['right']
top_b = datContent[player]['top']
print(bottom_b, left_b, right_b, top_b)
center = np.array(
[abs(left_b - right_b) / 2.,
abs(bottom_b - top_b) / 2.])
##############
j = 0
#center_dic={}
tic = time.time()
f = subdirectory + fi
print(f)
video_capture = cv2.VideoCapture(f)
#video_capture.set(cv2.CAP_PROP_POS_FRAMES, 100)
#x, y = centers[fi[:-4]]#np.array([abs(top_p+bottom_p)/2., abs(left_p+right_p)/2.])
# center = centers[fi[:-4]] # [1200,700] # center between hips - known position of target person
# print(fi, "center: ", center_dic["Pitcher"])
def f2_score(y_true, y_pred):
# fbeta_score throws a confusing error if inputs are not numpy arrays
y_true, y_pred, = np.array(y_true), np.array(y_pred)
# We need to use average='samples' here, any other average method will generate bogus results
return fbeta_score(y_true, y_pred, beta=2, average='samples')
def f2_score(y_true, y_pred):
# fbeta_score throws a confusing error if inputs are not numpy arrays
y_true, y_pred, = np.array(y_true), np.array(y_pred)
# We need to use average='samples' here, any other average method will generate bogus results
return fbeta_score(y_true, y_pred, beta=2, average='samples')
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
#print dic
#print dic.keys()
# ret, frame = video_capture.read()
# fourcc = cv2.VideoWriter_fourcc(*'DIVX')
#datContent = [i.strip().split() for i in open(path_input_dat).readlines()]
#datContent=ast.literal_eval(datContent[0][0])
bottom_p=datContent['Pitcher']['bottom']
left_p=datContent['Pitcher']['left']
right_p=datContent['Pitcher']['right']
top_p=datContent['Pitcher']['top']
bottom_b=datContent['Batter']['bottom']
left_b=datContent['Batter']['left']
right_b=datContent['Batter']['right']
top_b=datContent['Batter']['top']
center_dic['Pitcher']=np.array([abs(top_p-bottom_p)/2., abs(left_p-right_p)/2.])
center_dic['Batter']=np.array([abs(top_b-bottom_b)/2., abs(left_b-right_b)/2.])
# NEWVID_pitcher = cv2.VideoWriter(output_folder+path_input_vid[path_input_vid.rfind('/')+1:][:-4]+"_video_keypoints_pitcher.mov", fourcc, 30 , (frame[top_p:bottom_p, left_p:right_p].shape[1], frame[top_p:bottom_p, left_p:right_p].shape[0]),1) #2 frames / sec
# NEWVID_batter = cv2.VideoWriter(output_folder+path_input_vid[path_input_vid.rfind('/')+1:][:-4]+"_video_keypoints_batter.mov", fourcc, 30 ,(frame[top_b:bottom_b, left_b:right_b].shape[1], frame[top_b:bottom_b, left_b:right_b].shape[0]),1) #2 frames / sec
df = pd.DataFrame(columns=['Frame','Pitcher','Batter'])
p=0
while True:
# Capture frame-by-frame
ret, frame = video_capture.read()
if frame == None:
break
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
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])
norm = max(norm, MIN_FLT)
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]
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]):
