Update app.py

app.py

@@ -29,6 +29,19 @@ def frame_capture(video_path, num_frames=5):
     vidObj.release()
     return frames
 
+# Function for railway defect analysis with a set prompt
+def railway_defect_analysis(video_path):
+    frames = frame_capture(video_path)
+    images = [PIL.Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) for frame in frames]
+    
+    prompt = "Describe any potential railway defects or issues visible in these video frames."
+    images_with_prompt = [prompt] + images
+    
+    responses = model.generate_content(images_with_prompt)
+    descriptions = [response.text for response in responses]
+    
+    return format_descriptions(descriptions)
+
 # Function to generate text descriptions for frames or answer a specific question
 def analyze_video(video_path, user_question):
     frames = frame_capture(video_path)
@@ -58,15 +71,28 @@ def chat_interaction(video_path, chatbot, user_message):
 
 # Define Gradio interface
 video_input = gr.Video(label="Upload Video", autoplay=True)
-chatbot = gr.Chatbot(label="Video Analysis Chatbot")
-user_input = gr.Textbox(label="Ask something specific about the video", placeholder="E.g., Are there any cars in this video?")
 
-# Create Gradio app
-with gr.Blocks() as demo:
+# Railway defect analysis tab
+with gr.Blocks() as railway_tab:
     with gr.Column():
-        video_input.render()
-        chatbot.render()
-        user_input.render()
-        user_input.submit(fn=chat_interaction, inputs=[video_input, chatbot, user_input], outputs=[user_input, chatbot])
+        gr.Markdown("### Railway Defect Analysis")
+        railway_video_input = video_input.render()
+        defect_analysis_output = gr.Textbox(label="Defect Analysis")
+        analyze_button = gr.Button("Analyze for Railway Defects")
+        analyze_button.click(fn=railway_defect_analysis, inputs=railway_video_input, outputs=defect_analysis_output)
+
+# Interactive chat tab
+with gr.Blocks() as chat_tab:
+    with gr.Column():
+        gr.Markdown("### Interactive Chat Mode")
+        chat_video_input = video_input.render()
+        chatbot = gr.Chatbot(label="Video Analysis Chatbot")
+        user_input = gr.Textbox(label="Ask something specific about the video", placeholder="E.g., Are there any cars in this video?")
+        user_input.submit(fn=chat_interaction, inputs=[chat_video_input, chatbot, user_input], outputs=[user_input, chatbot])
+
+# Combine the two tabs into a single interface
+with gr.Blocks() as demo:
+    with gr.TabbedInterface([railway_tab, chat_tab], ["Railway Defect Analysis", "Interactive Chat"]):
+        pass
 
 demo.launch()