def prepare_annotation(project):
"""
Prepare all files needed for annotating the videos in the given project.
"""
project = urllib.parse.unquote(project)
dataset_path = utils.lookup_project_path(project)
# load feature extractor
inference_engine, model_config = utils.load_feature_extractor(dataset_path)
for split in SPLITS:
print(f'\n\tPreparing videos in the {split}-set')
for label in os.listdir(directories.get_videos_dir(
dataset_path, split)):
videos_dir = directories.get_videos_dir(dataset_path, split, label)
frames_dir = directories.get_frames_dir(dataset_path, split, label)
features_dir = directories.get_features_dir(dataset_path,
split,
model_config,
label=label)
compute_frames_features(inference_engine=inference_engine,
videos_dir=videos_dir,
frames_dir=frames_dir,
features_dir=features_dir)
return redirect(url_for("project_details", project=project))
def project_details(project):
"""
Show the details for the selected project.
"""
project = urllib.parse.unquote(project)
path = utils.lookup_project_path(project)
config = utils.load_project_config(path)
stats = {}
for class_name, tags in config['classes'].items():
stats[class_name] = {}
for split in SPLITS:
videos_dir = directories.get_videos_dir(path, split, class_name)
tags_dir = directories.get_tags_dir(path, split, class_name)
stats[class_name][split] = {
'total':
len(os.listdir(videos_dir)),
'tagged':
len(os.listdir(tags_dir)) if os.path.exists(tags_dir) else 0,
}
return render_template('project_details.html',
config=config,
path=path,
stats=stats,
project=config['name'])
def extract_features(path_in, model_config, net, num_layers_finetune, use_gpu, num_timesteps=1, log_fn=print):
# Create inference engine
inference_engine = engine.InferenceEngine(net, use_gpu=use_gpu)
# extract features
for split in SPLITS:
videos_dir = directories.get_videos_dir(path_in, split)
features_dir = directories.get_features_dir(path_in, split, model_config, num_layers_finetune)
video_files = glob.glob(os.path.join(videos_dir, "*", "*.mp4"))
num_videos = len(video_files)
log_fn(f"\nFound {num_videos} videos to process in the {split}-set")
for video_index, video_path in enumerate(video_files):
log_fn(f'\rExtract features from video {video_index + 1} / {num_videos}')
path_features = video_path.replace(videos_dir, features_dir).replace(".mp4", ".npy")
if os.path.isfile(path_features):
log_fn("\tSkipped - feature was already precomputed.")
else:
# Read all frames
frames = extract_frames(video_path=video_path,
inference_engine=inference_engine)
compute_features(path_features=path_features,
inference_engine=inference_engine,
frames=frames,
batch_size=16,
num_timesteps=num_timesteps)
log_fn('\n')
def toggle_project_setting():
"""
Toggle boolean project setting.
"""
data = request.json
path = data['path']
setting = data['setting']
new_status = project_utils.toggle_project_setting(path, setting)
# Update logreg model if assisted tagging was just enabled
if setting == 'assisted_tagging' and new_status:
split = data['split']
label = data['label']
inference_engine, model_config = utils.load_feature_extractor(path)
videos_dir = directories.get_videos_dir(path, split, label)
frames_dir = directories.get_frames_dir(path, split, label)
features_dir = directories.get_features_dir(path,
split,
model_config,
label=label)
# Compute the respective frames and features
compute_frames_and_features(inference_engine=inference_engine,
project_path=path,
videos_dir=videos_dir,
frames_dir=frames_dir,
features_dir=features_dir)
# Re-train the logistic regression model
utils.train_logreg(path=path, split=split, label=label)
return jsonify(setting_status=new_status)
def edit_class(project, class_name):
"""
Edit the class name and tags for an existing class in the given project.
"""
project = urllib.parse.unquote(project)
class_name = urllib.parse.unquote(class_name)
path = project_utils.lookup_project_path(project)
# Get new class name and tags
new_class_name, new_tag1, new_tag2 = utils.get_class_name_and_tags(
request.form)
# Update project config
config = project_utils.load_project_config(path)
del config['classes'][class_name]
config['classes'][new_class_name] = [new_tag1, new_tag2]
project_utils.write_project_config(path, config)
# Update directory names
data_dirs = []
for split in SPLITS:
data_dirs.extend([
directories.get_videos_dir(path, split),
directories.get_frames_dir(path, split),
directories.get_tags_dir(path, split),
])
# Feature directories follow the format <dataset_dir>/<split>/<model>/<num_layers_to_finetune>/<label>
features_dir = directories.get_features_dir(path, split)
if os.path.exists(features_dir):
model_dirs = [
os.path.join(features_dir, model_dir)
for model_dir in os.listdir(features_dir)
]
data_dirs.extend([
os.path.join(model_dir, tuned_layers)
for model_dir in model_dirs
for tuned_layers in os.listdir(model_dir)
])
logreg_dir = directories.get_logreg_dir(path)
if os.path.exists(logreg_dir):
data_dirs.extend([
os.path.join(logreg_dir, model_dir)
for model_dir in os.listdir(logreg_dir)
])
for base_dir in data_dirs:
class_dir = os.path.join(base_dir, class_name)
if os.path.exists(class_dir):
new_class_dir = os.path.join(base_dir, new_class_name)
os.rename(class_dir, new_class_dir)
return redirect(url_for('project_details', project=project))
def setup_project():
"""
Add a new project to the config file. Can also be used for updating an existing project.
"""
data = request.form
name = data['projectName']
path = data['path']
# Initialize project directory
if not os.path.exists(path):
os.mkdir(path)
# Update project config
try:
# Check for existing config file
config = utils.load_project_config(path)
old_name = config['name']
config['name'] = name
except FileNotFoundError:
# Setup new project config
config = {
'name': name,
'date_created': datetime.date.today().isoformat(),
'classes': {},
'use_gpu': False,
'temporal': False,
'video_recording': {
'countdown': 3,
'recording': 5,
},
}
old_name = None
utils.write_project_config(path, config)
# Setup directory structure
for split in SPLITS:
videos_dir = directories.get_videos_dir(path, split)
if not os.path.exists(videos_dir):
os.mkdir(videos_dir)
# Update overall projects config file
projects = utils.load_project_overview_config()
if old_name and old_name in projects:
del projects[old_name]
projects[name] = {
'path': path,
}
utils.write_project_overview_config(projects)
return redirect(url_for('project_details', project=name))
def show_video_list(project, split, label):
"""
Show the list of videos for the given split, class label and project.
If the necessary files for annotation haven't been prepared yet, this is done now.
"""
project = urllib.parse.unquote(project)
path = project_utils.lookup_project_path(project)
split = urllib.parse.unquote(split)
label = urllib.parse.unquote(label)
# load feature extractor
inference_engine, model_config = utils.load_feature_extractor(path)
videos_dir = directories.get_videos_dir(path, split, label)
frames_dir = directories.get_frames_dir(path, split, label)
features_dir = directories.get_features_dir(path,
split,
model_config,
label=label)
tags_dir = directories.get_tags_dir(path, split, label)
logreg_dir = directories.get_logreg_dir(path, model_config, label)
os.makedirs(logreg_dir, exist_ok=True)
os.makedirs(tags_dir, exist_ok=True)
# compute the features and frames missing
compute_frames_and_features(inference_engine=inference_engine,
project_path=path,
videos_dir=videos_dir,
frames_dir=frames_dir,
features_dir=features_dir)
videos = os.listdir(frames_dir)
videos = natsorted(videos, alg=ns.IC)
tagged_list = set(os.listdir(tags_dir))
tagged = [f'{video}.json' in tagged_list for video in videos]
num_videos = len(videos)
num_tagged = len(tagged_list)
num_untagged = num_videos - num_tagged
video_list = zip(videos, tagged, list(range(len(videos))))
return render_template('video_list.html',
video_list=video_list,
split=split,
label=label,
path=path,
project=project,
num_videos=num_videos,
num_tagged=num_tagged,
num_untagged=num_untagged)
def setup_new_project(project_name, path, config=None):
"""
Setup a project directory with a config file and the directories for train and valid videos and
add an entry to the projects overview config.
If an existing project config is given, this one will be used and the project name in there will
be updated.
"""
if not config:
# Setup new project config
config = {
'name': project_name,
'date_created': datetime.date.today().isoformat(),
'tags': {},
'max_tag_index': 0,
'classes': {},
'use_gpu': False,
'temporal': False,
'assisted_tagging': False,
'video_recording': {
'countdown': 3,
'recording': 5,
},
}
else:
config['name'] = project_name
write_project_config(path, config)
# Setup directory structure
for split in SPLITS:
videos_dir = directories.get_videos_dir(path, split)
if not os.path.exists(videos_dir):
os.mkdir(videos_dir)
# Update overall projects config file
projects = load_project_overview_config()
projects[project_name] = {
'path': path,
}
write_project_overview_config(projects)
return config
def add_class(project):
"""
Add a new class to the given project.
"""
data = request.form
project = urllib.parse.unquote(project)
path = project_utils.lookup_project_path(project)
class_name = data['className']
# Update project config
config = project_utils.load_project_config(path)
config['classes'][class_name] = []
project_utils.write_project_config(path, config)
# Setup directory structure
for split in SPLITS:
videos_dir = directories.get_videos_dir(path, split, class_name)
if not os.path.exists(videos_dir):
os.mkdir(videos_dir)
return redirect(url_for("project_details", project=project))
def project_details(project):
"""
Show the details for the selected project.
"""
project = urllib.parse.unquote(project)
path = project_utils.lookup_project_path(project)
config = project_utils.load_project_config(path)
stats = {}
for class_name in config['classes']:
stats[class_name] = {}
for split in SPLITS:
videos_dir = directories.get_videos_dir(path, split, class_name)
tags_dir = directories.get_tags_dir(path, split, class_name)
stats[class_name][split] = {
'total': len(os.listdir(videos_dir)),
'tagged': len(os.listdir(tags_dir)) if os.path.exists(tags_dir) else 0,
'videos': natsorted([video for video in os.listdir(videos_dir) if video.endswith(VIDEO_EXT)], alg=ns.IC)
}
tags = config['tags']
return render_template('project_details.html', config=config, path=path, stats=stats, project=config['name'],
tags=tags)
def add_class(project):
"""
Add a new class to the given project.
"""
project = urllib.parse.unquote(project)
path = utils.lookup_project_path(project)
# Get class name and tags
class_name, tag1, tag2 = utils.get_class_name_and_tags(request.form)
# Update project config
config = utils.load_project_config(path)
config['classes'][class_name] = [tag1, tag2]
utils.write_project_config(path, config)
# Setup directory structure
for split in SPLITS:
videos_dir = directories.get_videos_dir(path, split, class_name)
if not os.path.exists(videos_dir):
os.mkdir(videos_dir)
return redirect(url_for("project_details", project=project))
def train_model(path_in,
path_out,
model_name,
model_version,
num_layers_to_finetune,
epochs,
use_gpu=True,
overwrite=True,
temporal_training=None,
resume=False,
log_fn=print,
confmat_event=None):
os.makedirs(path_out, exist_ok=True)
# Check for existing files
saved_files = [
"last_classifier.checkpoint", "best_classifier.checkpoint",
"config.json", "label2int.json", "confusion_matrix.png",
"confusion_matrix.npy"
]
if not overwrite and any(
os.path.exists(os.path.join(path_out, file))
for file in saved_files):
print(f"Warning: This operation will overwrite files in {path_out}")
while True:
confirmation = input(
"Are you sure? Add --overwrite to hide this warning. (Y/N) ")
if confirmation.lower() == "y":
break
elif confirmation.lower() == "n":
sys.exit()
else:
print('Invalid input')
# Load weights
selected_config, weights = get_relevant_weights(
SUPPORTED_MODEL_CONFIGURATIONS,
model_name,
model_version,
log_fn,
)
backbone_weights = weights['backbone']
if resume:
# Load the last classifier
checkpoint_classifier = torch.load(
os.path.join(path_out, 'last_classifier.checkpoint'))
# Update original weights in case some intermediate layers have been finetuned
update_backbone_weights(backbone_weights, checkpoint_classifier)
# Load backbone network
backbone_network = build_backbone_network(selected_config,
backbone_weights)
# Get the required temporal dimension of feature tensors in order to
# finetune the provided number of layers
if num_layers_to_finetune > 0:
num_timesteps = backbone_network.num_required_frames_per_layer.get(
-num_layers_to_finetune)
if not num_timesteps:
# Remove 1 because we added 0 to temporal_dependencies
num_layers = len(
backbone_network.num_required_frames_per_layer) - 1
msg = (f'ERROR - Num of layers to finetune not compatible. '
f'Must be an integer between 0 and {num_layers}')
log_fn(msg)
raise IndexError(msg)
else:
num_timesteps = 1
# Extract layers to finetune
if num_layers_to_finetune > 0:
fine_tuned_layers = backbone_network.cnn[-num_layers_to_finetune:]
backbone_network.cnn = backbone_network.cnn[0:-num_layers_to_finetune]
# finetune the model
extract_features(path_in,
selected_config,
backbone_network,
num_layers_to_finetune,
use_gpu,
num_timesteps=num_timesteps,
log_fn=log_fn)
# Find label names
label_names = os.listdir(directories.get_videos_dir(path_in, 'train'))
label_names = [x for x in label_names if not x.startswith('.')]
label_names_temporal = ['background']
project_config = load_project_config(path_in)
if project_config:
for temporal_tags in project_config['classes'].values():
label_names_temporal.extend(temporal_tags)
else:
for label in label_names:
label_names_temporal.extend([f'{label}_tag1', f'{label}_tag2'])
label_names_temporal = sorted(set(label_names_temporal))
label2int_temporal_annotation = {
name: index
for index, name in enumerate(label_names_temporal)
}
label2int = {name: index for index, name in enumerate(label_names)}
extractor_stride = backbone_network.num_required_frames_per_layer_padding[
0]
# Create the data loaders
features_dir = directories.get_features_dir(path_in, 'train',
selected_config,
num_layers_to_finetune)
tags_dir = directories.get_tags_dir(path_in, 'train')
train_loader = generate_data_loader(
project_config,
features_dir,
tags_dir,
label_names,
label2int,
label2int_temporal_annotation,
num_timesteps=num_timesteps,
stride=extractor_stride,
temporal_annotation_only=temporal_training,
)
features_dir = directories.get_features_dir(path_in, 'valid',
selected_config,
num_layers_to_finetune)
tags_dir = directories.get_tags_dir(path_in, 'valid')
valid_loader = generate_data_loader(
project_config,
features_dir,
tags_dir,
label_names,
label2int,
label2int_temporal_annotation,
num_timesteps=None,
batch_size=1,
shuffle=False,
stride=extractor_stride,
temporal_annotation_only=temporal_training,
)
# Check if the data is loaded fully
if not train_loader or not valid_loader:
log_fn(
"ERROR - \n "
"\tMissing annotations for train or valid set.\n"
"\tHint: Check if tags_train and tags_valid directories exist.\n")
return
# Modify the network to generate the training network on top of the features
if temporal_training:
num_output = len(label_names_temporal)
else:
num_output = len(label_names)
# modify the network to generate the training network on top of the features
gesture_classifier = LogisticRegression(
num_in=backbone_network.feature_dim,
num_out=num_output,
use_softmax=False)
if resume:
gesture_classifier.load_state_dict(checkpoint_classifier)
if num_layers_to_finetune > 0:
# remove internal padding for training
fine_tuned_layers.apply(set_internal_padding_false)
net = Pipe(fine_tuned_layers, gesture_classifier)
else:
net = gesture_classifier
net.train()
if use_gpu:
net = net.cuda()
lr_schedule = {
0: 0.0001,
int(epochs / 2): 0.00001
} if epochs > 1 else {
0: 0.0001
}
num_epochs = epochs
# Save training config and label2int dictionary
config = {
'backbone_name': selected_config.model_name,
'backbone_version': selected_config.version,
'num_layers_to_finetune': num_layers_to_finetune,
'classifier': str(gesture_classifier),
'temporal_training': temporal_training,
'lr_schedule': lr_schedule,
'num_epochs': num_epochs,
'start_time': str(datetime.datetime.now()),
'end_time': '',
}
with open(os.path.join(path_out, 'config.json'), 'w') as f:
json.dump(config, f, indent=2)
with open(os.path.join(path_out, 'label2int.json'), 'w') as f:
json.dump(
label2int_temporal_annotation if temporal_training else label2int,
f,
indent=2)
# Train model
best_model_state_dict = training_loops(
net,
train_loader,
valid_loader,
use_gpu,
num_epochs,
lr_schedule,
label_names,
path_out,
temporal_annotation_training=temporal_training,
log_fn=log_fn,
confmat_event=confmat_event)
# Save best model
if isinstance(net, Pipe):
best_model_state_dict = {
clean_pipe_state_dict_key(key): value
for key, value in best_model_state_dict.items()
}
torch.save(best_model_state_dict,
os.path.join(path_out, "best_classifier.checkpoint"))
config['end_time'] = str(datetime.datetime.now())
with open(os.path.join(path_out, 'config.json'), 'w') as f:
json.dump(config, f, indent=2)
def submit_annotation():
"""
Submit annotated tags for all frames and save them to a json file.
"""
data = request.form # a multi-dict containing POST data
idx = int(data['idx'])
fps = float(data['fps'])
path = data['path']
project = data['project']
split = data['split']
label = data['label']
video = data['video']
next_frame_idx = idx + 1
frames_dir = directories.get_frames_dir(path, split, label)
tags_dir = directories.get_tags_dir(path, split, label)
description = {'file': f'{video}.mp4', 'fps': fps}
out_annotation = os.path.join(tags_dir, f'{video}.json')
time_annotation = []
for frame_idx in range(int(data['n_images'])):
time_annotation.append(int(data[f'{frame_idx}_tag']))
description['time_annotation'] = time_annotation
with open(out_annotation, 'w') as f:
json.dump(description, f, indent=2)
# Automatic re-training of the logistic regression model
if utils.get_project_setting(path, 'assisted_tagging'):
inference_engine, model_config = utils.load_feature_extractor(path)
videos_dir = directories.get_videos_dir(path, split, label)
frames_dir = directories.get_frames_dir(path, split, label)
features_dir = directories.get_features_dir(path,
split,
model_config,
label=label)
# Compute the respective frames and features
compute_frames_and_features(inference_engine=inference_engine,
project_path=path,
videos_dir=videos_dir,
frames_dir=frames_dir,
features_dir=features_dir)
# Re-train the logistic regression model
utils.train_logreg(path=path, split=split, label=label)
if next_frame_idx >= len(os.listdir(frames_dir)):
return redirect(
url_for('.show_video_list',
project=project,
split=split,
label=label))
return redirect(
url_for('.annotate',
split=split,
label=label,
project=project,
idx=next_frame_idx))
def train_logreg(path):
"""
(Re-)Train a logistic regression model on all annotations that have been submitted so far.
"""
inference_engine, model_config = utils.load_feature_extractor(path)
logreg_dir = directories.get_logreg_dir(path, model_config)
logreg_path = os.path.join(logreg_dir, 'logreg.joblib')
project_config = project_utils.load_project_config(path)
classes = project_config['classes']
all_features = []
all_annotations = []
for split in SPLITS:
for label, class_tags in classes.items():
videos_dir = directories.get_videos_dir(path, split, label)
frames_dir = directories.get_frames_dir(path, split, label)
features_dir = directories.get_features_dir(path,
split,
model_config,
label=label)
tags_dir = directories.get_tags_dir(path, split, label)
if not os.path.exists(tags_dir):
continue
# Compute the respective frames and features
compute_frames_and_features(inference_engine=inference_engine,
project_path=path,
videos_dir=videos_dir,
frames_dir=frames_dir,
features_dir=features_dir)
video_tag_files = os.listdir(tags_dir)
for video_tag_file in video_tag_files:
feature_file = os.path.join(
features_dir, video_tag_file.replace('.json', '.npy'))
annotation_file = os.path.join(tags_dir, video_tag_file)
features = np.load(feature_file)
for f in features:
all_features.append(f.mean(axis=(1, 2)))
with open(annotation_file, 'r') as f:
annotations = json.load(f)['time_annotation']
# Reset tags that have been removed from the class to 'background'
annotations = [
tag_idx if tag_idx in class_tags else 0
for tag_idx in annotations
]
all_annotations.extend(annotations)
# Use low class weight for background and higher weight for all present tags
annotated_tags = set(all_annotations)
class_weight = {tag: 2 for tag in annotated_tags}
class_weight[0] = 0.5
all_features = np.array(all_features)
all_annotations = np.array(all_annotations)
if len(annotated_tags) > 1:
os.makedirs(logreg_dir, exist_ok=True)
logreg = LogisticRegression(C=0.1, class_weight=class_weight)
logreg.fit(all_features, all_annotations)
dump(logreg, logreg_path)
