Update app.py

app.py

@@ -22,54 +22,53 @@ try:
         username=SF_USERNAME,
         password=SF_PASSWORD,
         security_token=SF_SECURITY_TOKEN,
-        domain='login'  # Use 'test' for sandbox
+        domain='login'  # Use 'test' for sandbox orgs
     )
 except Exception as e:
     print(f"Salesforce connection failed: {str(e)}")
     raise
 
-# Mock AI model for milestone detection (since we can't train a real model here)
+# Safe milestone values (must match Salesforce picklist exactly)
+MILESTONE_MAPPING = {
+    "In Progress": "In_Progress",
+    "Started": "In_Progress",
+    "Completed": "Completed",
+    "Planning": "Planning"
+}
+
+# Mock AI model for milestone detection
 def mock_ai_model(image):
-    # Preprocessing: Resize, normalize (simulated)
     img = image.convert("RGB")
     max_size = 1024
     img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
 
-    # Feature Extraction and Milestone Detection (simulated)
-    # In a real scenario, this would use a CNN model trained on construction images
-    # Using a generic milestone value that is likely to be in a Salesforce picklist
-    milestone = "In Progress"  # Changed to a more generic value to avoid picklist error
-    completion_percent = 30  # Estimated based on the image
-    confidence_score = round(random.uniform(0.85, 0.95), 2)  # Random confidence between 85-95%
+    # Simulated AI result
+    predicted_label = "In Progress"  # This label is mapped to valid picklist
+    milestone = MILESTONE_MAPPING.get(predicted_label, "Planning")  # Safe fallback
+    completion_percent = 30
+    confidence_score = round(random.uniform(0.85, 0.95), 2)
 
     return milestone, completion_percent, confidence_score
 
-# Function for Gradio UI to process the image
+# Function for Gradio UI
 def process_image(image):
     try:
-        # Validate inputs
         if image is None:
             return "Error: Please upload an image to proceed.", "Pending", "", "", 0
 
-        # Open and validate image
         img = Image.open(image)
 
-        # Validate image size and type
         image_size_mb = os.path.getsize(image) / (1024 * 1024)
-        if image_size_mb > 20:  # Fixed the typo: image_size_mb_tb to image_size_mb
+        if image_size_mb > 20:
             return "Error: Image size exceeds 20MB.", "Failure", "", "", 0
         if not str(image).lower().endswith(('.jpg', '.jpeg', '.png')):
             return "Error: Only JPG/PNG images are supported.", "Failure", "", "", 0
 
-        # Run mock AI model
         milestone, percent_complete, confidence_score = mock_ai_model(img)
-
-        # Generate a unique identifier for the Salesforce record
         unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
 
-        # Store result in Salesforce
         record = {
-            "Name__c": unique_id,  # Use a timestamp-based unique ID
+            "Name__c": unique_id,
             "Current_Milestone__c": milestone,
             "Completion_Percentage__c": percent_complete,
             "Last_Updated_On__c": datetime.now().isoformat(),
@@ -78,7 +77,7 @@ def process_image(image):
         }
 
         try:
-            sf.Construction__c.create(record)  # Always create a new record
+            sf.Construction__c.create(record)
         except Exception as e:
             return f"Error: Failed to update Salesforce - {str(e)}", "Failure", "", "", 0
 
@@ -93,7 +92,7 @@ def process_image(image):
     except Exception as e:
         return f"Error: {str(e)}", "Failure", "", "", 0
 
-# Gradio interface for testing
+# Gradio interface
 with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: Arial;} .title {color: #2c3e50; font-size: 24px; text-align: center;}") as demo:
     gr.Markdown("<h1 class='title'>Construction Milestone Detector</h1>")
     image_input = gr.Image(type="filepath", label="Upload Construction Site Photo (JPG/PNG, ≤ 20MB)")
@@ -111,4 +110,4 @@ with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: A
     )
 
 # Launch the Gradio app
-demo.launch(share=False)  # Note: share=True may not work in all environments; you may need to run this locally
+demo.launch(share=False)