Update app.py

app.py

@@ -16,72 +16,61 @@ load_dotenv()
 def validate_photo_size(image_file):
     max_size_mb = 20
     if isinstance(image_file, Image.Image):
-        # Convert PIL Image to bytes for size check
         img_byte_arr = io.BytesIO()
         image_file.save(img_byte_arr, format='JPEG')
-        file_size_mb = img_byte_arr.tell() / (1024 * 1024)  # Convert bytes to MB
+        file_size_mb = img_byte_arr.tell() / (1024 * 1024)
         return file_size_mb <= max_size_mb, None
     return False, "Invalid image format"
 
 # Function to process image with AI and predict milestone
 def predict_milestone(image):
     try:
-        # Simulate AI processing time (ensure < 5 seconds)
         start_time = time.time()
-        
-        # Process image with Hugging Face model
         model = pipeline("image-classification", model="microsoft/resnet-50")
         predictions = model(image)
-        
-        # Placeholder logic: Map model output to construction milestones
-        milestone = predictions[0]["label"]  # Example: "positive" -> "Walls Erected"
+        milestone = predictions[0]["label"]
         confidence = predictions[0]["score"]
-        
-        # Map model output to construction milestones (customize this)
+
         milestone_map = {
             "positive": "Walls Erected",
             "negative": "Foundation Completed",
-            # Add more mappings based on your model
         }
         completion_map = {
-            "positive": 60.00,  # Example: Walls = 60% complete
-            "negative": 20.00,  # Example: Foundation = 20% complete
+            "positive": 60.00,
+            "negative": 20.00,
         }
-        
+
         predicted_milestone = milestone_map.get(milestone, "Unknown Milestone")
         completion_percentage = completion_map.get(milestone, 0.00)
-        
+
         processing_time = time.time() - start_time
         if processing_time > 5:
             return None, None, "AI took too long to process (> 5 seconds)."
-        
+
         return predicted_milestone, completion_percentage, None
     except Exception as e:
         return None, None, f"AI failed to process the image: {str(e)}"
 
-# Function to upload image to Salesforce and get a URL
+# Placeholder for image upload to Salesforce (you can replace with actual logic)
 def upload_image_to_salesforce(image, project_name):
     try:
-        # Placeholder: Simulate uploading image to Salesforce ContentVersion
-        image_url = f"https://your-salesforce-instance.com/file/{project_name}.jpg"  # Simulated URL
+        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_Project__c object and fetch fields
+# Function to update Salesforce Construction_Progress__c record
 def update_salesforce_record(sf, project_name, milestone, percentage, image_url, status, comments):
     try:
-        # Query to check if the project exists
-        query = f"SELECT Id FROM Construction_Project__c WHERE Name = '{project_name}'"
+        query = f"SELECT Id FROM Construction_Progress__c WHERE Name = '{project_name}'"
         result = sf.query(query)
-        
+
         if result['totalSize'] == 0:
             return None, f"No project found with Name: {project_name}"
-        
+
         record_id = result['records'][0]['Id']
-        
-        # Update the record
-        sf.Construction_Project__c.update(record_id, {
+
+        sf.Construction_Progress__c.update(record_id, {
             'Current_Milestone__c': milestone,
             'Completion_Percentage__c': percentage,
             'Last_Updated_Image__c': image_url,
@@ -89,14 +78,16 @@ def update_salesforce_record(sf, project_name, milestone, percentage, image_url,
             'Upload_Status__c': status,
             'Comments__c': comments
         })
-        
-        # Fetch the updated record to get the specified fields
-        updated_query = f"SELECT Current_Milestone__c, Last_Updated_Image__c, Last_Updated_On__c, Upload_Status__c FROM Construction_Project__c WHERE Id = '{record_id}'"
+
+        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}'
+        """
         updated_result = sf.query(updated_query)
-        
+
         if updated_result['totalSize'] == 0:
             return None, "Failed to retrieve updated record."
-        
+
         record = updated_result['records'][0]
         fields_output = {
             'Current_Milestone__c': record.get('Current_Milestone__c', 'N/A'),
@@ -112,8 +103,7 @@ def update_salesforce_record(sf, project_name, milestone, percentage, image_url,
 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."
-    
-    # Connect to Salesforce
+
     try:
         sf = Salesforce(
             username=os.getenv('SALESFORCE_USERNAME'),
@@ -124,14 +114,12 @@ def process_construction_photo(project_name, image):
     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."
 
-    # Process the image with AI
     milestone, percentage, error = predict_milestone(image)
-    
+
     if error:
         fields, error_message = update_salesforce_record(
             sf=sf,
@@ -150,10 +138,9 @@ def process_construction_photo(project_name, image):
             for field, value in fields.items():
                 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:
         fields, error_message = update_salesforce_record(
             sf=sf,
@@ -172,8 +159,7 @@ def process_construction_photo(project_name, image):
             for field, value in fields.items():
                 error_text += f"{field}: {value}\n"
         return None, error_text
-    
-    # Update Salesforce with success
+
     fields, error_message = update_salesforce_record(
         sf=sf,
         project_name=project_name,
@@ -183,15 +169,14 @@ def process_construction_photo(project_name, image):
         status="Success",
         comments="Photo processed successfully"
     )
-    
+
     if error_message:
         return None, f"Salesforce Error: {error_message}"
-    
-    # Prepare output with AI results and Salesforce fields
+
     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"
-    
+
     return image, result_text
 
 # Gradio interface
@@ -209,7 +194,5 @@ iface = gr.Interface(
     description="Upload a photo of your construction site, and the AI will tell you the progress!"
 )
 
-# Launch the Gradio app
 if __name__ == "__main__":
     iface.launch()
-