app.py

from transformers import AutoProcessor, AutoModelForImageClassification
from PIL import Image
import gradio as gr
import torch
from datetime import datetime
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as PDFImage
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib import colors
from simple_salesforce import Salesforce
import os
from dotenv import load_dotenv
import base64
import io
import concurrent.futures

# Load environment variables from .env file
load_dotenv()

# Salesforce credentials
SF_USERNAME = os.getenv('SF_USERNAME')
SF_PASSWORD = os.getenv('SF_PASSWORD')
SF_SECURITY_TOKEN = os.getenv('SF_SECURITY_TOKEN')

# Initialize Salesforce connection
try:
    sf = Salesforce(username=SF_USERNAME, password=SF_PASSWORD, security_token=SF_SECURITY_TOKEN)
except Exception as e:
    sf = None
    print(f"Failed to connect to Salesforce: {str(e)}")

# Load ViT model and processor (generic ImageNet pretrained)
processor = AutoProcessor.from_pretrained("google/vit-base-patch16-224")
model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224")
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# Inference function to classify image and get predicted label
def generate_captions_from_image(image):
    if image.mode != "RGB":
        image = image.convert("RGB")
    inputs = processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    predicted_label = model.config.id2label[predicted_class_idx]
    return predicted_label

# Function to save DPR text to a PDF file
def save_dpr_to_pdf(dpr_text, image_paths, captions, filename):
    try:
        # Create a PDF document
        doc = SimpleDocTemplate(filename, pagesize=letter)
        styles = getSampleStyleSheet()
        
        # Define custom styles
        title_style = ParagraphStyle(
            name='Title',
            fontSize=16,
            leading=20,
            alignment=1,  # Center
            spaceAfter=20,
            textColor=colors.black,
            fontName='Helvetica-Bold'
        )
        body_style = ParagraphStyle(
            name='Body',
            fontSize=12,
            leading=14,
            spaceAfter=10,
            textColor=colors.black,
            fontName='Helvetica'
        )
        
        # Build the PDF content
        flowables = []
        
        # Add title
        flowables.append(Paragraph("Daily Progress Report", title_style))
        
        # Split DPR text into lines and add as paragraphs (excluding descriptions for images)
        for line in dpr_text.split('\n'):
            # Replace problematic characters for PDF
            line = line.replace('\u2019', "'").replace('\u2018', "'")
            if line.strip():
                flowables.append(Paragraph(line, body_style))
            else:
                flowables.append(Spacer(1, 12))
        
        # Add images and captions in the correct order
        for img_path, caption in zip(image_paths, captions):
            try:
                img = PDFImage(img_path, width=200, height=150)  # Adjust image size if needed
                flowables.append(img)
                description = f"Description: {caption}"
                flowables.append(Paragraph(description, body_style))
                flowables.append(Spacer(1, 12))
            except Exception as e:
                flowables.append(Paragraph(f"Error loading image: {str(e)}", body_style))
        
        # Build the PDF
        doc.build(flowables)
        return f"PDF saved successfully as {filename}", filename
    except Exception as e:
        return f"Error saving PDF: {str(e)}", None

# Function to upload file to Salesforce as ContentVersion
def upload_file_to_salesforce(file_path, filename, sf_connection, file_type):
    try:
        with open(file_path, 'rb') as f:
            file_content = f.read()
        file_content_b64 = base64.b64encode(file_content).decode('utf-8')
        
        description = "Daily Progress Report PDF" if file_type == "pdf" else "Site Image"
        
        content_version = sf_connection.ContentVersion.create({
            'Title': filename,
            'PathOnClient': filename,
            'VersionData': file_content_b64,
            'Description': description
        })
        
        content_version_id = content_version['id']
        content_document = sf_connection.query(
            f"SELECT ContentDocumentId FROM ContentVersion WHERE Id = '{content_version_id}'"
        )
        content_document_id = content_document['records'][0]['ContentDocumentId']
        
        content_document_url = f"https://{sf_connection.sf_instance}/sfc/servlet.shepherd/version/download/{content_version_id}"
    
        return content_document_id, content_document_url, f"File {filename} uploaded successfully"
    except Exception as e:
        return None, None, f"Error uploading {filename} to Salesforce: {str(e)}"

# Generate DPR, save PDF, upload to Salesforce
def generate_dpr(files):
    dpr_text = []
    captions = []
    image_paths = []
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    dpr_text.append(f"Daily Progress Report\nGenerated on: {current_time}\n")
    
    with concurrent.futures.ThreadPoolExecutor() as executor:
        results = list(executor.map(lambda file: generate_captions_from_image(Image.open(file.name)), files))
    
    for i, file in enumerate(files):
        caption = results[i]
        captions.append(caption)
        dpr_section = f"\nImage: {file.name}\nDescription: {caption}\n"
        dpr_text.append(dpr_section)
        image_paths.append(file.name)
    
    dpr_output = "\n".join(dpr_text)
    pdf_filename = f"DPR_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.pdf"
    
    pdf_result, pdf_filepath = save_dpr_to_pdf(dpr_output, image_paths, captions, pdf_filename)
    
    salesforce_result = ""
    pdf_content_document_id = None
    pdf_url = None
    
    if sf and pdf_filepath:
        try:
            report_description = "; ".join(captions)[:255]
            dpr_record = sf.Daily_Progress_Reports__c.create({
                'Detected_Activities__c': report_description
            })
            dpr_record_id = dpr_record['id']
            salesforce_result += f"Created Daily_Progress_Reports__c record with ID: {dpr_record_id}\n"
            
            pdf_content_document_id, pdf_url, pdf_upload_result = upload_file_to_salesforce(
                pdf_filepath, pdf_filename, sf, "pdf"
            )
            salesforce_result += pdf_upload_result + "\n"
            
            if pdf_content_document_id:
                sf.ContentDocumentLink.create({
                    'ContentDocumentId': pdf_content_document_id,
                    'LinkedEntityId': dpr_record_id,
                    'ShareType': 'V'
                })
            
            if pdf_url:
                sf.Daily_Progress_Reports__c.update(dpr_record_id, {
                    'PDF_URL__c': pdf_url
                })
                salesforce_result += f"Updated PDF URL for record ID {dpr_record_id}\n"
            
            for file in files:
                image_filename = os.path.basename(file.name)
                image_content_document_id, image_url, image_upload_result = upload_file_to_salesforce(
                    file.name, image_filename, sf, "image"
                )
                if image_content_document_id:
                    sf.ContentDocumentLink.create({
                        'ContentDocumentId': image_content_document_id,
                        'LinkedEntityId': dpr_record_id,
                        'ShareType': 'V'
                    })
                    sf.Daily_Progress_Reports__c.update(dpr_record_id, {
                        'Site_Images__c': image_content_document_id
                    })
                    salesforce_result += image_upload_result + "\n"
                        
        except Exception as e:
            salesforce_result += f"Error interacting with Salesforce: {str(e)}\n"
    else:
        salesforce_result = "Salesforce connection not available or PDF generation failed.\n"
    
    return (
        dpr_output + f"\n\n{pdf_result}\n\nSalesforce Upload Status:\n{salesforce_result}",
        pdf_filepath
    )

iface = gr.Interface(
    fn=generate_dpr,
    inputs=gr.Files(type="filepath", label="Upload Site Photos"),
    outputs=[
        gr.Textbox(label="Daily Progress Report"),
        gr.File(label="Download PDF")
    ],
    title="Daily Progress Report Generator",
    description="Upload up to 10 site photos. The AI model will generate a text-based Daily Progress Report (DPR), save it as a PDF, and upload the PDF and images to Salesforce under Daily_Progress_Reports__c in the Files related list. Download the PDF locally if needed.",
    allow_flagging="never"
)

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