from PIL import Image
import gradio as gr
import torch
from datetime import datetime
from ultralytics import YOLO

torch.serialization.add_safe_globals([torch.nn.Module, 'ultralytics.nn.tasks.DetectionModel'])

# Load YOLOv8 model (trained on construction dataset)
model = YOLO('yolov8n.pt')  # Path to pre-trained model on construction dataset

# Function to generate DPR text based on detections
def generate_dpr(files):
    dpr_text = []
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # Add header to the DPR
    dpr_text.append(f"Daily Progress Report\nGenerated on: {current_time}\n")
    
    # Process each uploaded file (image)
    for file in files:
        # Open the image from file path
        image = Image.open(file.name)
        
        # Perform object detection with YOLOv8
        results = model(image)  # Perform detection
        
        # Parse detections (activities, materials, etc.)
        detected_objects = results.names  # Object names detected by the model
        detections = results.pandas().xywh  # Get the dataframe with detection results
        
        detected_activities = []
        detected_materials = []
        
        # Define construction activity and material categories
        construction_activities = ['scaffolding', 'concrete pouring', 'welding', 'excavation']
        construction_materials = ['concrete', 'steel', 'bricks', 'cement', 'sand']
        
        # Check the detected objects and categorize them
        for obj in detected_objects:
            if obj.lower() in construction_activities:
                detected_activities.append(obj)
            elif obj.lower() in construction_materials:
                detected_materials.append(obj)
        
        # Build a detailed report for this image
        dpr_section = f"\nImage: {file.name}\n"
        
        if detected_activities:
            dpr_section += f"Detected Activities: {', '.join(detected_activities)}\n"
        else:
            dpr_section += "No construction activities detected.\n"
        
        if detected_materials:
            dpr_section += f"Detected Materials: {', '.join(detected_materials)}\n"
        else:
            dpr_section += "No materials detected.\n"
        
        dpr_text.append(dpr_section)
    
    # Return the generated DPR as a text output
    return "\n".join(dpr_text)

# Gradio interface for uploading multiple files and displaying the text-based DPR
iface = gr.Interface(
    fn=generate_dpr,
    inputs=gr.Files(type="filepath", label="Upload Site Photos"),  # Handle batch upload of images
    outputs="text",  # Display the DPR as text in the output section
    title="Daily Progress Report Generator",
    description="Upload up to 10 site photos. The AI model will detect construction activities, materials, and progress and generate a text-based Daily Progress Report (DPR).",
    allow_flagging="never"  # Optional: Disable flagging
)

iface.launch()