Spaces:
Sleeping
Sleeping
| from transformers import BlipProcessor, BlipForConditionalGeneration | |
| 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 concurrent.futures | |
| import gradio | |
| print("Gradio version:", gradio.__version__) | |
| # 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 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) | |
| def generate_captions_from_image(image): | |
| if image.mode != "RGB": | |
| image = image.convert("RGB") | |
| image = image.resize((320, 320)) | |
| 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 | |
| def save_dpr_to_pdf(dpr_text, image_paths, captions, filename): | |
| try: | |
| doc = SimpleDocTemplate(filename, pagesize=letter) | |
| styles = getSampleStyleSheet() | |
| title_style = ParagraphStyle( | |
| name='Title', | |
| fontSize=16, | |
| leading=20, | |
| alignment=1, | |
| spaceAfter=20, | |
| textColor=colors.black, | |
| fontName='Helvetica-Bold' | |
| ) | |
| body_style = ParagraphStyle( | |
| name='Body', | |
| fontSize=12, | |
| leading=14, | |
| spaceAfter=10, | |
| textColor=colors.black, | |
| fontName='Helvetica' | |
| ) | |
| flowables = [] | |
| flowables.append(Paragraph("Daily Progress Report", title_style)) | |
| for line in dpr_text.split('\n'): | |
| line = line.replace('\u2019', "'").replace('\u2018', "'") | |
| if line.strip(): | |
| flowables.append(Paragraph(line, body_style)) | |
| else: | |
| flowables.append(Spacer(1, 12)) | |
| for img_path, caption in zip(image_paths, captions): | |
| try: | |
| img = PDFImage(img_path, width=200, height=150) | |
| 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)) | |
| doc.build(flowables) | |
| return f"PDF saved successfully as {filename}", filename | |
| except Exception as e: | |
| return f"Error saving PDF: {str(e)}", None | |
| 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)}" | |
| 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() | |