Update app.py

app.py

@@ -39,7 +39,7 @@ class VideoClassifier:
     def setup_paths(self):
         self.path = './results'
         if os.path.exists(self.path):
-            shutil.rmtree(self.path)  # Remove the directory if it exists
+            shutil.rmtree(self.path)  
         os.mkdir(self.path)
 
     def setup_gemini_model(self):
@@ -105,11 +105,8 @@ class VideoClassifier:
         return transcript
        
 
-    def classify_video(self,video_input,model_selection=None):
-        if model_selection is not None:
-            self.model = model_selection
-            print(f"Model set to: {self.model}")
-            
+    def classify_video(self,video_input):
+    
         transcript=self.audio_extraction(video_input)
 
         video = cv2.VideoCapture(video_input)
@@ -253,7 +250,16 @@ class VideoClassifier:
 
         # return final_answer, first_video, second_video
         return final_answer
-
+    
+    def save_model_choice(self,model_name):
+        self.model_name = model_name
+        if self.model_name=='mistral':
+            print("Setting up Mistral model for Class Selection")
+            self.setup_mistral_model()
+        else :
+            print("Setting up Gemini model for Class Selection")
+            self.setup_gemini_model()
+        return "Model selected: " + model_name
 
     def launch_interface(self):
         css_code = """
@@ -264,24 +270,33 @@ class VideoClassifier:
         --body-background-fill: #FFFFFF; /* New value */
         }
         """
-        model_dropdown = gr.inputs.Dropdown(choices=['gemini', 'mistral'], label="Select Model", default='gemini')
-        demo = gr.Interface(fn=self.classify_video, inputs=[model_dropdown,"playablevideo" ],allow_flagging='never', 
-                            # examples=[  os.path.join(os.path.dirname(__file__),"American_football_heads_to_India_clip.mp4"),
-                            #           os.path.join(os.path.dirname(__file__), "videos/PersonalFinance_clip.mp4"),
-                            #              os.path.join(os.path.dirname(__file__), "Motorcycle_clip.mp4"),
-                            #              os.path.join(os.path.dirname(__file__), "Spirituality_1_clip.mp4"),
-                            #              os.path.join(os.path.dirname(__file__), "Science_clip.mp4")], 
-                            examples = [
-                                    ['gemini',"American_football_heads_to_India_clip.mp4"],
-                                    ['mistral',"PersonalFinance_clip.mp4"],
-                                    ['gemini',"Motorcycle_clip.mp4"],
-                                    ['mistral',"Spirituality_1_clip.mp4"],
-                                    ['gemini',"Science_clip.mp4"],
-                                ],
+        css_code += """
+        :root {
+        --body-background-fill: #000000; /* New value */
+        }
+        """
+
+        interface_1 = gr.Interface(
+        self.save_model_choice,
+        inputs=gr.Dropdown(choices=['gemini', 'mistral'], label="Select Model", info="Default model: Gemini"),
+        # outputs=interface_1_output,
+        outputs="text"
+        
+    )
+
+
+        demo = gr.Interface(fn=self.classify_video, inputs="playablevideo",allow_flagging='never', examples=[
+                            os.path.join(os.path.dirname(__file__), 
+                                         "American_football_heads_to_India_clip.mp4"),os.path.join(os.path.dirname(__file__), "PersonalFinance_clip.mp4"),
+                                         os.path.join(os.path.dirname(__file__), "Motorcycle_clip.mp4"),
+                                         os.path.join(os.path.dirname(__file__), "Spirituality_1_clip.mp4"),
+                                         os.path.join(os.path.dirname(__file__), "Science_clip.mp4")], 
                         cache_examples=False, outputs=["text"],
                         css=css_code, title="Interactive Advertising Bureau (IAB) compliant Video-Ad classification")
-        demo.launch(debug=True)
-    
+        # demo.launch(debug=True)
+
+        gr.TabbedInterface([interface_1, demo], ["Model Selection", "Video Classification"]).launch(debug=True)
+
     def run_inference(self, video_path,model):
         result = self.classify_video(video_path)
         print(result)
@@ -304,3 +319,7 @@ if __name__ == "__main__":
         vc.run_inference(args.video_path,args.model)
     else:
         print("Error: No video path/model provided for inference mode.")
+
+#Usage
+### python main.py --mode interface
+### python main.py videos/Spirituality_1_clip.mp4 -n 3 --mode inference --model gemini