from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
import gradio as gr
import torch
from datetime import datetime

# Load BLIP model and processor
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# Inference function to generate captions dynamically based on image content
def generate_captions_from_image(image):
    if image.mode != "RGB":
        image = image.convert("RGB")
    
    # Preprocess the image and generate a caption
    inputs = processor(image, return_tensors="pt").to(device, torch.float16)
    output = model.generate(**inputs, max_new_tokens=50)
    caption = processor.decode(output[0], skip_special_tokens=True)
    
    return caption

# Function to generate the daily progress report (DPR) text
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)  # Using file.name for filepath
        
        if image.mode != "RGB":
            image = image.convert("RGB")
        
        # Dynamically generate a caption based on the image
        caption = generate_captions_from_image(image)
        
        # Generate DPR section for this image with dynamic caption
        dpr_section = f"\nImage: {file.name}\nDescription: {caption}\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 dynamically detect construction activities, materials, and progress and generate a text-based Daily Progress Report (DPR).",
    allow_flagging="never"  # Optional: Disable flagging
)

iface.launch()