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)