Update app.py

app.py

@@ -6,6 +6,8 @@ import io
 import json
 from groq import Groq
 import logging
+import cv2
+import numpy as np
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG)
@@ -29,68 +31,106 @@ def encode_image(image):
             buffered = io.BytesIO()
             image.save(buffered, format="PNG")
             return base64.b64encode(buffered.getvalue()).decode('utf-8')
+        elif isinstance(image, np.ndarray):  # If image is a numpy array (from video)
+            is_success, buffer = cv2.imencode(".png", image)
+            if is_success:
+                return base64.b64encode(buffer).decode('utf-8')
         else:
             raise ValueError(f"Unsupported image type: {type(image)}")
     except Exception as e:
         logger.error(f"Error encoding image: {str(e)}")
         raise
 
-def analyze_construction_image(image):
-    if image is None:
-        logger.warning("No image provided")
-        return [("No image uploaded", "Error: Please upload an image for analysis.")]
+def analyze_construction_image(images, video=None):
+    if not images and video is None:
+        logger.warning("No images or video provided")
+        return [("No input", "Error: Please upload images or a video for analysis.")]
 
     try:
-        logger.info("Starting image analysis")
-        image_data_url = f"data:image/png;base64,{encode_image(image)}"
-        logger.debug("Image encoded successfully")
+        logger.info("Starting analysis")
+        results = []
 
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": "Analyze this construction site image. Identify any safety issues or hazards, categorize them, provide a detailed description, and suggest steps to resolve them."
-                    },
+        if images:
+            for i, image in enumerate(images):
+                image_data_url = f"data:image/png;base64,{encode_image(image)}"
+                messages = [
                     {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": image_data_url
-                        }
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": f"Analyze this construction site image (Image {i+1}/{len(images)}). Identify any safety issues or hazards, categorize them, provide a detailed description, and suggest steps to resolve them."
+                            },
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": image_data_url
+                                }
+                            }
+                        ]
                     }
                 ]
-            }
-        ]
-
-        logger.info("Sending request to Groq API")
-        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
-        )
-
-        logger.info("Received response from Groq API")
-        result = completion.choices[0].message.content
-        logger.debug(f"Raw API response: {result}")
-
-        if not result:
-            logger.warning("Received empty response from API")
-            return [("Image analysis request", "Error: Received empty response from API")]
+                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))
+
+        if video:
+            cap = cv2.VideoCapture(video.name)
+            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            fps = int(cap.get(cv2.CAP_PROP_FPS))
+            duration = frame_count / fps
+            
+            # Analyze frames at 0%, 25%, 50%, 75%, and 100% of the video duration
+            for i, time_point in enumerate([0, 0.25, 0.5, 0.75, 1]):
+                cap.set(cv2.CAP_PROP_POS_MSEC, time_point * duration * 1000)
+                ret, frame = cap.read()
+                if ret:
+                    image_data_url = f"data:image/png;base64,{encode_image(frame)}"
+                    messages = [
+                        {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": f"Analyze this frame from a construction site video (Frame {i+1}/5 at {time_point*100}% of video duration). Identify any safety issues or hazards, categorize them, provide a detailed description, and suggest steps to resolve them."
+                                },
+                                {
+                                    "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))
+            cap.release()
 
         logger.info("Analysis completed successfully")
-        
-        # Return the result directly, without any parsing
-        return [("Image analysis request", result)]
+        return results
     except Exception as e:
-        logger.error(f"Error during image analysis: {str(e)}")
+        logger.error(f"Error during analysis: {str(e)}")
         logger.error(traceback.format_exc())
         error_message = f"Error during analysis: {str(e)}. Please try again or contact support if the issue persists."
-        return [("Image analysis request", error_message)]
+        return [("Analysis error", error_message)]
         
 def chat_about_image(message, chat_history):
     try:
@@ -164,14 +204,15 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as iface:
             <div class="header">
                 <h1>🏗️ Construction Site Safety Analyzer</h1>
             </div>
-            <p class="subheader">Enhance workplace safety and compliance with AI-powered image analysis using Llama 3.2 90B Vision and expert chat assistance.</p>
+            <p class="subheader">Enhance workplace safety and compliance with AI-powered image and video analysis using Llama 3.2 90B Vision and expert chat assistance.</p>
         </div>
         """
     )
     
     with gr.Row():
         with gr.Column(scale=1):
-            image_input = gr.Image(type="pil", label="Upload Construction Site Image", elem_classes="image-container")
+            image_input = gr.File(label="Upload Construction Site Images", file_count="multiple", type="file", elem_classes="image-container")
+            video_input = gr.Video(label="Upload Construction Site Video", elem_classes="image-container")
             analyze_button = gr.Button("🔍 Analyze Safety Hazards", elem_classes="analyze-button")
         with gr.Column(scale=2):
             with gr.Group(elem_classes="chat-container"):
@@ -185,16 +226,16 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as iface:
                     )
                     clear = gr.Button("🗑️ Clear", elem_classes="clear-button")
 
-    def update_chat(history, new_message):
+    def update_chat(history, new_messages):
         history = history or []
-        history.append(new_message)
+        history.extend(new_messages)
         return history
 
     analyze_button.click(
         analyze_construction_image,
-        inputs=[image_input],
+        inputs=[image_input, video_input],
         outputs=[chatbot],
-        postprocess=lambda x: update_chat(chatbot.value, x[0])
+        postprocess=lambda x: update_chat(chatbot.value, x)
     )
 
     msg.submit(chat_about_image, [msg, chatbot], [msg, chatbot])