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()