Cleaned Names

app.py

@@ -184,14 +184,18 @@ decision_frameworks = {
 def predict(video_file, video_model_name, audio_model_name, framework_name):
 
     image_processor = VivitImageProcessor.from_pretrained("google/vivit-b-16x2-kinetics400")
-    video_model = torch.load('./' + video_model_name, map_location=torch.device('cpu'))
+    if video_model_name == "60% Accuracy":
+        video_model = torch.load("video_model_60_acc.pth", map_location=torch.device('cpu'))
+    elif video_model_name == "80% Accuracy":
+        video_model = torch.load("video_model_80_acc.pth", map_location=torch.device('cpu'))
 
     model_id = "facebook/wav2vec2-large"
     config = AutoConfig.from_pretrained(model_id, num_labels=6)
     audio_processor = AutoFeatureExtractor.from_pretrained(model_id)
     audio_model = Wav2Vec2ForSequenceClassification.from_pretrained(model_id, config=config)
-    audio_model.load_state_dict(torch.load('./' + audio_model_name, map_location=torch.device('cpu')))
-    audio_model.eval()
+    if audio_model_name == "60% Accuracy":
+        audio_model.load_state_dict(torch.load("audio_model_state_dict_6e.pth", map_location=torch.device('cpu')))
+        audio_model.eval()
 
     delete_directory_path = "./temp/"
 
@@ -219,8 +223,8 @@ def predict(video_file, video_model_name, audio_model_name, framework_name):
 
 inputs = [
     gr.File(label="Upload Video"),
-    gr.Dropdown(["video_model_60_acc.pth", "video_model_80_acc.pth"], label="Select Video Model"),
-    gr.Dropdown(["audio_model_state_dict_6e.pth"], label="Select Audio Model"),
+    gr.Dropdown(["60% Accuracy", "80% Accuracy"], label="Select Video Model"),
+    gr.Dropdown(["60% Accuracy"], label="Select Audio Model"),
     gr.Dropdown(list(decision_frameworks.keys()), label="Select Decision Framework")
 ]