Update app.py

app.py

@@ -1,9 +1,7 @@
-from fastapi import FastAPI, File, UploadFile
-import requests
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
-import torch
 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
@@ -19,8 +17,6 @@ import concurrent.futures
 # Load environment variables from .env file
 load_dotenv()
 
-app = FastAPI()
-
 # Salesforce credentials
 SF_USERNAME = os.getenv('SF_USERNAME')
 SF_PASSWORD = os.getenv('SF_PASSWORD')
@@ -40,7 +36,7 @@ model.eval()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-# Function to generate captions dynamically based on image content
+# Inference function to generate captions dynamically based on image content
 def generate_captions_from_image(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
@@ -55,48 +51,99 @@ def generate_captions_from_image(image):
     
     return caption
 
-# FastAPI endpoint to handle image upload and forward it to Hugging Face API for caption generation
-HUGGING_FACE_ENDPOINT = 'https://huggingface.co/spaces/Rammohan0504/DPR-4/predict'
-
-@app.post("/predict/")
-async def predict(image: UploadFile = File(...)):
+# Function to save DPR text to a PDF file
+def save_dpr_to_pdf(dpr_text, image_paths, captions, filename):
     try:
-        # Read the image from the request
-        image_bytes = await image.read()
-        image = Image.open(io.BytesIO(image_bytes))
-
-        # Forward the image to Hugging Face endpoint
-        response = forward_image_to_huggingface(image)
-
-        # Check the response from Hugging Face
-        if response.status_code == 200:
-            result = response.json()
-            caption = result.get("caption", "No caption found.")
-            return {"caption": caption}
-        else:
-            return {"error": f"Failed to get prediction from Hugging Face Space. Status code: {response.status_code}"}
+        # 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 (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 f"PDF saved successfully as {filename}", filename
     except Exception as e:
-        return {"error": str(e)}
+        return f"Error saving PDF: {str(e)}", None
 
-# Function to forward the image to Hugging Face API
-def forward_image_to_huggingface(image: Image):
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    
-    # Resize image for faster processing
-    image = image.resize((640, 640))
-    
-    # Convert image to bytes for API request
-    img_byte_arr = io.BytesIO()
-    image.save(img_byte_arr, format='JPEG')
-    img_byte_arr = img_byte_arr.getvalue()
+# 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(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': filename,
+            'PathOnClient': filename,
+            'VersionData': file_content_b64,
+            'Description': description
+        })
+        
+        # Get ContentDocumentId
+        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']
+        
+        # Generate a valid Salesforce URL for the ContentDocument
+        content_document_url = f"https://{sf_connection.sf_instance}.salesforce.com/{content_document_id}"
 
-    # Create the payload to send to Hugging Face (it expects a file)
-    files = {'file': ('image.jpg', img_byte_arr, 'image/jpeg')}
-    
-    # Make the POST request to Hugging Face Space
-    response = requests.post(HUGGING_FACE_ENDPOINT, files=files)
-    return response
+        # Ensure the link is valid
+        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)}"
 
 # Function to generate the daily progress report (DPR), save as PDF, and upload to Salesforce
 def generate_dpr(files):
@@ -118,6 +165,8 @@ def generate_dpr(files):
         
         # 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
@@ -204,66 +253,6 @@ def generate_dpr(files):
         pdf_filepath
     )
 
-# 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 (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 f"PDF saved successfully as {filename}", filename
-    except Exception as e:
-        return f"Error saving PDF: {str(e)}", None
-
 # Gradio interface for uploading multiple files, displaying DPR, and downloading PDF
 iface = gr.Interface(
     fn=generate_dpr,
@@ -278,5 +267,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-        iface.launch()
-
+    iface.launch()