Update app.py

app.py

@@ -1,10 +1,9 @@
 from fastapi import FastAPI, File, UploadFile
 import requests
-
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
-import gradio as gr
 import torch
+import gradio as gr
 from datetime import datetime
 from reportlab.lib.pagesizes import letter
 from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as PDFImage
@@ -22,7 +21,6 @@ load_dotenv()
 
 app = FastAPI()
 
-
 # Salesforce credentials
 SF_USERNAME = os.getenv('SF_USERNAME')
 SF_PASSWORD = os.getenv('SF_PASSWORD')
@@ -42,35 +40,54 @@ model.eval()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
+# 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'
 
-# FastAPI endpoint to handle image upload and caption generation
 @app.post("/predict/")
 async def predict(image: UploadFile = File(...)):
     try:
         # Read the image from the request
         image_bytes = await image.read()
-        image = Image.open(BytesIO(image_bytes))
-        
-        # Generate caption from the image
-        caption = generate_captions_from_image(image)
-        return {"caption": caption}
+        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}"}
     except Exception as e:
         return {"error": str(e)}
 
-
-# Inference function to generate captions dynamically based on image content
-def generate_captions_from_image(image):
+# 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))
     
-    # Preprocess the image and generate a caption
+    # 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()
+
+    # 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
+
+# Inference function to generate captions dynamically based on image content
+def generate_captions_from_image(image):
     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
 
 # Function to save DPR text to a PDF file
@@ -133,148 +150,6 @@ def save_dpr_to_pdf(dpr_text, image_paths, captions, filename):
     except Exception as e:
         return f"Error saving PDF: {str(e)}", None
 
-# 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}"
-
-        # 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):
-    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,
@@ -289,6 +164,5 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
-    
+        iface.launch()
+