# app.py

import gradio as gr
from PIL import Image
import torch
from transformers import pipeline
import os

# Load YOLOv5 model from Ultralytics
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# Load GPT-2 text generation pipeline
text_generator = pipeline('text-generation', model='gpt2')

def generate_report(image):
    # Run YOLOv5 detection
    results = model(image)
    labels = results.pandas().xyxy[0]['name'].tolist()

    if not labels:
        return "No relevant objects detected in the image. Please try a clearer construction site photo."

    # Prepare prompt for report generation
    prompt = (
        f"The following objects were detected at the construction site: {', '.join(labels)}.\n"
        "Based on this, write a brief site progress report or highlight any visible safety concerns."
    )

    # Generate report using GPT-2
    report = text_generator(prompt, max_length=150, do_sample=True, temperature=0.7)[0]['generated_text']

    return report

# Gradio interface
interface = gr.Interface(
    fn=generate_report,
    inputs=gr.Image(type="pil", label="Upload Construction Site Photo"),
    outputs=gr.Textbox(label="AI Generated Report"),
    title="📸 Photo to Construction Report Generator",
    description="Upload a construction site photo and get an AI-generated progress report or safety checklist based on detected objects.",
    theme="default"
)

if __name__ == "__main__":
    interface.launch()