def delete_video(video_id): from multitracker.be import dbconnection db = dbconnection.DatabaseConnection() print(''' TODO: delete all entries from db delete all files from disk ''')
def get_roi_crop_dim(data_dir, project_id, video_id, target_frame_height): """ we need to crop around the centers of bounding boxes the crop dimension should be high enough to see the whole animal but as small as possible different videos show animals in different sizes, so we scan the db for all bboxes and take 98% median as size """ [Hframe, Wframe, _] = cv.imread( glob( os.path.join(data_dir, 'projects', str(project_id), str(video_id), 'frames', 'train', '*.png'))[0]).shape db = dbconnection.DatabaseConnection( file_db=os.path.join(data_dir, 'data.db')) db.execute("select * from bboxes where video_id=%s and is_visible=true;" % video_id) db_boxxes = [x for x in db.cur.fetchall()] ref_vid_id = None #assert len(db_boxxes) > 0, "[*] ERROR: no labeled bounding boxes found! please label at least one bounding box for video " + str(video_id) if len(db_boxxes) == 0: while ref_vid_id is None: try: txtinp = input( ' Sorry, I have not found any labeled bounding boxes. What video has the closest sized animals in videos, that you have already labeled? ' ) ref_vid_id = int(txtinp) except: print( ' ... :( sorry, I did not understand you Dave. Could you give me the video id? (it is a number)' ) db.execute( "select * from bboxes where video_id=%i and is_visible=true;" % ref_vid_id) db_boxxes = [x for x in db.cur.fetchall()] deltas = [] for i, [_, _, frame_idx, individual_id, x1, y1, x2, y2, is_visible] in enumerate(db_boxxes): if is_visible: deltas.extend([x2 - x1, y2 - y1]) deltas = np.array(deltas) if 0: print('deltas', deltas.shape, deltas.mean(), deltas.std(), 'min/max', deltas.min(), deltas.max()) for pp in [50, 75, 90, 95, 96, 97, 98, 99]: print('perc', pp, np.percentile(deltas, pp)) # take 97% percentile with margin of 20% crop_dim = np.percentile(deltas, 97) * 1.1 # scale to target frame height crop_dim = int(crop_dim * Hframe / Hframe) # not bigger than image crop_dim = min(crop_dim, Hframe) crop_dim = int(crop_dim * target_frame_height / Hframe) return crop_dim
def add_video_to_project(base_dir, project_id, source_video_file, fixed_number, half_resolution=False): project_dir = get_project_dir(base_dir, project_id) video_dir = os.path.join(project_dir, "videos") if not os.path.isdir(video_dir): os.makedirs(video_dir) video_file = os.path.join(video_dir, source_video_file.split('/')[-1]) # save video to db query = "insert into videos (name, project_id, fixed_number) values (?,?,?)" conn = dbconnection.DatabaseConnection() video_id = conn.insert(query, (video_file, int(project_id), int(fixed_number))) frames_dir = get_frames_dir(project_dir, video_id) if not os.path.isdir(frames_dir): os.makedirs(os.path.join(frames_dir, "train")) os.makedirs(os.path.join(frames_dir, "test")) #os.makedirs(os.path.join(project_dir,"/data")) # copy file to project directory if not os.path.isfile(video_file): shutil.copy(source_video_file, video_file) '''# sample frames if half_resolution: subprocess.call(['ffmpeg','-i',video_file, '-vf', 'fps=30','-vf', "scale=iw/2:ih/2", frames_dir+'/%05d.png']) else: subprocess.call(['ffmpeg','-i',video_file, '-vf', 'fps=30', frames_dir+'/%05d.png']) # split frames into train/test half/half frames = sorted(glob(os.path.join(frames_dir, '*.png'))) num_frames = len(frames) # split_frame_idx = num_frames // 2 # 50/50 split_frame_idx = len(frames) - 100 for i in range(split_frame_idx): os.rename(frames[i], os.path.join(frames_dir,'train',frames[i].split('/')[-1])) for i in range(split_frame_idx,len(frames),1): os.rename(frames[i], os.path.join(frames_dir,'test',frames[i].split('/')[-1])) ''' print('[*] added video %s to project %i with new video id %i.' % (source_video_file, project_id, video_id)) return video_id
def create_project(name, manager, keypoint_names): conn = dbconnection.DatabaseConnection() query = """ insert into projects (name, manager, keypoint_names, created_at) values(?,?,?,?); """ print('keypoint_names_str', keypoint_names) keypoint_names_str = conn.list_sep.join(keypoint_names) values = (name, manager, keypoint_names_str, util.get_now()) project_id = conn.insert(query, values) print( '[*] created project %i: name: %s, manager: %s, keypoint_names: [%s]' % (project_id, name, manager, ', '.join( keypoint_names_str.split(conn.list_sep)))) return project_id
def export_annotation(args): print('[*] export', args) tstart = time.time() project_id = int(args.project_id) video_ids = [int(iid) for iid in args.video_ids.split(',')] # load source database db = dbconnection.DatabaseConnection( file_db=os.path.join(args.data_dir, 'data.db')) # hacky: make sure that all the frames are written to disk finetune.get_bbox_data( { 'project_id': project_id, 'data_dir': args.data_dir }, args.video_ids, abort_early=True) ## fetch project, video, bounding boxes and keypoints data from source database db.execute("select * from projects where id = %i;" % project_id) project_data = db.cur.fetchone() video_data = {} box_data = {} keypoints_data = {} labeled_files = [] for video_id in video_ids: db.execute("select * from videos where id = %i;" % video_id) video_data[video_id] = [x for x in db.cur.fetchall()][0] db.execute('select * from bboxes where video_id = %i;' % video_id) box_data[video_id] = [x for x in db.cur.fetchall()] db.execute('select * from keypoint_positions where video_id = %i;' % video_id) keypoints_data[video_id] = [x for x in db.cur.fetchall()] for _d in box_data[video_id]: _, _, frame_idx, individual_id, x1, y1, x2, y2, is_visible = _d frame_file = os.path.join(args.data_dir, 'projects', str(project_id), str(video_id), 'frames', 'train', '%05d.png' % int(frame_idx)) labeled_files.append(frame_file) for _d in keypoints_data[video_id]: labid, video_id, frame_idx, keypoint_name, individual_id, keypoint_x, keypoint_y, is_visible = _d frame_file = os.path.join(args.data_dir, 'projects', str(project_id), str(video_id), 'frames', 'train', '%05d.png' % int(frame_idx)) labeled_files.append(frame_file) labeled_files = list(set(labeled_files)) workdir = setup_workdir() file_pickle = os.path.join(workdir, 'data.pkl') pickle.dump([project_data, video_data, box_data, keypoints_data], open(file_pickle, 'wb')) ## copy files to output directory print('[*] copying %i files' % len(labeled_files)) for f in labeled_files: fo = workdir + '/projects/' + f.split('/projects/')[1] dd = os.path.split(fo)[0] if not os.path.isdir(dd): os.makedirs(dd) shutil.copy(f, fo) ## create zip file shutil.make_archive(args.zip[:-4], 'zip', workdir) ## delete work directory shutil.rmtree(workdir) print('[*] export data to %s successful after %i minutes' % (args.zip, int((time.time() - tstart) / 60.)))
def import_annotation(args): print('[*] import', args) tstart = time.time() workdir = setup_workdir() ## extract zip to workdirectory shutil.unpack_archive(args.zip, workdir) [project_data, video_data, box_data, keypoints_data ] = pickle.load(open(os.path.join(workdir, 'data.pkl'), "rb")) old_project_id = project_data[0] old_video_ids = list(video_data.keys()) ## create new database in output directory and insert all the data db_out = dbconnection.DatabaseConnection( file_db=os.path.join(args.data_dir, 'data.db')) # insert new project if args.project_id is None: query = "insert into projects (name, manager, keypoint_names, created_at) values(?,?,?,?);" new_project_id = db_out.insert(query, project_data[1:]) else: new_project_id = int(args.project_id) # insert new videos new_video_ids = [] for video_id in old_video_ids: query = "insert into videos (name, project_id, fixed_number) values (?,?,?)" viddat = video_data[video_id][1:] # cut id viddat = tuple([viddat[0]] + [new_project_id] + [viddat[2]]) new_video_id = db_out.insert(query, viddat) new_video_ids.append(new_video_id) # insert boxes for old_video_id, new_video_id in zip(old_video_ids, new_video_ids): query = "insert into bboxes (video_id, frame_idx, individual_id, x1, y1, x2, y2, is_visible) values (?,?,?,?,?,?,?,?);" for dat in box_data[old_video_id]: dat = dat[2:] # cut id + video_id frame_idx = dat[0] fin = os.path.join(workdir, 'projects', str(old_project_id), str(old_video_id), 'frames', 'train', '%05d.png' % int(frame_idx)) fout = os.path.join(args.data_dir, 'projects', str(new_project_id), str(new_video_id), 'frames', 'train', '%05d.png' % int(frame_idx)) if not os.path.isdir(os.path.split(fout)[0]): os.makedirs(os.path.split(fout)[0]) if not os.path.isfile(fout): shutil.copy(fin, fout) dat = tuple([new_video_id] + list(dat)) db_out.insert(query, dat) # insert keypoints for old_video_id, new_video_id in zip(old_video_ids, new_video_ids): query = "insert into keypoint_positions (video_id, frame_idx, keypoint_name, individual_id, keypoint_x, keypoint_y, is_visible) values (?,?,?,?,?,?,?);" for dat in keypoints_data[old_video_id]: dat = dat[2:] # cut id + video_id frame_idx = dat[0] fin = os.path.join(workdir, 'projects', str(old_project_id), str(old_video_id), 'frames', 'train', '%05d.png' % int(frame_idx)) fout = os.path.join(args.data_dir, 'projects', str(new_project_id), str(new_video_id), 'frames', 'train', '%05d.png' % int(frame_idx)) if not os.path.isdir(os.path.split(fout)[0]): os.makedirs(os.path.split(fout)[0]) if not os.path.isfile(fout): shutil.copy(fin, fout) dat = tuple([new_video_id] + list(dat)) db_out.insert(query, dat) shutil.rmtree(workdir) print('[*] import data from %s to %s successful after %i minutes' % (args.zip, args.data_dir, int((time.time() - tstart) / 60.)))
""" plot results of experiments """ import os from glob import glob import numpy as np import matplotlib.pyplot as plt import json from multitracker.keypoint_detection import model from multitracker.experiments import bg_accuracy from multitracker.experiments import roi_curve from multitracker.be import dbconnection db = dbconnection.DatabaseConnection() project_experiment_dir = os.path.expanduser( '~/checkpoints/experiments/MiceTop') train_ids = [9, 14] test_ids = [13, 14] test_id_in = 14 test_id_out = 13 if 0: train_ids = [1, 3] test_ids = [2, 3] test_id_in = 3 test_id_out = 2 colors = [ 'tab:brown', 'tab:blue', 'tab:orange', 'tab:green', 'tab:red',
parser = argparse.ArgumentParser() parser.add_argument('--model',default=None) parser.add_argument('--project_id',type=int,default=None) parser.add_argument('--video_id',type=int,default=None) parser.add_argument('--upper_bound',type=int,default=4) parser.add_argument('--num_labeling_base',type=int,default=100) parser.add_argument('--open_gallery', dest='open_gallery', action='store_true') parser.add_argument('--postprocess_video', type=str,default=None) parser.add_argument('--postprocess_csv', type=str,default=None) parser.add_argument('--postprocess_upperbound', type=int,default=None) parser.add_argument('--data_dir', required=False, default = os.path.expanduser('~/data/multitracker')) args = parser.parse_args() os.environ['MULTITRACKER_DATA_DIR'] = args.data_dir from multitracker.be import dbconnection db = dbconnection.DatabaseConnection(file_db=os.path.join(args.data_dir, 'data.db')) # load neural network from disk (or init new one) if args.model is not None and os.path.isfile(args.model): assert args.project_id is not None import h5py import tensorflow as tf tf.get_logger().setLevel(logging.ERROR) training_model = tf.keras.models.load_model(h5py.File(args.model, 'r')) print('[*] loaded model %s from disk' % args.model) config = model.get_config(project_id=args.project_id) optimizer = tf.keras.optimizers.Adam(config['kp_lr']) config = model.get_config(int(args.project_id)) else: training_model = None
def randomly_drop_visualiztions(config, project_id, video_id, dst_dir='/tmp/keypoint_heatmap_vis', num=-1, horistack=True, max_height=None, random_maps=False): # take random frames from the db and show their labeling as gaussian heatmaps if not os.path.isdir(dst_dir): os.makedirs(dst_dir) db = dbconnection.DatabaseConnection( file_db=os.path.join(config['data_dir'], 'data.db')) keypoint_names = db.get_keypoint_names(project_id) frame_idxs = [] while len(frame_idxs) == 0: #video_id = db.get_random_project_video(project_id) if video_id is None: raise Exception("[ERROR] no video found for project!") # first get all frames q = "select frame_idx, keypoint_name, individual_id, keypoint_x, keypoint_y from keypoint_positions where video_id=%i and is_visible=true;" % video_id db.execute(q) #frame_idxs = [x[0] for x in db.cur.fetchall()] #frame_idxs = list(set(frame_idxs)) frame_data = {} frame_idxs = [] for [frame_idx, keypoint_name, individual_id, keypoint_x, keypoint_y] in db.cur.fetchall(): #print('frame_idx',frame_idx, 'keypoint_name',keypoint_name, 'individual_id',individual_id, keypoint_x, keypoint_y) if not frame_idx in frame_idxs: frame_idxs.append(frame_idx) frame_data[frame_idx] = [] if random_maps == True: # apply random noise to coordinates as augmentation KA = 400 keypoint_x += np.random.uniform(-KA, KA) keypoint_y += np.random.uniform(-KA, KA) frame_data[frame_idx].append( [keypoint_name, individual_id, keypoint_x, keypoint_y]) shuffle(frame_idxs) print('[*] found keypoint annotation for %i frames' % len(list(frame_idxs))) #shuffle(frame_data) if num > 0: frame_idxs = frame_idxs[:min(num, len(frame_idxs))] for mode in ['train', 'test']: mode_dir = os.path.join(dst_dir, mode) if not os.path.isdir(mode_dir): os.makedirs(mode_dir) frames_dir = os.path.join( config['data_dir'], 'projects', str(project_id), str(video_id), 'frames', 'train' ) #os.path.join(video.get_frames_dir(video.get_project_dir(video.base_dir_default, project_id), video_id)) with Pool(processes=os.cpu_count()) as pool: result_objs = [] for i, frame_idx in enumerate(frame_data.keys()): obj = { 'frame_idx': frame_idx, 'frame_data': frame_data[frame_idx], 'frames_dir': frames_dir, 'project_id': project_id, 'video_id': video_id, 'random_maps': random_maps, 'dst_dir': dst_dir, 'keypoint_names': keypoint_names, 'horistack': horistack, 'max_height': max_height, 'data_dir': config['data_dir'] } result_objs.append(pool.apply_async(write_mp, (obj, ))) #write_mp(obj) results = [result.get() for result in result_objs]
def get_bbox_data(config, video_ids, vis_input_data=0, abort_early=False): seed = 2305 train_image_tensors = [] train_gt_classes_one_hot_tensors = [] train_gt_box_tensors = [] test_image_tensors = [] test_gt_classes_one_hot_tensors = [] test_gt_box_tensors = [] frames_dir = os.path.join(config['data_dir'], 'projects', '%i' % config['project_id']) db = dbconnection.DatabaseConnection( file_db=os.path.join(config['data_dir'], 'data.db')) frame_bboxes = {} for _video_id in video_ids.split(','): _video_id = int(_video_id) db.execute( "select * from bboxes where video_id=%i and is_visible=true order by id;" % _video_id) db_boxxes = [x for x in db.cur.fetchall()] random.Random(4).shuffle(db_boxxes) for dbbox in db_boxxes: _, _, frame_idx, individual_id, x1, y1, x2, y2, is_visible = dbbox frame_idx = '%05d' % int(frame_idx) _key = '%i_%s' % (_video_id, frame_idx) if not _key in frame_bboxes: frame_bboxes[_key] = [] frame_bboxes[_key].append( np.array([float(z) for z in [y1, x1, y2, x2]])) for i, _key in enumerate(frame_bboxes.keys()): frame_bboxes[_key] = np.array(frame_bboxes[_key]) ## open video, check if annotated frames are written to disk, if not, write them frames_missing_on_disk = [] for i, _key in enumerate(frame_bboxes.keys()): video_id, frame_idx = _key.split('_') frame_path = os.path.join(frames_dir, video_id, 'frames', 'train', '%05d.png' % int(frame_idx)) if not os.path.isfile(frame_path): frames_missing_on_disk.append( [video_id, frame_idx, frame_path]) if len(frames_missing_on_disk) > 0: frames_missing_on_disk = sorted(frames_missing_on_disk, key=lambda x: int(x[1])) video_name = db.get_video_name(int(video_id)) video_path = os.path.join(frames_dir, 'videos', video_name) print('sampling %i frames' % len(frames_missing_on_disk), ' from video', video_id, video_name, video_path, os.path.isfile(video_path)) video = cv.VideoCapture(video_path) frame_cnt = 0 frame = 1 while len(frames_missing_on_disk) > 0 and frame is not None: next_target_frame = int(frames_missing_on_disk[0][1]) _, frame = video.read() if frame_cnt == next_target_frame: # write to disk cv.imwrite(frames_missing_on_disk[0][2], frame) print('[*] writing annotated frame %s' % frames_missing_on_disk[0][2]) frames_missing_on_disk = frames_missing_on_disk[1:] frame_cnt += 1 if abort_early: return True # read one arbitray frame to extract height and width of video frames sample_fim = '' while not os.path.isfile(sample_fim): k = int(np.random.uniform(len(list(frame_bboxes.keys())))) sample_fim = os.path.join( frames_dir, video_ids.split(',')[0], 'frames', 'train', '%s.png' % list(frame_bboxes.keys())[k].split('_')[1]) H, W, _ = cv.imread(sample_fim).shape frames = list(frame_bboxes.keys()) #random.Random(4).shuffle(frames) frames = sorted(frames) for i, _key in enumerate(frames): # key is video_id_frame_idx frame_bboxes[_key] = frame_bboxes[_key] / np.array([H, W, H, W]) bboxes = frame_bboxes[_key] if i % 10 > 0: #np.random.uniform() > 0.2: train_image_tensors.append(_key) # str(frame_idx).zfill(5) train_gt_box_tensors.append( tf.convert_to_tensor(bboxes, dtype=tf.float32)) train_gt_classes_one_hot_tensors.append( tf.one_hot( tf.convert_to_tensor( np.ones(shape=[bboxes.shape[0]], dtype=np.int32) - label_id_offset), num_classes)) else: test_image_tensors.append(_key) test_gt_box_tensors.append( tf.convert_to_tensor(bboxes, dtype=tf.float32)) test_gt_classes_one_hot_tensors.append( tf.one_hot( tf.convert_to_tensor( np.ones(shape=[bboxes.shape[0]], dtype=np.int32) - label_id_offset), num_classes)) # maybe use only a part of the train set ddata_train, ddata_test = [train_image_tensors[0]], [test_image_tensors[0]] for i in range(len(train_image_tensors)): if config['use_all_data4train'] or ( not ('data_ratio' in config and np.random.uniform() > config['data_ratio'])): ddata_train.append(train_image_tensors[i]) for i in range(len(test_image_tensors)): ddata_test.append(test_image_tensors[i]) if not 'data_ratio' in config: config['data_ratio'] = 1.0 print( '[*] loaded object detection %s data: training on %i samples, testing on %i samples (data_ratio %f)' % (config['object_detection_backbone'], len(ddata_train), len(ddata_test), config['data_ratio'])) labeling_list_train = tf.data.Dataset.from_tensor_slices(ddata_train) labeling_list_test = tf.data.Dataset.from_tensor_slices(ddata_test) @tf.function def load_im(_key): _video_id = tf.strings.split(_key, '_')[0] frame_idx = tf.strings.split(_key, '_')[1] image_file = tf.strings.join( [frames_dir, '/', _video_id, '/frames/train/', frame_idx, '.png']) image = tf.io.read_file(image_file) image = tf.image.decode_png(image, channels=3) image = tf.cast(image, tf.float32) #image = tf.image.resize(image,[H//2,W//2]) image = tf.image.resize(image, (config['object_detection_resolution'][1], config['object_detection_resolution'][0])) return _key, image data_train = labeling_list_train.map( load_im, num_parallel_calls=tf.data.experimental.AUTOTUNE).shuffle(256).batch( config['object_detection_batch_size']).prefetch( 4 * config['object_detection_batch_size']) #.cache() data_test = labeling_list_test.map( load_im, num_parallel_calls=tf.data.experimental.AUTOTUNE).shuffle(256).batch( config['object_detection_batch_size']).prefetch( 4 * config['object_detection_batch_size']) #.cache() return frame_bboxes, data_train, data_test
def load_roi_dataset(config, mode='train', batch_size=None, video_id=None): video_ids = config['%s_video_ids' % mode] if video_id is not None: video_ids = str(video_id) max_height = None if 'max_height' in config: max_height = config['max_height'] ## create heatmaps for whole frames if len(glob(os.path.join(config['kp_data_dir'], 'train', '*.png'))) == 0: _video_ids = ','.join( list( set(config['train_video_ids'].split(',') + config['test_video_ids'].split(',')))) print('[*] creating training keypoint data for video ids', _video_ids.split(',')) for _video_id in _video_ids.split(','): _video_id = int(_video_id) heatmap_drawing.randomly_drop_visualiztions( config, config['project_id'], _video_id, dst_dir=config['kp_data_dir'], max_height=max_height) image_directory = os.path.join(config['kp_data_dir'], '%s' % mode) print('[*] loading keypoint image data from', image_directory) [Hframe, Wframe, _] = cv.imread( glob( os.path.join(config['data_dir'], 'projects', str(config['project_id']), video_ids.split(',')[0], 'frames', 'train', '*.png'))[0]).shape [Hcomp, Wcomp, _] = cv.imread(glob(os.path.join(image_directory, '*.png'))[0]).shape H = Hcomp w = int(Wframe * Hcomp / Hframe) len_parts = Wcomp // w crop_dim = get_roi_crop_dim(config['data_dir'], config['project_id'], video_ids.split(',')[0], Hcomp) #Hframe) crop_dim_extended_ratio = 1.5 crop_dim_extended = min(Hcomp, crop_dim * crop_dim_extended_ratio) crop_dim_extended = int(crop_dim_extended) crop_dim_extended_ratio = crop_dim_extended / crop_dim #print('crop_dim_extended_ratio',crop_dim_extended_ratio,'crop_dim',crop_dim,'crop_dim_extended',crop_dim_extended) if batch_size is None: batch_size = config['batch_size'] for _mode in ['train', 'test']: if not os.path.isdir(os.path.join(config['kp_roi_dir'], _mode)): os.makedirs(os.path.join(config['kp_roi_dir'], _mode)) ## create data sets if not present for region of interest if len(glob(os.path.join(config['kp_roi_dir'], mode, '*.png'))) == 0: print( '[*] creating cropped regions for each animal to train keypoint prediction ...' ) _video_ids = ','.join( list( set(config['train_video_ids'].split(',') + config['test_video_ids'].split(',')))) #print('_video_ids',_video_ids) db = dbconnection.DatabaseConnection( file_db=os.path.join(config['data_dir'], 'data.db')) for _video_id in _video_ids.split(','): _video_id = int(_video_id) frame_bboxes = {} db.execute( "select * from bboxes where video_id=%i and is_visible=true;" % _video_id) db_boxxes = [x for x in db.cur.fetchall()] shuffle(db_boxxes) for dbbox in db_boxxes: _, _, frame_idx, individual_id, x1, y1, x2, y2, is_visible = dbbox if not frame_idx in frame_bboxes: frame_bboxes[frame_idx] = [] frame_bboxes[frame_idx].append( np.array([float(z) for z in [y1, x1, y2, x2]])) #print('[*] found %i frames for video %s' % ( len(list(frame_bboxes.keys())), _video_id)) with Pool(processes=os.cpu_count()) as pool: result_objs = [] for i, frame_idx in enumerate(frame_bboxes.keys()): frame_bboxes[frame_idx] = np.array(frame_bboxes[frame_idx]) #f = os.path.join(os.path.join(config['kp_data_dir'],mode,'%i_%s.png' % (_video_id,frame_idx)) ) f = os.path.join(config['kp_data_dir'], mode, '%i_%i.png' % (_video_id, int(frame_idx))) if not os.path.isfile(f): f = os.path.join( config['kp_data_dir'], mode, '%i_%05d.png' % (_video_id, int(frame_idx))) obj = { 'Hframe': Hframe, 'Hcomp': Hcomp, 'w': w, 'f': f, 'video_id': _video_id, 'config': config, 'crop_dim_extended': crop_dim_extended, 'len_parts': len_parts, 'frame_idx': frame_idx, 'boxes': frame_bboxes[frame_idx] } result_objs.append( pool.apply_async(write_crop_to_disk, (obj, ))) #write_crop_to_disk(obj) results = [result.get() for result in result_objs] with Pool(processes=os.cpu_count()) as pool: # check homogenous image sizes results_shapefilter = [] sampled_shapes = [] for _mode in ['train', 'test']: _roicrops = glob( os.path.join(config['kp_roi_dir'], _mode, '*.png')) if len(_roicrops) > 0: for ii in range(min(100, len(_roicrops))): sampled_shapes.append( cv.imread( _roicrops[int(np.random.uniform() * len(_roicrops))]).shape[:2]) sampled_H, sampled_W = np.median( np.array(sampled_shapes)[:, 0]), np.median( np.array(sampled_shapes)[:, 1]) for ii, ff in enumerate( glob( os.path.join(config['kp_roi_dir'], _mode, '*.png'))): results_shapefilter.append( pool.apply_async(filter_crop_shape, ({ 'f': ff, 'H': sampled_H, 'W': sampled_W }, ))) #mean_rgb = calc_mean_rgb(config) #print(' globs',glob(os.path.join(config['kp_roi_dir'],'train','*.png'))) hroi, wroicomp = cv.imread( glob(os.path.join(config['kp_roi_dir'], 'train', '*.png'))[0]).shape[:2] wroi = hroi #wroicomp // (1+len(config['keypoint_names'])//3) h = hroi #int(hroi * 0.98) def load_im(image_file): image = tf.io.read_file(image_file) image = tf.image.decode_png(image, channels=3) image = tf.cast(image, tf.float32) comp = [ image[:, ii * wroi:(ii + 1) * wroi, :] for ii in range(wroicomp // wroi) ] # from hstacked to depth stacked comp = [comp[0]] + [c[:, :, ::-1] for c in comp[1:] ] # keypoint heatmaps are in BGR not RGB # preprocess rgb image comp[0] = unet.preprocess(config, comp[0]) comp = tf.concat(comp, axis=2) comp = comp[:, :, :( 3 + len(config['keypoint_names']) )] # cut 'overhanging' channels, that were filled up to reach 3channel png image hh = h if mode == 'train' and config['kp_rotation_augmentation']: # random scale augmentation #if tf.random.uniform([]) > 0.3: # hh = h + int(tf.random.uniform([],-h/10,h/10)) #h += int(tf.random.uniform([],-h/7,h/7)) # apply augmentations # random rotation if tf.random.uniform([]) > 0.5: random_angle = tf.random.uniform([1], -35. * np.pi / 180., 35. * np.pi / 180.) comp = tfa.image.rotate(comp, random_angle) # add background of heatmap background = 255 - tf.reduce_sum(comp[:, :, 3:], axis=2) comp = tf.concat((comp, tf.expand_dims(background, axis=2)), axis=2) # scale down a bit bigger than need before random cropping comp = tf.image.resize(comp, [ int(config['img_height'] * crop_dim_extended_ratio), int(config['img_width'] * crop_dim_extended_ratio) ]) # random crop to counteract imperfect bounding box centers crop = tf.image.random_crop(comp, [ config['img_height'], config['img_width'], 1 + 3 + len(config['keypoint_names']) ]) #crop = comp # split stack into images and heatmaps image = crop[:, :, :3] #image = image - mean_rgb heatmaps = crop[:, :, 3:] / 255. return image, heatmaps ## read data files of all videos file_list = [] for _video_id in video_ids.split(','): file_list.extend( glob( os.path.join(config['kp_roi_dir'], mode, '%i_*.png' % int(_video_id)))) shuffle(file_list) if mode == 'train' and 'data_ratio' in config and config[ 'data_ratio'] > 0 and config[ 'data_ratio'] < 1: #'experiment' in config and config['experiment']=='A': oldlen = len(file_list) file_list = file_list[:int(len(file_list) * config['data_ratio'])] file_list = sorted(file_list) print('[*] cutting training data from %i samples to %i samples' % (oldlen, len(file_list))) print('[*] loaded %i samples for mode %s for videos %s' % (len(file_list), mode, video_ids)) file_list_tf = tf.data.Dataset.from_tensor_slices(file_list) data = file_list_tf.map( load_im, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch( batch_size).prefetch(4 * batch_size) #.cache() #if mode == 'train': data = data.shuffle(512) return data
def main(args): os.environ['MULTITRACKER_DATA_DIR'] = args.data_dir from multitracker.be import dbconnection if args.minutes > 0 or args.video_resolution is not None: tpreprocessvideostart = time.time() sscale = args.video_resolution if args.video_resolution is not None else '' smins = str(args.minutes) if args.minutes > 0 else '' fvo = args.video[:-4] + sscale + smins + args.video[-4:] if not os.path.isfile(fvo): commands = ['ffmpeg'] if args.minutes > 0: commands.extend(['-t', str(int(60. * args.minutes))]) commands.extend(['-i', args.video]) if args.video_resolution is not None: commands.extend([ '-vf', 'scale=%s' % args.video_resolution.replace('x', ':') ]) commands.extend([fvo]) print('[*] preprocess video', ' '.join(commands)) subprocess.call(commands) tpreprocessvideoend = time.time() print('[*] preprocessing of video to %s took %f seconds' % (fvo, tpreprocessvideoend - tpreprocessvideostart)) args.video = fvo tstart = time.time() config = model.get_config(project_id=args.project_id) config['project_id'] = args.project_id config['video'] = args.video config['keypoint_model'] = args.keypoint_model config['autoencoder_model'] = args.autoencoder_model config['objectdetection_model'] = args.objectdetection_model config['train_video_ids'] = args.train_video_ids config['test_video_ids'] = args.test_video_ids config['minutes'] = args.minutes #config['upper_bound'] = db.get_video_fixednumber(args.video_id) #config['upper_bound'] = None config['upper_bound'] = args.upper_bound config['n_blocks'] = 4 config['tracking_method'] = args.tracking_method config['track_tail'] = args.track_tail config['sketch_file'] = args.sketch_file config['file_tracking_results'] = args.output_tracking_results config['use_all_data4train'] = args.use_all_data4train config['object_detection_backbone'] = args.objectdetection_method config = model.update_config_object_detection(config) config['object_detection_resolution'] = [ int(r) for r in args.objectdetection_resolution.split('x') ] config['keypoint_resolution'] = [ int(r) for r in args.keypoint_resolution.split('x') ] config['img_height'], config['img_width'] = config[ 'keypoint_resolution'][::-1] config['kp_backbone'] = args.keypoint_method if 'hourglass' in args.keypoint_method: config['kp_num_hourglass'] = int(args.keypoint_method[9:]) config['kp_backbone'] = 'efficientnetLarge' if args.inference_objectdetection_batchsize > 0: config[ 'inference_objectdetection_batchsize'] = args.inference_objectdetection_batchsize if args.inference_keypoint_batchsize > 0: config[ 'inference_keypoint_batchsize'] = args.inference_keypoint_batchsize if args.delete_all_checkpoints: if os.path.isdir(os.path.expanduser('~/checkpoints/multitracker')): shutil.rmtree(os.path.expanduser('~/checkpoints/multitracker')) if args.data_dir: db = dbconnection.DatabaseConnection( file_db=os.path.join(args.data_dir, 'data.db')) config['data_dir'] = args.data_dir config['kp_data_dir'] = os.path.join( args.data_dir, 'projects/%i/data' % config['project_id']) config['kp_roi_dir'] = os.path.join( args.data_dir, 'projects/%i/data_roi' % config['project_id']) config['keypoint_names'] = db.get_keypoint_names(config['project_id']) config['project_name'] = db.get_project_name(config['project_id']) # <train models> # 1) animal bounding box finetuning -> trains and inferences config['objectdetection_max_steps'] = 30000 # train object detector now = str(datetime.now()).replace(' ', '_').replace(':', '-').split('.')[0] checkpoint_directory_object_detection = os.path.expanduser( '~/checkpoints/multitracker/bbox/vids%s-%s' % (config['train_video_ids'], now)) object_detect_restore = None if 'objectdetection_model' in config and config[ 'objectdetection_model'] is not None: object_detect_restore = config['objectdetection_model'] detection_model = None if object_detect_restore is None: detection_model = finetune.finetune( config, checkpoint_directory_object_detection, checkpoint_restore=object_detect_restore) print('[*] trained object detection model', checkpoint_directory_object_detection) config[ 'object_detection_model'] = checkpoint_directory_object_detection ## crop bbox detections and train keypoint estimation on extracted regions #point_classification.calculate_keypoints(config, detection_file_bboxes) # 2) train autoencoder for tracking appearence vector if config['autoencoder_model'] is None and config[ 'tracking_method'] == 'DeepSORT': config_autoencoder = autoencoder.get_autoencoder_config() config_autoencoder['project_id'] = config['project_id'] config_autoencoder['video_ids'] = natsorted( list( set([int(iid) for iid in config['train_video_ids'].split(',')] + [int(iid) for iid in config['test_video_ids'].split(',')]))) config_autoencoder['project_name'] = config['project_name'] config_autoencoder['data_dir'] = config['data_dir'] config['autoencoder_model'] = autoencoder.train(config_autoencoder) print('[*] trained autoencoder model', config['autoencoder_model']) # 4) train keypoint estimator model if config['keypoint_model'] is None and not config['kp_backbone'] == 'none': config['kp_max_steps'] = 25000 config['keypoint_model'] = roi_segm.train(config) print('[*] trained keypoint_model', config['keypoint_model']) # </train models> # <load models> # load trained object detection model if detection_model is None: # load config json to know which backbone was used with open(os.path.join(config['objectdetection_model'], 'config.json')) as json_file: objconfig = json.load(json_file) objconfig['objectdetection_model'] = config['objectdetection_model'] detection_model = finetune.load_trained_model(objconfig) # load trained autoencoder model for Deep Sort Tracking encoder_model = None if config['tracking_method'] == 'DeepSORT': encoder_model = inference.load_autoencoder_feature_extractor(config) # load trained keypoint model if config['kp_backbone'] == 'none': keypoint_model = None else: keypoint_model = inference.load_keypoint_model( config['keypoint_model']) # </load models> # 3) run bbox tracking deep sort with fixed tracks nms_max_overlap = 1.0 # Non-maxima suppression threshold: Maximum detection overlap max_cosine_distance = 0.2 # Gating threshold for cosine distance metric (object appearance). nn_budget = None # Maximum size of the appearance descriptors gallery. If None, no budget is enforced. print(4 * '\n', config) ttrack_start = time.time() if config['tracking_method'] == 'DeepSORT': deep_sort_app.run(config, detection_model, encoder_model, keypoint_model, args.min_confidence_boxes, args.min_confidence_keypoints, nms_max_overlap, max_cosine_distance, nn_budget) elif config['tracking_method'] == 'UpperBound': upperbound_tracker.run(config, detection_model, encoder_model, keypoint_model, args.min_confidence_boxes, args.min_confidence_keypoints) elif config['tracking_method'] == 'VIoU': viou_tracker.run(config, detection_model, encoder_model, keypoint_model, args.min_confidence_boxes, args.min_confidence_keypoints) ttrack_end = time.time() ugly_big_video_file_out = inference.get_video_output_filepath(config) video_file_out = ugly_big_video_file_out.replace('.avi', '.mp4') convert_video_h265(ugly_big_video_file_out, video_file_out) print( '[*] done tracking after %f minutes. outputting file' % float(int((ttrack_end - ttrack_start) * 10.) / 10.), video_file_out)