Update app.py

app.py

@@ -70,9 +70,9 @@ def extract_frames_from_video(video, frame_points=[0, 0.5, 1], max_size=(800, 80
     cap.release()
     return frames
 
-def analyze_construction_image(images=None, video=None):
-    if not images and video is None:
-        logger.warning("No images or video provided")
+def analyze_construction_media(media):
+    if not media:
+        logger.warning("No media provided")
         return [("No input", "Error: Please upload images or a video for analysis.")]
 
     try:
@@ -84,72 +84,52 @@ def analyze_construction_image(images=None, video=None):
                        "If it does, identify any safety issues or hazards, categorize them, and provide a detailed description, "
                        "and suggest steps to resolve them. If it's not a construction site, simply state that")
 
-        if images:
-            for i, image_file in enumerate(images):
-                image = Image.open(image_file.name)
+        for i, file in enumerate(media):
+            file_type = file.name.split('.')[-1].lower()
+            
+            if file_type in ['jpg', 'jpeg', 'png', 'gif']:
+                # Handle image
+                image = Image.open(file.name)
                 resized_image = resize_image(image)
                 image_data_url = f"data:image/png;base64,{encode_image(resized_image)}"
-                messages = [
-                    {
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "text",
-                                "text": f"{instruction}\n\nAnalyze this image (Image {i+1}/{len(images)}). First, determine if it's a construction site. If it is, explain the image in detail, focusing on safety aspects. If it's not, briefly describe what you see."
-                            },
-                            {
-                                "type": "image_url",
-                                "image_url": {
-                                    "url": image_data_url
-                                }
-                            }
-                        ]
-                    }
-                ]
-                completion = client.chat.completions.create(
-                    model="llama-3.2-90b-vision-preview",
-                    messages=messages,
-                    temperature=0.7,
-                    max_tokens=1000,
-                    top_p=1,
-                    stream=False,
-                    stop=None
-                )
-                result = completion.choices[0].message.content
-                results.append((f"Image {i+1} analysis", result))
+                content_type = "image"
+            elif file_type in ['mp4', 'avi', 'mov', 'wmv']:
+                # Handle video
+                frames = extract_frames_from_video(file.name)
+                image_data_url = f"data:image/png;base64,{encode_image(frames[0])}"  # Use first frame
+                content_type = "video"
+            else:
+                results.append((f"File {i+1} analysis", f"Unsupported file type: {file_type}"))
+                continue
 
-        if video:
-            frames = extract_frames_from_video(video)
-            for i, frame in enumerate(frames):
-                image_data_url = f"data:image/png;base64,{encode_image(frame)}"
-                messages = [
-                    {
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "text",
-                                "text": f"{instruction}\n\nAnalyze this frame from a video (Frame {i+1}/{len(frames)}). First, explain the video and then determine if it's a construction site. If it is, explain what you observe, focusing on safety aspects. If it's not, briefly describe what you see."
-                            },
-                            {
-                                "type": "image_url",
-                                "image_url": {
-                                    "url": image_data_url
-                                }
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": f"{instruction}\n\nAnalyze this {content_type} (File {i+1}/{len(media)}). First, determine if it's a construction site. If it is, explain the {content_type} in detail, focusing on safety aspects. If it's not, briefly describe what you see."
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": image_data_url
                             }
-                        ]
-                    }
-                ]
-                completion = client.chat.completions.create(
-                    model="llama-3.2-90b-vision-preview",
-                    messages=messages,
-                    temperature=0.7,
-                    max_tokens=1000,
-                    top_p=1,
-                    stream=False,
-                    stop=None
-                )
-                result = completion.choices[0].message.content
-                results.append((f"Video frame {i+1} analysis", result))
+                        }
+                    ]
+                }
+            ]
+            completion = client.chat.completions.create(
+                model="llama-3.2-90b-vision-preview",
+                messages=messages,
+                temperature=0.7,
+                max_tokens=1000,
+                top_p=1,
+                stream=False,
+                stop=None
+            )
+            result = completion.choices[0].message.content
+            results.append((f"File {i+1} analysis", result))
 
         logger.info("Analysis completed successfully")
         return results
@@ -268,22 +248,19 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as iface:
         """
     )
     
-    # First row: Upload Image and Video side by side
+    # Combined upload for images and videos
     with gr.Row():
-        with gr.Column(scale=1):
-            image_input = gr.File(label="Upload Construction Site Images", file_count="multiple", type="filepath", elem_classes="image-container")
-        with gr.Column(scale=1):
-            video_input = gr.Video(label="Upload Construction Site Video", elem_classes="image-container")
+        media_input = gr.File(label="Upload Construction Site Images or Videos", file_count="multiple", type="filepath", elem_classes="image-container")
     
-    # Second row: Analyze Safety Hazards Button
+    # Analyze Safety Hazards Button
     with gr.Row():
         analyze_button = gr.Button("🔍 Analyze Safety Hazards", elem_classes="analyze-button")
     
-    # Third row: Chat Interface (Safety Analysis Results)
+    # Chat Interface (Safety Analysis Results)
     with gr.Row():
         chatbot = gr.Chatbot(label="Safety Analysis Results and Expert Chat", elem_classes="chatbot")
     
-    # Fourth row: Question Bar
+    # Question Bar
     with gr.Row():
         msg = gr.Textbox(
             label="Ask about safety measures or regulations",
@@ -292,15 +269,14 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as iface:
             elem_classes="chat-input"
         )
 
-    # Fifth row: Clear Chat and Download Report Buttons
+    # Clear Chat and Download Report Buttons
     with gr.Row():
         clear = gr.Button("🗑️ Clear Chat", elem_classes="clear-button")
         download_button = gr.Button("📥 Download Report", elem_classes="download-button")
 
-    # Add a file component to handle the download
+    # File component to handle the download
     report_file = gr.File(label="Download Safety Analysis Report")
 
-
     def update_chat(history, new_messages):
         history = history or []
         for title, content in new_messages:
@@ -308,8 +284,8 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as iface:
         return history
 
     analyze_button.click(
-        analyze_construction_image,
-        inputs=[image_input, video_input],
+        analyze_construction_media,
+        inputs=[media_input],
         outputs=[chatbot],
         postprocess=lambda x: update_chat(chatbot.value, x)
     )
@@ -317,7 +293,6 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as iface:
     msg.submit(chat_about_image, [msg, chatbot], [msg, chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
-
     download_button.click(
         download_report,
         inputs=[chatbot],
@@ -337,4 +312,4 @@ if __name__ == "__main__":
     except Exception as e:
         logger.error(f"Error when trying to launch the interface: {str(e)}")
         logger.error(traceback.format_exc())
-        print("Failed to launch the Gradio interface. Please check the logs for more information.")
+        print("Failed to launch the Gradio interface. Please check the logs for more information.")