Update app.py

app.py

@@ -163,87 +163,99 @@ def plot_query(query_idx, data, key_list, val_list, data_path="videos/", fps=15,
 	
 	return plot
 
-def retrieve(file_choice, target_word, query_idx, rows=5, cols=4):
-	
-#     file_choice = fileobj.name
-	print(file_choice)
-	
-	if file_choice=="big_little":
-		file="features/gestsync_feats_biglittle_train.pkl"
-		word_map={"big":0, "little":1}
-		print("Target word: ", target_word)
-		target_label=word_map[target_word]
-
-	print("Target label: ", target_label)
-	key_list = list(word_map.keys())
-	val_list = list(word_map.values())
-		
-	x_train, y_train_actual, files_train, data_rows = load_feats(file)
-	print("X train: {} | Y train: {} | Files train: {} | Data rows test: {}".format(x_train.shape, y_train_actual.shape, files_train.shape, len(data_rows)))
-
-	print("Query idx orig: ", query_idx)
-	if query_idx=="Random":
-		pos_indices = find_ind(target_word, y_train_actual, word_map)
-		query_idx = random.choice(pos_indices)
-	else:
-		query_idx = int(query_idx)
-		
-	print("Query idx: ", query_idx)
-
-	y_train = get_binary_labels_exemplar(y_train_actual, target_label)
-
-	svm_x_train = []
-	svm_y_train = []
-	svm_data_rows = []
-
-	for i in range(len(x_train)):
-		if y_train[i]==-1:
-			svm_x_train.append(x_train[i])
-			svm_y_train.append(y_train[i])
-			svm_data_rows.append(data_rows[i])
-		if i==query_idx:
-			print(query_idx, y_train[i])
-			svm_x_train.append(x_train[i])
-			svm_y_train.append(y_train[i])
-			svm_data_rows.append(data_rows[i])
-			
-	svm_x_train = np.array(svm_x_train).astype(float)
-	svm_y_train = np.array(svm_y_train).astype(float)
-	print("SVM dataset - X: {} | Y: {}".format(svm_x_train.shape,svm_y_train.shape))
-
-	clf = svm.SVC(C=1, kernel="linear")
-	clf.fit(svm_x_train, svm_y_train)
-
-	scores = get_scores_sklearn(x_train, clf)
-#     print(scores)
-	
-	query_plot = plot_query(query_idx, data_rows, key_list, val_list, target_label=target_label)
-	retrieved_plots = plot_ranked_videos(y_train_actual, torch.tensor(scores), data_rows, key_list, val_list, row=rows, col=cols, target_label=target_label)
-	
-	return query_plot, retrieved_plots
+def retrieve(file_choice, target_word, query_idx, rows=6, cols=4):
+
+    print(file_choice)
+
+    if file_choice=="big_little":
+        file="features/gestsync_feats_biglittle_train.pkl"
+        word_map={"big":0, "little":1}
+        print("Target word: ", target_word)
+        if target_word not in list(word_map.keys()):
+            msg="The selected target word is not present in the word classes. For '{}' word classes, please select the target word from the following list: {}".format(file_choice, list(word_map.keys()))
+            return msg, ""
+        target_label=word_map[target_word]
+    elif file_choice=="5_words":
+        file="features/gestsync_feats_5words_balanced_train.pkl"
+        word_map = {'big': 0, 'little': 1, 'next': 2, 'i': 3, 'you': 4}
+        print("Target word: ", target_word)
+        if target_word not in list(word_map.keys()):
+            msg="The selected target word is not present in the word classes. For '{}' word classes, please select the target word from the following list: {}".format(file_choice, list(word_map.keys()))
+            return msg, ""
+        target_label=word_map[target_word]
+
+    print("Target label: ", target_label)
+    key_list = list(word_map.keys())
+    val_list = list(word_map.values())
+
+    x_train, y_train_actual, files_train, data_rows = load_feats(file)
+    print("X train: {} | Y train: {} | Files train: {} | Data rows test: {}".format(x_train.shape, y_train_actual.shape, files_train.shape, len(data_rows)))
+
+    print("Query idx orig: ", query_idx)
+    pos_indices = find_ind(target_word, y_train_actual, word_map)
+    if query_idx=="Random":
+        query_idx = random.choice(pos_indices)
+    else:
+        query_idx = int(query_idx)
+        if query_idx not in pos_indices:
+            msg = "The input query index selected ({}) is not the right index! Please select an index of the positive sample to include in Exemplar-SVM training. <br> For the selected target word '{}', the following are the positive indices in the dataset that can be selected: <br> {}".format(query_idx, target_word, pos_indices)
+            return msg, ""
+
+    print("Query idx: ", query_idx)
+
+    y_train = get_binary_labels_exemplar(y_train_actual, target_label)
+
+    svm_x_train = []
+    svm_y_train = []
+    svm_data_rows = []
+
+    for i in range(len(x_train)):
+        if y_train[i]==-1:
+            svm_x_train.append(x_train[i])
+            svm_y_train.append(y_train[i])
+            svm_data_rows.append(data_rows[i])
+        if i==query_idx:
+            print(query_idx, y_train[i])
+            svm_x_train.append(x_train[i])
+            svm_y_train.append(y_train[i])
+            svm_data_rows.append(data_rows[i])
+
+    svm_x_train = np.array(svm_x_train).astype(float)
+    svm_y_train = np.array(svm_y_train).astype(float)
+    print("SVM dataset - X: {} | Y: {}".format(svm_x_train.shape,svm_y_train.shape))
+
+    clf = svm.SVC(C=1, kernel="linear")
+    clf.fit(svm_x_train, svm_y_train)
+
+    scores = get_scores_sklearn(x_train, clf)
+    #     print(scores)
+
+    query_plot = plot_query(query_idx, data_rows, key_list, val_list, target_label=target_label)
+    retrieved_plots = plot_ranked_videos(y_train_actual, torch.tensor(scores), data_rows, key_list, val_list, row=rows, col=cols, target_label=target_label)
+
+    return query_plot, retrieved_plots
 
 
 if __name__ == "__main__":
 
-	query_choices = ["Random"]+list(range(0, 268))
 	with gr.Blocks() as demo:
-		with gr.Row():
-			gr.HTML("<b> <font size=5> <center> Welcome to Exemplar-SVM Training and Visualization of top-k videos </center> </font> </b>")
-
-		with gr.Row():
-			features = gr.Dropdown(["big_little"], label="Word classes", value="big_little")
-			word = gr.Radio(["big", "little"], label="Target word", value="big")
-			query_idx = gr.Dropdown(query_choices, value="Random", label="Positive sample index to train Exemplar SVM")
-			
-		with gr.Row():
-			submit = gr.Button("Retrieve")
-
-		with gr.Row():
-			query = gr.HTML()
-			
-		with gr.Row():
-			ret_videos = gr.HTML()
-
-		submit.click(retrieve, inputs=[features, word, query_idx], outputs=[query, ret_videos])
+        with gr.Row():
+            gr.HTML("<b> <font size=5> <center> Welcome to Exemplar-SVM Training and Visualization of top-k videos </center> </font> </b>")
+    
+        with gr.Row():
+            features = gr.Dropdown(["big_little", "5_words"], label="Word classes", value="big_little")
+            word = gr.Radio(["big", "little", "i", "you", "next"], label="Target word", value="big")
+            query_idx = gr.Textbox(value="Random", label="Positive sample index to train Exemplar-SVM (Type 'Random' to randomly select the index)")
+    
+        with gr.Row():
+            submit = gr.Button("Retrieve")
+    
+        with gr.Row():
+            query = gr.HTML()
+    
+        with gr.Row():
+            ret_videos = gr.HTML()
+    
+        submit.click(retrieve, inputs=[features, word, query_idx], outputs=[query, ret_videos])
 
 	demo.launch(allowed_paths=["."])