def create_dataframe_faces(): # 取读所有的mask图,构成一个dataframe faces = pd.DataFrame([]) i = 0 arg = setting.args_class() for video_path,json_path,contour_path in zip(arg.videodir,arg.tracked_json,arg.contdir): img_dict = contour_path data = contour_utils.load_json(json_path) for img_name in tqdm(os.listdir(img_dict)): if not(img_name[-1] == 'g'): print(img_name+' is not used') continue if np.random.randn()>0.8: print('add %d th image '%(len(faces+1))) # print(i,img_dict + img_name) frame = plt.imread(img_dict + img_name) frame = np.asarray(frame[:,:,0],dtype = 'uint8') frame[frame == 255] = 1 # flat = frame.flatten() # print(flat.shape) face = pd.Series(frame.flatten(),name = img_name) # print(face) faces = faces.append(face) i += 1 if len(faces)>2500: break if len(faces)>2500: break # ## 可视化 # width, height = frame.shape # fig, axes = plt.subplots(10,10,figsize=(9,9), # subplot_kw = {'xticks':[], 'yticks':[]}, # gridspec_kw = dict(hspace=0.01, wspace=0.01)) # for i, ax in enumerate(axes.flat): # ax.imshow(faces.iloc[i].values.reshape(height,width),cmap='gray') # plt.savefig('love.png') import pickle print('writing file...') with open('./faces.pckl', 'wb') as f: pickle.dump(faces, f)
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))
# for video_path,json_path,contour_path in zip(arg.videodir,arg.tracked_json,arg.contdir): # for i in range(len(arg.videodir)): for i in [1, 2, 3, 4, 5]: video_path = arg.videodir[i] json_path = arg.tracked_json[i] contour_path = arg.contdir[i] # get_samples(video_path,json_path,contour_path) # ------------------- video for debug ----------------------- out_video_path = '{}/contour_{}_debug.mp4'.format('./tmp', i) print(out_video_path) fourcc = cv2.VideoWriter_fourcc(*'mp4v') out_all = cv2.VideoWriter(out_video_path, fourcc, 9.0, (400, 200)) data = contour_utils.load_json(json_path) # ------------------- Read First Frame ----------------------- # cap = cv2.VideoCapture(dir_name, video_name) print(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) db = False show_count = 0 plt.figure(figsize=(15, 5))
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))
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()