def get_samples(video_path, json_path, contour_path, arg, targetMouseID):
# data = contour_utils.load_json(json_name)
data = contour_utils.load_json(json_path)
# contour_path = dir_name + 'contour_zexin'
# ------------------- Read First Frame -----------------------
# cap = cv2.VideoCapture(dir_name, video_name)
print('getting contour of %s' % video_path)
cap = cv2.VideoCapture(video_path)
read_flag, frame = cap.read()
width, height, depth = np.asarray(frame).shape
i = 0
if not os.path.exists(contour_path):
mkdir_p(contour_path)
# ----------------- Sample all the mouse mask ----------------------------
count_lessMouse = 0
count_out = 0
count_used = 0
count_frame = 0
a, b = 0, 0
while (read_flag):
count_frame += 1
# if (i<625):
# read_flag, frame = cap.read()
# i += 1
# continue
bad_clip = False
try:
mouses = data['frame_{}'.format(i)]
except:
read_flag, frame = cap.read()
i += 1
continue
for m in mouses:
if not ('keypoints' in m):
bad_clip = True
break
if bad_clip:
count_lessMouse += 1
print(
'\r frame ' + str(i) + ' of ' + video_path +
' does not have enough mice! (%d less, %d out, %d used,%d frames)'
% (count_lessMouse, count_out, count_used, count_frame),
end='')
read_flag, frame = cap.read()
i += 1
continue
for m_id in range(2):
p = np.asarray(mouses[m_id]['keypoints']).reshape(
(arg.joint_num, 3))
for p_id in [0, 3]:
if p[p_id, 0] < 0 or p[p_id, 1] > width or p[p_id, 1] < 0 or p[
p_id, 0] > height: ## bad frame
bad_clip = True
break
if bad_clip:
count_out += 1
print(
'\r frame ' + str(i) + ' of ' + video_path +
' has out picture point!(%d less, %d out, %d used,%d frames)' %
(count_lessMouse, count_out, count_used, count_frame),
end='')
read_flag, frame = cap.read()
i += 1
continue
# 当前frame的pose信息
mouses = data['frame_{}'.format(i)]
if targetMouseID == 0:
pose1 = np.asarray(mouses[0]['keypoints']).reshape(
(arg.joint_num, 3))
pose2 = np.asarray(mouses[1]['keypoints']).reshape(
(arg.joint_num, 3))
else:
pose1 = np.asarray(mouses[1]['keypoints']).reshape(
(arg.joint_num, 3))
pose2 = np.asarray(mouses[0]['keypoints']).reshape(
(arg.joint_num, 3))
# 当前frame中寻找contours
frame = gaussian_filter(frame, sigma=3)
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
ret, thre = cv2.threshold(gray, 50, 255, 0)
contours, hierarchy = cv2.findContours(thre, cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)[-2:]
# 遍历每个contour，看是否符合要求
for contour_id, contour in enumerate(contours):
# if (contour.size>150) and (contour.size<600):
if (contour.size > 150):
# 把contour以binary mask的形式呈现
mask = np.zeros((width, height, depth), dtype='uint8')
cv2.drawContours(mask, contours, contour_id, (255, 255, 255),
-1)
if (np.sum(mask == 255) > width * height / 2):
# mask too large, may be the background
continue
# 假设当前的contour符合要求，但发现有任意一个keypoint不在mask内，就放弃
flag = True
for j in [0, 3]:
if (mask[int(pose1[j, 1]), int(pose1[j, 0]), 0] == 0):
continue
if flag:
# 首先把mask平移到中心
rows, cols, depth = mask.shape
x, y, w, h = cv2.boundingRect(contour)
# M = np.float32([[1,0,w/2-(x+w/2)],[0,1,h/2-(y+h/2)]])
mouse_center_y = int((pose1[0, 0] + pose1[3, 0]) / 2)
mouse_center_x = int((pose1[0, 1] + pose1[3, 1]) / 2)
M = np.float32([[1, 0, height / 2 - mouse_center_y],
[0, 1, width / 2 - mouse_center_x]])
tra = cv2.warpAffine(mask, M, (cols, rows))
# 旋转到身体的轴在x轴上
body = pose1[3, 0:2] - pose1[0, 0:2]
rho, phi = contour_utils.cart2pol(body[0], body[1])
angle = math.degrees(phi)
M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
rot = cv2.warpAffine(tra, M, (cols, rows))
# 裁剪成 200 * 200
ori_crop_halfSize = 400 / 2
final_crop_size = 200
crop = rot[
int(width / 2 -
ori_crop_halfSize):int(width / 2 +
ori_crop_halfSize),
int(height / 2 -
ori_crop_halfSize):int(height / 2 +
ori_crop_halfSize)].copy()
crop = cv2.resize(crop, (final_crop_size, final_crop_size),
interpolation=cv2.INTER_CUBIC)
cv2.imwrite(
contour_path +
'/mask_mouse{}_{}.png'.format(targetMouseID, i), crop)
print('\r' + contour_path +
'/mask_mouse{}_{}.png'.format(targetMouseID, i),
end='',
flush=True)
count_used += 1
continue
read_flag, frame = cap.read()
i += 1
cap.release()
print(
'%d frames dont have 2 mice, %d frames have points outside pic. %d frame will be used, %d frame in total'
% (count_lessMouse, count_out, count_used, count_frame))
def get_samples(video_path,json_path,contour_path,arg):
# data = contour_utils.load_json(json_name)
data = contour_utils.load_json(json_path)
# contour_path = dir_name + 'contour_zexin'
# ------------------- Read First Frame -----------------------
# cap = cv2.VideoCapture(dir_name, video_name)
print('getting contour of %s'%video_path)
cap = cv2.VideoCapture(video_path)
read_flag, frame = cap.read()
width, height,depth = np.asarray(frame).shape
i = 0
if not os.path.exists(contour_path):
mkdir_p(contour_path)
# ----------------- Sample all the mouse mask ----------------------------
count_lessMouse = 0
count_out = 0
count_used = 0
count_frame = 0
a,b = 0,0
while(read_flag):
count_frame += 1
bad_clip = False
try:
# 当前frame的pose信息
mouses = data['frame_{}'.format(i)]
except:
read_flag, frame = cap.read()
i += 1
continue
for m in mouses:
if not ('keypoints' in m):
bad_clip = True
break
if len(mouses)< arg.mice_num:
bad_clip = True
if bad_clip:
count_lessMouse += 1
print('\r frame '+str(i)+' of '+video_path+' does not have enough mice! (%d less, %d out, %d used,%d frames)'%(count_lessMouse,count_out,count_used,count_frame),end='')
read_flag, frame = cap.read()
i += 1
continue
for m_id in range(arg.mice_num):
p = np.asarray(mouses[m_id]['keypoints']).reshape((arg.joint_num,3))
for p_id in range(arg.joint_num):
if p[p_id,0]<0 or p[p_id,1]>width or p[p_id,1]<0 or p[p_id,0]>height: ## bad frame
bad_clip = True
break
if bad_clip:
count_out += 1
print('\r frame '+str(i)+' of '+video_path+' has out picture point!(%d less, %d out, %d used,%d frames)'%(count_lessMouse,count_out,count_used,count_frame),end='')
read_flag, frame = cap.read()
i += 1
continue
# 当前frame中寻找contours
frame = gaussian_filter(frame, sigma=3)
gray = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
ret,thre = cv2.threshold(gray,50,255,0)
contours, hierarchy = cv2.findContours(thre,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)[-2:]
# 遍历每个contour，看是否符合要求
mouseIdUsed = []
mouseMatchedMask = [{'mask':None, 'maxOverlabJoinNum':0,'contour':None} for iii in range(arg.mice_num)]
for contour_id, contour in enumerate(contours):
if (contour.size>150):
# 把contour以binary mask的形式呈现
mask = np.zeros((width,height,depth),dtype = 'uint8')
cv2.drawContours(mask, contours, contour_id, (255,255,255), -1)
if(np.sum(mask==255)>width*height/2):
# mask too large, may be the background
continue
# find the coresponding mouse id for the counter:
for mouseId in range(arg.mice_num):
overlabJoinNum = 0
pose = np.asarray(mouses[mouseId]['keypoints']).reshape((arg.joint_num,3))
for j in range(arg.joint_num):
if mask[int(pose[j,1]),int(pose[j,0]),0] != 0:
overlabJoinNum += 1
if overlabJoinNum > mouseMatchedMask[mouseId]['maxOverlabJoinNum']:
mouseMatchedMask[mouseId]['maxOverlabJoinNum']=overlabJoinNum
mouseMatchedMask[mouseId]['mask']=mask
mouseMatchedMask[mouseId]['contour']=contour
targetMouseIDGenerate = []
for targetMouseID in range(len(mouseMatchedMask)):
if mouseMatchedMask[targetMouseID]['maxOverlabJoinNum'] >0:
pose1 = np.asarray(mouses[targetMouseID]['keypoints']).reshape((arg.joint_num,3))
mask = mouseMatchedMask[targetMouseID]['mask']
# 首先把mask平移到中心
rows,cols,depth = mask.shape
x,y,w,h = cv2.boundingRect(contour)
# M = np.float32([[1,0,w/2-(x+w/2)],[0,1,h/2-(y+h/2)]])
mouse_center_y = int((pose1[0,0]+pose1[3,0])/2)
mouse_center_x = int((pose1[0,1]+pose1[3,1])/2)
M = np.float32([[1,0,height/2-mouse_center_y],[0,1,width/2-mouse_center_x]])
tra = cv2.warpAffine(mask,M,(cols,rows))
# 旋转到身体的轴在x轴上
body = pose1[3,0:2]-pose1[0,0:2]
rho,phi = contour_utils.cart2pol(body[0],body[1])
angle = math.degrees(phi)
M = cv2.getRotationMatrix2D((cols/2,rows/2),angle,1)
rot = cv2.warpAffine(tra,M,(cols,rows))
# 裁剪成 200 * 200
ori_crop_halfSize = 400/2
final_crop_size = 200
crop = rot[int(width/2-ori_crop_halfSize):int(width/2+ori_crop_halfSize),int(height/2-ori_crop_halfSize):int(height/2+ori_crop_halfSize)].copy()
crop = cv2.resize(crop, (final_crop_size, final_crop_size), interpolation=cv2.INTER_CUBIC)
cv2.imwrite(contour_path+ '/mask_mouse{}_{}.png'.format(targetMouseID,i),crop)
targetMouseIDGenerate.append(targetMouseID)
# print('\r' + contour_path+ '/mask_mouse{}_{}.png'.format(targetMouseID,i),end = '',flush=True)
count_used += 1
print('\r counters images generated in:' + contour_path+ ': {}'.format(['mask_mouse{}_{}.png'.format(mid,i) for mid in targetMouseIDGenerate]),end = '',flush=True)
read_flag, frame = cap.read()
i += 1
cap.release()
print('%d frames dont have enough mice, %d frames have points outside pic. %d counter is generated, %d frame in total'%(count_lessMouse,count_out,count_used,count_frame))
flag2 = False
continue
if flag2:
# 首先把mask平移到中心
rows, cols, depth = mask.shape
x, y, w, h = cv2.boundingRect(contour)
M = np.float32([[1, 0, 640 - (x + w / 2)],
[0, 1, 360 - (y + h / 2)]])
tra = cv2.warpAffine(mask, M, (cols, rows))
tra_f = cv2.warpAffine(frame_ori, M, (cols, rows))
# 旋转到身体的轴在x轴上
body = pose1[3, 0:2] - pose1[0, 0:2]
rho, phi = contour_utils.cart2pol(body[0], body[1])
angle = math.degrees(phi)
M = cv2.getRotationMatrix2D((cols / 2, rows / 2),
angle, 1)
rot = cv2.warpAffine(tra, M, (cols, rows))
rot_f = cv2.warpAffine(tra_f, M, (cols, rows))
# 裁剪成 200 * 200
crop = rot[260:460, 540:740].copy()
crop_f = rot_f[260:460, 540:740].copy()
# # show image
# plt.imshow(crop)
# plt.show()
# # plt.close()
def get_samples(video_path, json_path, contour_path):
# data = contour_utils.load_json(json_name)
data = contour_utils.load_json(json_path)
# contour_path = dir_name + 'contour_zexin'
# ------------------- Read First Frame -----------------------
# cap = cv2.VideoCapture(dir_name, video_name)
print('getting contour of %s' % video_path)
cap = cv2.VideoCapture(video_path)
read_flag, frame = cap.read()
width, height, depth = np.asarray(frame).shape
i = 0
if not os.path.exists(contour_path):
mkdir_p(contour_path)
# ----------------- Sample all the mouse mask ----------------------------
while (read_flag):
# if (i<625):
# read_flag, frame = cap.read()
# i += 1
# continue
bad_clip = False
mouses = data['frame_{}'.format(i)]
for m in mouses:
if not ('keypoints' in m):
bad_clip = True
break
if bad_clip:
print('\n frame ' + str(i) + ' of ' + video_path +
' does not have enough mice!')
read_flag, frame = cap.read()
i += 1
continue
# if (len(mouses)<2) : ## bad frame
# read_flag, frame = cap.read()
# i += 1
# continue
for m_id in range(2):
p = np.asarray(mouses[m_id]['keypoints']).reshape((4, 3))
for p_id in [0, 3]:
if p[p_id, 0] < 0 or p[p_id, 1] > width or p[p_id, 1] < 0 or p[
p_id, 0] > height: ## bad frame
bad_clip = True
break
if bad_clip:
print('\n frame ' + str(i) + ' of ' + video_path +
' has out picture point!')
read_flag, frame = cap.read()
i += 1
continue
# 当前frame的pose信息
mouses = data['frame_{}'.format(i)]
# print(mouses)
# print(len(mouses))
# print(i)
pose1 = np.asarray(mouses[0]['keypoints']).reshape((4, 3))
pose2 = np.asarray(mouses[1]['keypoints']).reshape((4, 3))
# 当前frame中寻找contours
frame = gaussian_filter(frame, sigma=3)
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
ret, thre = cv2.threshold(gray, 40, 255, 0)
contours, hierarchy = cv2.findContours(thre, cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)[-2:]
mouse_processed = []
# 遍历每个contour，看是否符合要求
for contour_id, contour in enumerate(contours):
if (contour.size > 150) and (contour.size < 600):
# 把contour以binary mask的形式呈现
mask = np.zeros((width, height, depth), dtype='uint8')
cv2.drawContours(mask, contours, contour_id, (255, 255, 255),
-1)
# 假设当前的contour符合要求，但发现有任意一个keypoint不在mask内，就放弃
flag = True
# 如果pose1的头和尾都在mask里面则mask算第一只老鼠的，pose2类似，否则放弃
if (mask[int(pose1[0, 1]),
int(pose1[0, 0]), 0] !=
0) and (mask[int(pose1[3, 1]),
int(pose1[3, 0]), 0] !=
0) and (0 not in mouse_processed):
mouse_id = 0
mouse_processed.append(0)
elif (mask[int(pose2[0, 1]),
int(pose2[0, 0]), 0] !=
0) and (mask[int(pose2[3, 1]),
int(pose2[3, 0]), 0] !=
0) and (1 not in mouse_processed):
mouse_id = 1
mouse_processed.append(1)
elif (mask[int(pose1[0, 1]),
int(pose1[0, 0]), 0] !=
0) or (mask[int(pose1[3, 1]),
int(pose1[3, 0]), 0] !=
0) and (0 not in mouse_processed):
mouse_id = 0
mouse_processed.append(0)
elif (mask[int(pose2[0, 1]),
int(pose2[0, 0]), 0] !=
0) or (mask[int(pose2[3, 1]),
int(pose2[3, 0]), 0] !=
0) and (1 not in mouse_processed):
mouse_id = 1
mouse_processed.append(1)
else:
continue
# 首先把mask平移到中心
rows, cols, depth = mask.shape
x, y, w, h = cv2.boundingRect(contour)
M = np.float32([[1, 0, 640 - (x + w / 2)],
[0, 1, 360 - (y + h / 2)]])
tra = cv2.warpAffine(mask, M, (cols, rows))
# 旋转到身体的轴在x轴上
body = pose1[3, 0:2] - pose1[0, 0:2]
rho, phi = contour_utils.cart2pol(body[0], body[1])
angle = math.degrees(phi)
M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
rot = cv2.warpAffine(tra, M, (cols, rows))
# 裁剪成 200 * 200
crop = rot[260:460, 540:740].copy()
cv2.imwrite(
contour_path + '/mask_{}_mouse_{}.png'.format(i, mouse_id),
crop)
print('\r' + contour_path +
'/mask_{}_mouse_{}.png'.format(i, mouse_id),
end='',
flush=True)
read_flag, frame = cap.read()
i += 1
cap.release()