app.py

import os
import base64
import gradio as gr
from PIL import Image, ImageOps
import io
import json
from groq import Groq
import logging
import cv2
import numpy as np
import traceback
from datetime import datetime
import tempfile

# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Load environment variables
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
if not GROQ_API_KEY:
    logger.error("GROQ_API_KEY is not set in environment variables")
    raise ValueError("GROQ_API_KEY is not set")

# Initialize Groq client
client = Groq(api_key=GROQ_API_KEY)

def encode_image(image):
    try:
        if isinstance(image, str):  # If image is a file path
            with open(image, "rb") as image_file:
                return base64.b64encode(image_file.read()).decode('utf-8')
        elif isinstance(image, Image.Image):  # If image is a PIL 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 resize_image(image, max_size=(800, 800)):
    """Resize image to avoid exceeding the API size limits."""
    try:
        image.thumbnail(max_size, Image.Resampling.LANCZOS)  # Use LANCZOS resampling for better quality
        return image
    except Exception as e:
        logger.error(f"Error resizing image: {str(e)}")
        raise
        
def extract_frames_from_video(video, frame_points=[0, 0.5, 1], max_size=(800, 800)):
    """Extract key frames from the video at specific time points."""
    cap = cv2.VideoCapture(video)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    duration = frame_count / fps

    frames = []
    for time_point in frame_points:
        cap.set(cv2.CAP_PROP_POS_MSEC, time_point * duration * 1000)
        ret, frame = cap.read()
        if ret:
            resized_frame = cv2.resize(frame, max_size)
            frames.append(resized_frame)
    cap.release()
    return frames

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:
        logger.info(f"Starting analysis of {len(media)} files")
        results = []

        instruction = ("You are an AI assistant specialized in analyzing images for safety issues. "
                       "Your task is first to explain what you see in the image and determine if the image shows a construction site. "
                       "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")

        for i, file in enumerate(media):
            try:
                file_path = file.name  # Get the file path
                logger.info(f"Processing file {i+1}/{len(media)}: {file_path}")
                
                if not os.path.exists(file_path):
                    logger.error(f"File does not exist: {file_path}")
                    results.append((f"File {i+1} analysis", f"Error: File does not exist: {file_path}"))
                    continue

                file_type = os.path.splitext(file_path)[1][1:].lower()
                
                if file_type in ['jpg', 'jpeg', 'png', 'gif']:
                    # Handle image
                    try:
                        image = Image.open(file_path)
                        logger.info(f"Successfully opened image: {file_path}")
                        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 (File {i+1}/{len(media)}). 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))
                        logger.info(f"Successfully analyzed image {i+1}")
                    except Exception as img_error:
                        logger.error(f"Error processing image {i+1}: {str(img_error)}")
                        results.append((f"Image {i+1} analysis", f"Error processing image: {str(img_error)}"))
                
                elif file_type in ['mp4', 'avi', 'mov', 'wmv']:
                    # Handle video
                    try:
                        frames = extract_frames_from_video(file_path)
                        logger.info(f"Extracted {len(frames)} frames from video: {file_path}")
                        for j, 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 (File {i+1}/{len(media)}, Frame {j+1}/{len(frames)}). First, 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
                                            }
                                        }
                                    ]
                                }
                            ]
                            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 {i+1}, Frame {j+1} analysis", result))
                        logger.info(f"Successfully analyzed video {i+1}")
                    except Exception as vid_error:
                        logger.error(f"Error processing video {i+1}: {str(vid_error)}")
                        results.append((f"Video {i+1} analysis", f"Error processing video: {str(vid_error)}"))
                
                else:
                    logger.warning(f"Unsupported file type: {file_type}")
                    results.append((f"File {i+1} analysis", f"Unsupported file type: {file_type}"))
            
            except Exception as file_error:
                logger.error(f"Error processing file {i+1}: {str(file_error)}")
                results.append((f"File {i+1} analysis", f"Error processing file: {str(file_error)}"))

        logger.info("Analysis completed successfully")
        return results
    except Exception as e:
        logger.error(f"Error during analysis: {str(e)}")
        logger.error(traceback.format_exc())
        return [("Analysis error", f"Error during analysis: {str(e)}")]

        
def chat_about_image(message, chat_history):
    try:
        # Prepare the conversation history for the API
        messages = [
            {"role": "system", "content": "You are an AI assistant specialized in analyzing construction site images and answering questions about them. Use the information from the initial analysis to answer user queries."},
        ]
        
        # Add chat history to messages
        for human, ai in chat_history:
            if human:
                messages.append({"role": "user", "content": human})
            if ai:
                messages.append({"role": "assistant", "content": ai})
        
        # Add the new user message
        messages.append({"role": "user", "content": message})
        
        # Make API call
        completion = client.chat.completions.create(
            model="llama-3.2-90b-vision-preview",
            messages=messages,
            temperature=0.7,
            max_tokens=500,
            top_p=1,
            stream=False,
            stop=None
        )
        
        response = completion.choices[0].message.content
        chat_history.append((message, response))
        
        return "", chat_history
    except Exception as e:
        logger.error(f"Error during chat: {str(e)}")
        return "", chat_history + [(message, f"Error: {str(e)}")]

def generate_summary_report(chat_history):
    """
    Generate a summary report from the chat history.
    """
    report = "Construction Site Safety Analysis Report\n"
    report += "=" * 40 + "\n"
    report += f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"

    for i, (user, ai) in enumerate(chat_history, 1):
        if user:
            report += f"Query {i}:\n{user}\n\n"
        if ai:
            report += f"Analysis {i}:\n{ai}\n\n"
        report += "-" * 40 + "\n"

    return report

def download_report(chat_history):
    """
    Generate and provide a download link for the summary report.
    """
    report = generate_summary_report(chat_history)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"safety_analysis_report_{timestamp}.txt"
    
    # Create a temporary file
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as temp_file:
        temp_file.write(report)
        temp_file_path = temp_file.name
    
    return temp_file_path


# Custom CSS for improved styling
custom_css = """
.container { max-width: 1200px; margin: auto; padding-top: 1.5rem; }
.header { text-align: center; margin-bottom: 1rem; }
.header h1 { color: #2c3e50; font-size: 2.5rem; }
.subheader { 
    color: #34495e; 
    font-size: 1rem; 
    line-height: 1.2; 
    margin-bottom: 1.5rem; 
    text-align: center; 
    padding: 0 15px;
    white-space: nowrap;
    overflow: hidden;
    text-overflow: ellipsis;
}
.image-container { border: 2px dashed #3498db; border-radius: 10px; padding: 1rem; text-align: center; margin-bottom: 1rem; }
.analyze-button { background-color: #2ecc71 !important; color: white !important; width: 100%; }
.clear-button { background-color: #e74c3c !important; color: white !important; width: 100px !important; }
.chatbot { border: 1px solid #bdc3c7; border-radius: 10px; padding: 1rem; height: 500px; overflow-y: auto; }
.chat-input { border: 1px solid #bdc3c7; border-radius: 5px; padding: 0.5rem; width: 100%; }
.groq-badge { position: fixed; bottom: 10px; right: 10px; background-color: #f39c12; color: white; padding: 5px 10px; border-radius: 5px; font-weight: bold; }
.chat-container { display: flex; flex-direction: column; height: 100%; }
.input-row { display: flex; align-items: center; margin-top: 10px; justify-content: space-between; }
.input-row > div:first-child { flex-grow: 1; margin-right: 10px; }
"""

# Create the Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as iface:
    gr.HTML(
        """
        <div class="container">
            <div class="header">
                <h1>🏗️ Construction Site Safety Analyzer</h1>
            </div>
            <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>
        """
    )
    
    # Combined upload for images and videos
    with gr.Row():
        media_input = gr.File(label="Upload Construction Site Images or Videos", file_count="multiple", type="filepath", elem_classes="image-container")
    
    # Analyze Safety Hazards Button
    with gr.Row():
        analyze_button = gr.Button("🔍 Analyze Safety Hazards", elem_classes="analyze-button")
    
    # Chat Interface (Safety Analysis Results)
    with gr.Row():
        chatbot = gr.Chatbot(label="Safety Analysis Results and Expert Chat", elem_classes="chatbot")
    
    # Question Bar
    with gr.Row():
        msg = gr.Textbox(
            label="Ask about safety measures or regulations",
            placeholder="E.g., 'What OSHA guidelines apply to this hazard?'",
            show_label=False,
            elem_classes="chat-input"
        )

    # 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")

    # 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:
            history.append((None, f"{title}\n\n{content}"))
        return history

    analyze_button.click(
        analyze_construction_media,
        inputs=[media_input],
        outputs=[chatbot],
        postprocess=lambda x: update_chat(chatbot.value, x)
    )

    msg.submit(chat_about_image, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

    download_button.click(
        download_report,
        inputs=[chatbot],
        outputs=[report_file]
    )

    gr.HTML(
        """
        <div class="groq-badge">Powered by Groq</div>
        """
    )

# Launch the app
if __name__ == "__main__":
    try:
        iface.launch(debug=True)
    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.")