Update app.py

app.py

@@ -4,7 +4,6 @@ import numpy as np
 from fpdf import FPDF
 import tempfile
 import os
-import shutil
 from paddleocr import PaddleOCR
 import time
 from simple_salesforce import Salesforce
@@ -88,69 +87,66 @@ def create_pdf_report(coverage_percent, extracted_texts, annotated_image_path, o
 
     pdf.output(output_path)
 
-def upload_image_and_get_url(image_path):
+# New function to upload image to Salesforce and get URL (adapted from reference code)
+def upload_image_to_salesforce(image_path, image_name, record_id=None):
     """
-    Upload the image to the uploads directory in Hugging Face Space and return a public URL.
-    Use /tmp as a temporary storage location, then copy to the persistent uploads folder.
+    Upload the image to Salesforce as a ContentVersion and return a public URL.
     """
-    import shutil
-    import os
-    from datetime import datetime
-    import uuid
-
-    # Step 1: Write to /tmp (guaranteed writable in Hugging Face Spaces)
-    temp_dir = "/tmp/uploads"
     try:
-        os.makedirs(temp_dir, exist_ok=True)
-        logger.info(f"Created temporary directory at: {temp_dir}")
+        sf = Salesforce(
+            username=os.environ['SF_USERNAME'],
+            password=os.environ['SF_PASSWORD'],
+            security_token=os.environ['SF_SECURITY_TOKEN'],
+            domain=os.environ.get('SF_DOMAIN', 'login')
+        )
+        logger.debug(f"Uploading image {image_name} for record ID: {record_id}")
+
+        # Read the image file and encode it as base64
+        with open(image_path, "rb") as f:
+            image_data = f.read()
+
+        encoded_image_data = base64.b64encode(image_data).decode('utf-8')
+
+        # Create a ContentVersion in Salesforce
+        content_version_data = {
+            "Title": image_name,
+            "PathOnClient": image_name,
+            "VersionData": encoded_image_data,
+        }
+        
+        if record_id:
+            content_version_data["FirstPublishLocationId"] = record_id
+
+        content_version = sf.ContentVersion.create(content_version_data)
+        content_version_id = content_version["id"]
+        logger.info(f"Image uploaded to Salesforce with ContentVersion ID: {content_version_id}")
+
+        # Generate the public URL for the image
+        image_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version_id}"
+        logger.debug(f"Generated image URL: {image_url}")
+        return image_url
     except Exception as e:
-        logger.error(f"Failed to create temporary directory {temp_dir}: {e}")
+        logger.error(f"Error uploading image to Salesforce: {str(e)}", exc_info=True)
         raise
 
+def upload_image_and_get_url(image_path):
+    """
+    Upload the image to Salesforce and return a public URL.
+    """
+    from datetime import datetime
+    import uuid
+
     # Generate a unique filename to avoid conflicts
     unique_filename = f"{uuid.uuid4().hex}_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.jpg"
-    temp_path = os.path.join(temp_dir, unique_filename)
-
-    # Copy the image to the temporary directory
-    try:
-        shutil.copy(image_path, temp_path)
-        os.chmod(temp_path, 0o664)
-        logger.info(f"Successfully copied image to temporary path: {temp_path}")
-    except Exception as e:
-        logger.error(f"Failed to copy image to {temp_path}: {e}")
-        raise
-
-    # Step 2: Define the uploads directory in the Space's repository root
-    # The repository root is typically the current working directory in Hugging Face Spaces
-    repo_root = os.getcwd()  # Should be the root of the Space (e.g., where app.py is)
-    public_dir = os.path.join(repo_root, "uploads")
-
-    # Create the uploads directory if it doesn't exist
-    try:
-        os.makedirs(public_dir, exist_ok=True)
-        os.chmod(public_dir, 0o775)
-        logger.info(f"Created or verified uploads directory at: {public_dir}")
-    except Exception as e:
-        logger.error(f"Failed to create uploads directory: {e}")
-        raise
-
-    # Step 3: Copy the file from /tmp to the persistent uploads folder
-    destination_path = os.path.join(public_dir, unique_filename)
+    
+    # Upload the image to Salesforce and get the URL
     try:
-        shutil.move(temp_path, destination_path)
-        os.chmod(destination_path, 0o664)
-        logger.info(f"Successfully moved image to: {destination_path}")
+        image_url = upload_image_to_salesforce(image_path, unique_filename)
+        return image_url
     except Exception as e:
-        logger.error(f"Failed to move image to {destination_path}: {e}")
+        logger.error(f"Failed to upload image to Salesforce: {e}")
         raise
 
-    # Step 4: Construct the public URL using the HF_SPACE_URL from .env
-    space_url = os.environ.get('HF_SPACE_URL', 'https://huggingface.co/spaces/pavansuresh/Smartlab_tech')
-    public_url = f"{space_url.rstrip('/')}/uploads/{unique_filename}"
-    logger.info(f"Generated public URL: {public_url}")
-
-    return public_url
-
 def save_record_to_salesforce(annotated_image_url, coverage_percent, original_image_pil, compliance_threshold=80):
     sf = Salesforce(
         username=os.environ['SF_USERNAME'],