def main_function(cam_url, image_url, queue_reader, area): ''' This is main function ''' print("Cam URL: {}".format(cam_url)) print("Area: {}".format(area)) # Variables for tracking faces frame_counter = 0 # Variables holding the correlation trackers and the name per faceid list_of_trackers = TrackersList() clear_tracking_folder() # Model for human detection print('Load YOLO model ...') options = { "model": "./cfg/yolo.cfg", "load": "../models/yolo.weights", "threshold": 0.5 } detector = TFNet(options) # Model for person re-id body_extractor = BodyExtractor() if image_url is not None: imgcv = cv2.imread(image_url) results = detector.return_predict(imgcv) print(results) imgcv = draw_results(imgcv, results) print('Result drawn as ../test-data/result.jpg') cv2.imwrite('../test-data/result.jpg', imgcv) track_results = TrackerResultsDict() predict_dict = {} if Config.CALC_FPS: start_time = time.time() if args.cam_url is not None: frame_reader = URLFrameReader(args.cam_url, scale_factor=1) elif queue_reader is not None: frame_reader = RabbitFrameReader(rabbit_mq, queue_reader) else: print('Empty Image Source') return -1 video_out_fps, video_out_w, video_out_h, = frame_reader.get_info() print(video_out_fps, video_out_w, video_out_h) center = (int(video_out_w / 2), int(video_out_h / 2)) bbox = [ int(center[0] - Config.Frame.ROI_CROP[0] * video_out_w), int(center[1] - video_out_h * Config.Frame.ROI_CROP[1]), int(center[0] + Config.Frame.ROI_CROP[2] * video_out_w), int(center[1] + Config.Frame.ROI_CROP[3] * video_out_h) ] video_out = None if Config.Track.TRACKING_VIDEO_OUT: # new_w = abs(bbox[0] - bbox[2]) # new_h = abs(bbox[1] - bbox[3]) # print(new_w, new_h) video_out = VideoHandle('../data/tracking_video_out.avi', video_out_fps, int(video_out_w), int(video_out_h)) dlib_c_tr_video = cv2.VideoWriter('../data/check_tracking_video.avi', cv2.VideoWriter_fourcc(*'XVID'), video_out_fps, (int(video_out_w), int(video_out_h))) try: while True: # frame_reader.has_next(): frame = frame_reader.next_frame() if frame is None: print("Waiting for the new image") trackers_return_dict, predict_trackers_dict = \ list_of_trackers.check_delete_trackers(None, rabbit_mq) track_results.update_two_dict(trackers_return_dict) predict_dict.update(predict_trackers_dict) # list_of_trackers.trackers_history.check_time(None) # if args.cam_url is not None: # frame_reader = URLFrameReader(args.cam_url, scale_factor=1) # elif queue_reader is not None: # frame_reader = RabbitFrameReader(rabbit_mq, queue_reader) # else: # print('Empty Image Source') # return -1 continue print("Frame ID: %d" % frame_counter) # frame = frame[bbox[1]:bbox[3], bbox[0]:bbox[2], :] if Config.Track.TRACKING_VIDEO_OUT: video_out.tmp_video_out(frame) if Config.CALC_FPS: fps_counter = time.time() tmpi = list_of_trackers.update_dlib_trackers(frame) dlib_c_tr_video.write(tmpi) if frame_counter % Config.Frame.FRAME_INTERVAL == 0: start_time = time.time() detected_bbs = [] results = detector.return_predict(frame) print('Detect time: {}'.format(1 / (time.time() - start_time))) for result in results: if result['label'] != 'person': continue x1 = result['topleft']['x'] y1 = result['topleft']['y'] x2 = result['bottomright']['x'] y2 = result['bottomright']['y'] origin_bb = [x1, y1, x2, y2] detected_bbs.append(origin_bb) print(is_inner_bb(bbox, origin_bb)) if not is_inner_bb(bbox, origin_bb): continue if calc_bb_size(origin_bb) < 10000: continue bb_size = calc_bb_size(origin_bb) body_image = frame[origin_bb[1]:origin_bb[3], origin_bb[0]:origin_bb[2], :] # body_emb = body_extractor.extract_feature(body_image) # TODO: refractor matching_detected_face_with_trackers matched_fid = list_of_trackers.matching_face_with_trackers( frame, frame_counter, origin_bb, None, body_image, body_extractor) # Update list_of_trackers list_of_trackers.update_trackers_list( matched_fid, time.time(), origin_bb, body_image, bb_size, area, frame_counter, None, body_extractor, rabbit_mq) # list_of_trackers.update_trackers_by_tracking(body_extractor, # frame, # area, # frame_counter, # detected_bbs, # rabbit_mq) trackers_return_dict, predict_trackers_dict = \ list_of_trackers.check_delete_trackers(None, rabbit_mq) track_results.update_two_dict(trackers_return_dict) predict_dict.update(predict_trackers_dict) # Check extract trackers history time (str(frame_counter) + '_' + str(i)) # list_of_trackers.trackers_history.check_time(None) frame_counter += 1 if Config.CALC_FPS: print("FPS: %f" % (1 / (time.time() - fps_counter))) if Config.Track.TRACKING_VIDEO_OUT: print('Write track video') video_out.write_track_video(track_results.tracker_results_dict) if Config.Track.PREDICT_DICT_OUT: PickleUtils.save_pickle(Config.PREDICTION_DICT_FILE, predict_dict) except KeyboardInterrupt: print('Keyboard Interrupt !!! Release All !!!') # list_of_trackers.trackers_history.check_time(matcher) if Config.CALC_FPS: print('Time elapsed: {}'.format(time.time() - start_time)) print('Avg FPS: {}'.format( (frame_counter + 1) / (time.time() - start_time))) frame_reader.release() if Config.Track.TRACKING_VIDEO_OUT: print('Write track video') video_out.write_track_video(track_results.tracker_results_dict) video_out.release() if Config.Track.PREDICT_DICT_OUT: PickleUtils.save_pickle(Config.PREDICTION_DICT_FILE, predict_dict)
def generic_function(cam_url, queue_reader, area, face_extractor_model, re_source, multi_thread): ''' This is main function ''' print("Generic function") print("Cam URL: {}".format(cam_url)) print("Area: {}".format(area)) # Variables for tracking faces frame_counter = 0 # Variables holding the correlation trackers and the name per faceid tracking_dirs = glob.glob(Config.TRACKING_DIR + '/*') if tracking_dirs == []: number_of_existing_trackers = 0 else: lof_int_trackid = [ int(tracking_dir.split('/')[-1]) for tracking_dir in tracking_dirs ] number_of_existing_trackers = max(lof_int_trackid) + 1 imageid_to_keyid = {} tracker_manager = TrackerManager( area, imageid_to_keyid=imageid_to_keyid, current_id=number_of_existing_trackers) if Config.Matcher.CLEAR_SESSION: clear_session_folder() global querying_top10_image_ids_queue # mongodb_faceinfo.remove({}) # reg_dict = PickleUtils.read_pickle(Config.REG_IMAGE_FACE_DICT_FILE) # if reg_dict is not None: # for fid in reg_dict: # mongodb_faceinfo.insert_one({'image_id': fid, 'face_id': reg_dict[fid]}) # print('Saved regdict in mongodb as collection regdict') matcher = FaissMatcher() matcher.build( mongodb_faceinfo, imageid_to_keyid=imageid_to_keyid, use_image_id=True) svm_matcher = None if Config.Matcher.CLOSE_SET_SVM: svm_matcher = SVMMatcher() svm_matcher.build(mongodb_faceinfo) track_results = TrackerResultsDict() if Config.CALC_FPS: start_time = time.time() if cam_url is not None: frame_reader = URLFrameReader(cam_url, scale_factor=1) elif queue_reader is not None: frame_reader = RabbitFrameReader(rabbit_mq, queue_reader) elif args.anno_mode: frame_reader = URLFrameReader('./nothing.mp4', scale_factor=1) else: print('Empty Image Source') return -1 if not args.anno_mode: video_out_fps, video_out_w, video_out_h, = frame_reader.get_info() print(video_out_fps, video_out_w, video_out_h) bbox = [ int(Config.Frame.ROI_CROP[0] * video_out_w), int(Config.Frame.ROI_CROP[1] * video_out_h), int(Config.Frame.ROI_CROP[2] * video_out_w), int(Config.Frame.ROI_CROP[3] * video_out_h) ] # bbox = [0, 0, video_out_w, video_out_h] video_out = None if Config.Track.TRACKING_VIDEO_OUT: video_out = VideoHandle(time.time(), video_out_fps, int(video_out_w), int(video_out_h)) # Turn on querying top 10 from queue if Config.QUERY_TOP10_MODE: thread = threading.Thread(target=give_this_id_10_closest_ids) thread.daemon = True thread.start() frame_queue = [] lock = threading.Lock() if multi_thread: is_on = [True] t = threading.Thread(target=(get_frames), args=(frame_queue, frame_reader, re_source, \ cam_url, queue_reader, lock, is_on, )) t.start() try: while True: ms_msg = rabbit_mq.receive_str(Config.Queues.MERGE) ms_flag = 'merge' if ms_msg is None: ms_msg = rabbit_mq.receive_str(Config.Queues.SPLIT) ms_flag = 'split' if ms_msg is not None: merge_anchor, merge_list = extract_info_from_json(ms_msg) while merge_list != []: image_id1 = merge_list.pop() split_merge_id(ms_flag, image_id1, merge_anchor, matcher, preprocessor, face_extractor, tracker_manager, mongodb_dashinfo, mongodb_faceinfo, mongodb_mslog) continue action_msg = rabbit_mq.receive_str(Config.Queues.ACTION) if action_msg is not None: return_dict = json.loads(action_msg) print('Receive: {}'.format(return_dict)) if return_dict['actionType'] == 'getNearest': lock.acquire() querying_top10_image_ids_queue.append(return_dict['data']) lock.release() continue if args.anno_mode: print('Annotation Mode, waiting for new tasks ...') time.sleep(1) continue if multi_thread: if len(frame_queue) > 0: lock.acquire() frame = frame_queue.pop(0) lock.release() else: frame = None else: frame = frame_reader.next_frame() tracker_manager.update_trackers(frame) #do this before check_and_recognize tracker (sync local matcher vs mongodb) trackers_return_dict, recognized_results = update_recognition( self, matcher, svm_matcher, mongodb_faceinfo) for tracker, matching_result_dict in recognized_results: handle_results(tracker, matching_result_dict, imageid_to_keyid = imageid_to_keyid, \ dump=False) # trackers_return_dict = tracker_manager.find_and_process_end_track(mongodb_faceinfo) track_results.merge(trackers_return_dict) tracker_manager.long_term_history.check_time( matcher, mongodb_faceinfo) if frame is None: print("Waiting for the new image") # if Config.Track.RECOGNIZE_FULL_TRACK: # overtime_track_ids = tracker_manager.find_overtime_current_trackers( # time_last=Config.Track.CURRENT_EXTRACR_TIMER-5, # find_unrecognized=True # ) # for overtime_track_id in overtime_track_ids: # checking_tracker, top_predicted_face_ids, matching_result_dict = \ # tracker_manager.check_and_recognize_tracker( # matcher, # overtime_track_id, # mongodb_faceinfo, # svm_matcher) # handle_results(checking_tracker, matching_result_dict, imageid_to_keyid = imageid_to_keyid,\ # dump=False) if re_source and not multi_thread: print('Trying to connect the stream again ...') if cam_url is not None: frame_reader = URLFrameReader(cam_url, scale_factor=1) elif queue_reader is not None: frame_reader = RabbitFrameReader( rabbit_mq, queue_reader) else: print('Empty Image Source') return -1 break print("Frame ID: %d" % frame_counter) if "_rotate" in video: # rotate cw rotation = int(video.split("_")[-1].split(".")[0]) frame = rotate_image_90(frame, rotation) if Config.Track.TRACKING_VIDEO_OUT: video_out.tmp_video_out(frame) if Config.CALC_FPS: fps_counter = time.time() if frame_counter % Config.Frame.FRAME_INTERVAL == 0: # crop frame #frame = frame[bbox[1]:bbox[3], bbox[0]:bbox[2],:] origin_bbs, points = detector.detect_face(frame) if len(origin_bbs) > 0: origin_bbs = [origin_bb[:4] for origin_bb in origin_bbs] display_and_padded_faces = [ CropperUtils.crop_display_face(frame, origin_bb) for origin_bb in origin_bbs ] #cropped_faces = [CropperUtils.crop_face(frame, origin_bb) for origin_bb in origin_bbs] preprocessed_images = [ preprocessor.process( CropperUtils.crop_face(frame, origin_bb)) for origin_bb in origin_bbs ] embeddings_array, _ = face_extractor.extract_features_all_at_once( preprocessed_images) for i, origin_bb in enumerate(origin_bbs): display_face, str_padded_bbox = display_and_padded_faces[i] #cropped_face = CropperUtils.crop_face(frame, origin_bb) # Calculate embedding preprocessed_image = preprocessed_images[i] # preprocessed_image = align_preprocessor.process(frame, points[:,i], aligner, 160) emb_array = np.asarray([embeddings_array[i]]) face_info = FaceInfo( #oigin_bb.tolist(), #emb_array, frame_counter, origin_bb, points[:, i] #display_face, #str_padded_bbox, #points[:,i].tolist() ) is_good_face = handle_filters(points[:, i], coeff_extractor, face_info, preprocessed_image) face_info.is_good = is_good_face # TODO: refractor matching_detected_face_with_trackers matched_track_id = tracker_manager.track(face_info) if not face_info.is_good: print('BAD FACE') continue # Update tracker_manager tracker_manager.update(matched_track_id, frame, face_info) checking_tracker = None # if not Config.Track.RECOGNIZE_FULL_TRACK: # checking_tracker, top_predicted_face_ids, matching_result_dict = \ # tracker_manager.check_and_recognize_tracker( # matcher, # matched_track_id, # mongodb_faceinfo, # svm_matcher) # handle_results(checking_tracker, matching_result_dict, imageid_to_keyid = imageid_to_keyid, \ # dump=True) # if Config.Track.RECOGNIZE_FULL_TRACK: # overtime_track_ids = tracker_manager.find_overtime_current_trackers( # time_last=Config.Track.CURRENT_EXTRACR_TIMER-5, # find_unrecognized=True # ) # for overtime_track_id in overtime_track_ids: # checking_tracker, top_predicted_face_ids, matching_result_dict = \ # tracker_manager.check_and_recognize_tracker( # matcher, # overtime_track_id, # mongodb_faceinfo, # svm_matcher) # handle_results(checking_tracker, matching_result_dict, imageid_to_keyid = imageid_to_keyid, \ # dump=False) frame_counter += 1 if Config.CALC_FPS: print("FPS: %f" % (1 / (time.time() - fps_counter))) if Config.Track.TRACKING_VIDEO_OUT: print('Write track video') video_out.write_track_video(track_results.tracker_results_dict) Config.Track.CURRENT_EXTRACR_TIMER = 0 trackers_return_dict = tracker_manager.find_and_process_end_track( mongodb_faceinfo) Config.Track.HISTORY_CHECK_TIMER = 0 Config.Track.HISTORY_EXTRACT_TIMER = 0 tracker_manager.long_term_history.check_time(matcher, mongodb_faceinfo) except KeyboardInterrupt: if multi_thread: is_on[0] = False t.join() print('Keyboard Interrupt !!! Release All !!!') Config.Track.CURRENT_EXTRACR_TIMER = 0 trackers_return_dict = tracker_manager.find_and_process_end_track( mongodb_faceinfo) Config.Track.HISTORY_CHECK_TIMER = 0 Config.Track.HISTORY_EXTRACT_TIMER = 0 tracker_manager.long_term_history.check_time(matcher, mongodb_faceinfo) if Config.CALC_FPS: print('Time elapsed: {}'.format(time.time() - start_time)) print('Avg FPS: {}'.format( (frame_counter + 1) / (time.time() - start_time))) frame_reader.release() if Config.Track.TRACKING_VIDEO_OUT: print('Write track video') video_out.write_track_video(track_results.tracker_results_dict) video_out.release()
def generic_function(cam_url, queue_reader, area, face_extractor_model, re_source): ''' This is main function ''' print("Generic function") print("Cam URL: {}".format(cam_url)) print("Area: {}".format(area)) # TODO: init logger, modulize this? # Variables for tracking faces frame_counter = 0 if Config.Matcher.CLEAR_SESSION: clear_session_folder() if Config.Mode.CALC_FPS: start_time = time.time() if cam_url is not None: frame_reader = URLFrameReader(cam_url) else: print('Empty Image Source') return -1 video_out_fps, video_out_w, video_out_h, = frame_reader.get_info() print(video_out_fps, video_out_w, video_out_h) video_out = None if Config.Track.TRACKING_VIDEO_OUT: video_out = VideoHandle(time.time(), video_out_fps, int(video_out_w), int(video_out_h)) db = DashboardDatabase(use_image_id=True) rabbit_mq = RabbitMQ((Config.Rabbit.USERNAME, Config.Rabbit.PASSWORD), (Config.Rabbit.IP_ADDRESS, Config.Rabbit.PORT)) matcher = KdTreeMatcher() matcher.build(db) # find current track import glob tracking_dirs = glob.glob(Config.Dir.TRACKING_DIR + '/*') if tracking_dirs == []: number_of_existing_trackers = 0 else: lof_int_trackid = [ int(tracking_dir.split('/')[-1]) for tracking_dir in tracking_dirs ] number_of_existing_trackers = max(lof_int_trackid) + 1 mode = 'video' # video, live ''' # Feature 1: Find Merge Split splitMerge = pipe.SplitMergeThread(database=db, rabbit_mq=rabbit_mq, matcher=matcher) splitMerge.daemon = True splitMerge.start() # Feature 2: Find similar findSimilarFaceThread = pipe.FindSimilarFaceThread(database=db, rabbit_mq=rabbit_mq) findSimilarFaceThread.daemon = True findSimilarFaceThread.start() ''' # main program stage stageDetectFace = pipe.Stage(pipe.FaceDetectWorker, 1) stagePreprocess = pipe.Stage(pipe.PreprocessDetectedFaceWorker, 1) stageDistributor = pipe.Stage(pipe.FaceDistributorWorker, 1) stageExtract = pipe.Stage(pipe.FaceExtractWorker, 1) stageTrack = pipe.Stage(pipe.FullTrackTrackingWorker, 1, area=area, database=db, matcher=matcher, init_tracker_id=number_of_existing_trackers) stageResultToTCH = pipe.Stage(pipe.SendToDashboardWorker, 1, database=db, rabbit_mq=rabbit_mq) stageStorage = pipe.Stage(pipe.DashboardStorageWorker, 1) stageDatabase = pipe.Stage(pipe.DashboardDatabaseWorker, 1, database=db) stageDetectFace.link(stagePreprocess) stagePreprocess.link(stageDistributor) stageDistributor.link(stageExtract) stageExtract.link(stageTrack) stageTrack.link(stageResultToTCH) stageTrack.link(stageStorage) stageTrack.link(stageDatabase) if Config.Track.TRACKING_VIDEO_OUT: stageVideoOut = pipe.Stage(pipe.VideoWriterWorker, 1, database=db, video_out=video_out) stageTrack.link(stageVideoOut) pipeline = pipe.Pipeline(stageDetectFace) print('Begin') try: while frame_reader.has_next(): #continue frame = frame_reader.next_frame() if frame is None: if mode == 'video': print("Wait for executor to finish it jobs") pipeline.put(None) break if mode == 'live': if re_source: print('Trying to connect the stream again ...') if cam_url is not None: frame_reader = URLFrameReader(cam_url, scale_factor=1, should_crop=True) continue print('Read frame', frame_counter, frame.shape) if frame_counter % Config.Frame.FRAME_INTERVAL == 0: # timer = Timer(frame_counter) task = pipe.Task(pipe.Task.Frame) task.package(frame=frame, frame_info=frame_counter) pipeline.put(task) # pipeline.put((frame, frame_counter, timer)) frame_counter += 1 print('Time elapsed: {}'.format(time.time() - start_time)) print('Avg FPS: {}'.format( (frame_counter + 1) / (time.time() - start_time))) frame_reader.release() ''' splitMerge.join() findSimilarFaceThread.join() ''' except KeyboardInterrupt: if Config.Track.TRACKING_VIDEO_OUT: video_out.release_tmp() pipeline.put(None) print('Keyboard Interrupt !!! Release All !!!') print('Time elapsed: {}'.format(time.time() - start_time)) print('Avg FPS: {}'.format( (frame_counter + 1) / (time.time() - start_time))) frame_reader.release() '''