Update app.py

app.py

@@ -1,17 +1,18 @@
-from fastapi import FastAPI, File, UploadFile
-import requests
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
-import io
-import base64
-import os
-from dotenv import load_dotenv
-from simple_salesforce import Salesforce
+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 datetime import datetime
+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()
@@ -28,45 +29,21 @@ except Exception as e:
     sf = None
     print(f"Failed to connect to Salesforce: {str(e)}")
 
-# Initialize Hugging Face model and processor
+# 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)
 
-# Define FastAPI instance
-app = FastAPI()
-
-# Endpoint for image upload and caption generation
-@app.post("/predict/")
-async def predict(file: UploadFile = File(...)):
-    image = Image.open(file.file)
-    
-    # Generate caption using Hugging Face model
-    caption = generate_captions_from_image(image)
-
-    # Save the image to a file
-    file_path = f"./uploaded_images/{file.filename}"
-    image.save(file_path)
-    
-    # Save the daily report as a PDF
-    pdf_filename = save_dpr_to_pdf(caption, file_path)
-    
-    # Upload to Salesforce
-    if sf:
-        salesforce_result = upload_file_to_salesforce(pdf_filename, sf)
-    else:
-        salesforce_result = "Salesforce connection is not available."
-    
-    return {"caption": caption, "pdf_filename": pdf_filename, "salesforce_result": salesforce_result}
-
-
-# Function to generate captions from an image
+# Inference function to generate captions dynamically based on image content
 def generate_captions_from_image(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
     
+    # Resize image for faster processing
+    image = image.resize((640, 640))
+    
     # 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)
@@ -74,14 +51,11 @@ def generate_captions_from_image(image):
     
     return caption
 
-# Function to save Daily Progress Report to a PDF file
-def save_dpr_to_pdf(caption, image_path):
+# Function to save DPR text to a PDF file
+def save_dpr_to_pdf(dpr_text, image_paths, captions, filename):
     try:
-        # PDF filename
-        pdf_filename = f"DPR_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.pdf"
-        
         # Create a PDF document
-        doc = SimpleDocTemplate(pdf_filename, pagesize=letter)
+        doc = SimpleDocTemplate(filename, pagesize=letter)
         styles = getSampleStyleSheet()
         
         # Define custom styles
@@ -109,33 +83,51 @@ def save_dpr_to_pdf(caption, image_path):
         # Add title
         flowables.append(Paragraph("Daily Progress Report", title_style))
         
-        # Add image and caption
-        flowables.append(PDFImage(image_path, width=200, height=150))  # Adjust image size
-        flowables.append(Paragraph(f"Description: {caption}", body_style))
-        flowables.append(Spacer(1, 12))  # Add some space between images
+        # 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 (no need to add description to dpr_text again)
+        for img_path, caption in zip(image_paths, captions):
+            try:
+                # Add image first
+                img = PDFImage(img_path, width=200, height=150)  # Adjust image size if needed
+                flowables.append(img)
+                # Add description below the image
+                description = f"Description: {caption}"
+                flowables.append(Paragraph(description, body_style))
+                flowables.append(Spacer(1, 12))  # Add some space between images
+            except Exception as e:
+                flowables.append(Paragraph(f"Error loading image: {str(e)}", body_style))
         
         # Build the PDF
         doc.build(flowables)
-        
-        return pdf_filename
+        return f"PDF saved successfully as {filename}", filename
     except Exception as e:
-        print(f"Error saving PDF: {str(e)}")
-        return None
+        return f"Error saving PDF: {str(e)}", None
 
-# Function to upload the file to Salesforce
-def upload_file_to_salesforce(pdf_filename, sf_connection):
+# Function to upload a file to Salesforce as ContentVersion
+def upload_file_to_salesforce(file_path, filename, sf_connection, file_type):
     try:
         # Read file content and encode in base64
-        with open(pdf_filename, 'rb') as f:
+        with open(file_path, 'rb') as f:
             file_content = f.read()
         file_content_b64 = base64.b64encode(file_content).decode('utf-8')
         
+        # Set description based on file type
+        description = "Daily Progress Report PDF" if file_type == "pdf" else "Site Image"
+        
         # Create ContentVersion
         content_version = sf_connection.ContentVersion.create({
-            'Title': pdf_filename,
-            'PathOnClient': pdf_filename,
+            'Title': filename,
+            'PathOnClient': filename,
             'VersionData': file_content_b64,
-            'Description': "Daily Progress Report PDF"
+            'Description': description
         })
         
         # Get ContentDocumentId
@@ -145,14 +137,134 @@ def upload_file_to_salesforce(pdf_filename, sf_connection):
         )
         content_document_id = content_document['records'][0]['ContentDocumentId']
         
-        # Create the record link in Salesforce
-        return f"File {pdf_filename} uploaded successfully with ContentDocumentId: {content_document_id}"
+        # Generate a valid Salesforce URL for the ContentDocument
+        content_document_url = f"https://{sf_connection.sf_instance}.salesforce.com/{content_document_id}"
+
+        # Ensure the link is valid
+        return content_document_id, content_document_url, f"File {filename} uploaded successfully"
     except Exception as e:
-        print(f"Error uploading to Salesforce: {str(e)}")
-        return f"Error uploading {pdf_filename} to Salesforce: {str(e)}"
+        return None, None, f"Error uploading {filename} to Salesforce: {str(e)}"
+
+# Function to generate the daily progress report (DPR), save as PDF, and upload to Salesforce
+def generate_dpr(files):
+    dpr_text = []
+    captions = []
+    image_paths = []
+    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 images in parallel for faster performance
+    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)
+        
+        # Generate DPR section for this image with dynamic caption
+        dpr_section = f"\nImage: {file.name}\nDescription: {caption}\n"
+        # Remove the description from the dpr_text section
+        # No need to add it again as the image and caption will be inserted in the PDF
+        dpr_text.append(dpr_section)
+        
+        # Save image path for embedding in the report
+        image_paths.append(file.name)
+    
+    # Combine DPR text (no redundant description here)
+    dpr_output = "\n".join(dpr_text)
+    
+    # Generate PDF filename with timestamp
+    pdf_filename = f"DPR_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.pdf"
+    
+    # Save DPR text to PDF
+    pdf_result, pdf_filepath = save_dpr_to_pdf(dpr_output, image_paths, captions, pdf_filename)
+    
+    # Salesforce upload
+    salesforce_result = ""
+    pdf_content_document_id = None
+    pdf_url = None
+    image_content_document_ids = []
+    
+    if sf and pdf_filepath:
+        try:
+            # Create Daily_Progress_Reports__c record
+            report_description = "; ".join(captions)[:255]  # Concatenate captions, limit to 255 chars
+            dpr_record = sf.Daily_Progress_Reports__c.create({
+                'Detected_Activities__c': report_description  # Store in Detected_Activities__c field
+            })
+            dpr_record_id = dpr_record['id']
+            salesforce_result += f"Created Daily_Progress_Reports__c record with ID: {dpr_record_id}\n"
+            
+            # Upload PDF to Salesforce
+            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"
+            
+            # Link PDF to DPR record
+            if pdf_content_document_id:
+                sf.ContentDocumentLink.create({
+                    'ContentDocumentId': pdf_content_document_id,
+                    'LinkedEntityId': dpr_record_id,
+                    'ShareType': 'V'
+                })
+            
+            # Update the DPR record with the PDF URL
+            if pdf_url:
+                sf.Daily_Progress_Reports__c.update(dpr_record_id, {
+                    'PDF_URL__c': pdf_url  # Storing the PDF URL correctly
+                })
+                salesforce_result += f"Updated PDF URL for record ID {dpr_record_id}\n"
+            
+            # Upload images to Salesforce and create Site_Images__c records
+            for file in files:
+                image_filename = os.path.basename(file.name)
+                image_content_document_id, image_upload_result = upload_file_to_salesforce(
+                    file.name, image_filename, sf, "image"
+                )
+                if image_content_document_id:
+                    image_content_document_ids.append(image_content_document_id)
+                
+                # Create Site_Images__c record and link to DPR
+                site_image_record = sf.Site_Images__c.create({
+                    'Image__c': image_content_document_id,
+                    'Related_Report__c': dpr_record_id  # Link image to DPR record
+                })
+                salesforce_result += image_upload_result + "\n"
+                
+                # Link image to DPR record
+                if image_content_document_id:
+                    sf.ContentDocumentLink.create({
+                        'ContentDocumentId': image_content_document_id,
+                        'LinkedEntityId': dpr_record_id,
+                        'ShareType': 'V'
+                    })
+            
+        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 text, PDF file, and Salesforce upload status
+    return (
+        dpr_output + f"\n\n{pdf_result}\n\nSalesforce Upload Status:\n{salesforce_result}",
+        pdf_filepath
+    )
 
+# Gradio interface for uploading multiple files, displaying DPR, and downloading PDF
+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"
+)
 
-# To run the app
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    iface.launch()