Update app.py

app.py

@@ -2,12 +2,11 @@ import gradio as gr
 import requests
 import os
 from PIL import Image
-import numpy as np
-from transformers import pipeline
-from simple_salesforce import Salesforce
 import io
 import time
 from dotenv import load_dotenv
+from simple_salesforce import Salesforce
+from transformers import pipeline
 
 # Load environment variables from .env file
 load_dotenv()
@@ -26,22 +25,30 @@ def validate_photo_size(image_file):
 def predict_milestone(image):
     try:
         start_time = time.time()
+        # Use a simpler model for demo purposes; replace with a construction-specific model if available
         model = pipeline("image-classification", model="microsoft/resnet-50")
         predictions = model(image)
-        milestone = predictions[0]["label"]
+        
+        # Debugging: Print predictions to understand the output
+        print("Model predictions:", predictions)
+        
+        # Assuming the model returns labels like 'building' or 'construction', map them accordingly
+        top_prediction = predictions[0]["label"]
         confidence = predictions[0]["score"]
 
+        # Updated milestone mapping based on possible ResNet-50 labels
         milestone_map = {
-            "positive": "Walls Erected",
-            "negative": "Foundation Completed",
+            "building": "Walls Erected",
+            "structure": "Foundation Completed",
         }
         completion_map = {
-            "positive": 60.00,
-            "negative": 20.00,
+            "building": 60.00,
+            "structure": 20.00,
         }
 
-        predicted_milestone = milestone_map.get(milestone, "Unknown Milestone")
-        completion_percentage = completion_map.get(milestone, 0.00)
+        # Fallback to a default if the label isn't in the map
+        predicted_milestone = milestone_map.get(top_prediction, "Unknown Milestone")
+        completion_percentage = completion_map.get(top_prediction, 0.00)
 
         processing_time = time.time() - start_time
         if processing_time > 5:
@@ -51,17 +58,20 @@ def predict_milestone(image):
     except Exception as e:
         return None, None, f"AI failed to process the image: {str(e)}"
 
-# Placeholder for image upload to Salesforce (you can replace with actual logic)
+# Placeholder for image upload to Salesforce (replace with actual logic if needed)
 def upload_image_to_salesforce(image, project_name):
     try:
+        # Simulate uploading the image to Salesforce (replace with actual Salesforce file upload logic)
         image_url = f"https://your-salesforce-instance.com/file/{project_name}.jpg"
         return image_url, None
     except Exception as e:
         return None, f"Failed to upload image to Salesforce: {str(e)}"
 
-# Function to update Salesforce Construction_Progress__c record
+# Function to update Salesforce custom object record
 def update_salesforce_record(sf, project_name, milestone, percentage, image_url, status, comments):
     try:
+        # Ensure the custom object name matches your Salesforce setup
+        # Replace 'Construction_Progress__c' with the correct API name of your custom object
         query = f"SELECT Id FROM Construction_Progress__c WHERE Name = '{project_name}'"
         result = sf.query(query)
 
@@ -70,15 +80,17 @@ def update_salesforce_record(sf, project_name, milestone, percentage, image_url,
 
         record_id = result['records'][0]['Id']
 
+        # Update the record with the correct field API names
         sf.Construction_Progress__c.update(record_id, {
-            'Current_Milestone__c': milestone,
-            'Completion_Percentage__c': percentage,
-            'Last_Updated_Image__c': image_url,
+            'Current_Milestone__c': milestone if milestone else 'N/A',
+            'Completion_Percentage__c': float(percentage) if percentage else 0.0,
+            'Last_Updated_Image__c': image_url if image_url else 'N/A',
             'Last_Updated_On__c': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
             'Upload_Status__c': status,
-            'Comments__c': comments
+            'Comments__c': comments if comments else 'No comments'
         })
 
+        # Fetch the updated record
         updated_query = f"""
         SELECT Current_Milestone__c, Last_Updated_Image__c, Last_Updated_On__c, Upload_Status__c
         FROM Construction_Progress__c WHERE Id = '{record_id}'
@@ -104,20 +116,23 @@ def process_construction_photo(project_name, image):
     if not project_name or not image:
         return None, "Please provide a project name and upload a photo."
 
+    # Authenticate with Salesforce
     try:
         sf = Salesforce(
             username=os.getenv('SALESFORCE_USERNAME'),
             password=os.getenv('SALESFORCE_PASSWORD'),
             security_token=os.getenv('SALESFORCE_SECURITY_TOKEN'),
-            domain=os.getenv('SALESFORCE_DOMAIN')
+            domain=os.getenv('SALESFORCE_DOMAIN', 'login')  # Default to 'login' if not specified
         )
     except Exception as e:
         return None, f"Failed to connect to Salesforce: {str(e)}"
 
+    # Validate photo size
     is_valid, error = validate_photo_size(image)
     if not is_valid:
         return None, error or "Photo is too large! Please upload a photo smaller than 20MB."
 
+    # Predict milestone using AI
     milestone, percentage, error = predict_milestone(image)
 
     if error:
@@ -139,6 +154,7 @@ def process_construction_photo(project_name, image):
                 error_text += f"{field}: {value}\n"
         return None, error_text
 
+    # Upload image to Salesforce
     image_url, upload_error = upload_image_to_salesforce(image, project_name)
 
     if upload_error:
@@ -160,6 +176,7 @@ def process_construction_photo(project_name, image):
                 error_text += f"{field}: {value}\n"
         return None, error_text
 
+    # Update Salesforce record
     fields, error_message = update_salesforce_record(
         sf=sf,
         project_name=project_name,
@@ -173,6 +190,7 @@ def process_construction_photo(project_name, image):
     if error_message:
         return None, f"Salesforce Error: {error_message}"
 
+    # Prepare success message
     result_text = f"Success! Milestone: {milestone}, Completion: {percentage}%\nProgress saved to Salesforce!\n\nSalesforce Fields:\n"
     for field, value in fields.items():
         result_text += f"{field}: {value}\n"
@@ -195,4 +213,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()