Update app.py

app.py

@@ -15,11 +15,9 @@ SF_USERNAME = os.getenv("SF_USERNAME")
 SF_PASSWORD = os.getenv("SF_PASSWORD")
 SF_SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")
 
-# Validate Salesforce credentials
 if not all([SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN]):
     raise ValueError("Missing Salesforce credentials. Set SF_USERNAME, SF_PASSWORD, and SF_SECURITY_TOKEN in environment variables.")
 
-# Initialize Salesforce connection
 try:
     sf = Salesforce(
         username=SF_USERNAME,
@@ -31,13 +29,9 @@ except Exception as e:
     print(f"Salesforce connection failed: {str(e)}")
     raise
 
-# Valid milestones
 VALID_MILESTONES = ["Planning", "Foundation", "Walls Erected", "Completed"]
-
-# Adjust the timezone to your local timezone
 local_timezone = pytz.timezone("Asia/Kolkata")
 
-# Image processing and Salesforce upload
 def process_image(images, project_name):
     try:
         if not images or len(images) == 0:
@@ -46,15 +40,13 @@ def process_image(images, project_name):
         milestones = []
         for image in images:
             img = Image.open(image)
-            
-            # Check if the image exceeds the size limit (20MB)
+
             image_size_mb = os.path.getsize(image) / (1024 * 1024)
             if image_size_mb > 20:
                 return "Error: One or more images exceed 20MB.", "Failure", "", "", 0
             if not str(image).lower().endswith(('.jpg', '.jpeg', '.png')):
                 return "Error: Only JPG/PNG images are supported.", "Failure", "", "", 0
 
-            # Save image to public folder temporarily before uploading to Salesforce
             upload_dir = "public_uploads"
             os.makedirs(upload_dir, exist_ok=True)
             unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
@@ -62,38 +54,40 @@ def process_image(images, project_name):
             saved_image_path = os.path.join(upload_dir, image_filename)
             shutil.copy(image, saved_image_path)
 
-            # Convert image to base64 before uploading to Salesforce
             with open(saved_image_path, 'rb') as image_file:
                 image_data = base64.b64encode(image_file.read()).decode('utf-8')
 
-            # Create the ContentVersion record in Salesforce
-            content_version = {
-                'Title': image_filename,
-                'PathOnClient': saved_image_path,
-                'VersionData': image_data
-            }
-
-            # Upload the file to Salesforce
             try:
+                content_version = {
+                    'Title': image_filename,
+                    'PathOnClient': saved_image_path,
+                    'VersionData': image_data
+                }
                 content_version_result = sf.ContentVersion.create(content_version)
                 content_version_id = content_version_result['id']
                 file_url = f"https://sathkruthatechsolutionspri8-dev-ed.develop.lightning.force.com/{content_version_id}"
             except Exception as e:
                 return f"Error: Failed to upload image to Salesforce - {str(e)}", "Failure", "", "", 0
 
-            # Now check for completion criteria
-            if "completed" in os.path.basename(image).lower() or "finished" in os.path.basename(image).lower():
-                milestone = "Completed"  # Directly detect completion if such features exist
+            # 👇 UPDATED milestone detection based on filename
+            filename = os.path.basename(image).lower()
+            if "completed" in filename or "finished" in filename:
+                milestone = "Completed"
+            elif "walls" in filename or "erected" in filename:
+                milestone = "Walls Erected"
+            elif "foundation" in filename:
+                milestone = "Foundation"
+            elif "planning" in filename or "plan" in filename:
+                milestone = "Planning"
             else:
-                # Default milestone detection logic (mock AI model or image-based hash)
                 img_bytes = img.tobytes()
                 img_hash = int(hashlib.sha256(img_bytes).hexdigest(), 16)
                 milestone_index = img_hash % len(VALID_MILESTONES)
                 milestone = VALID_MILESTONES[milestone_index]
+            # 👆 END update
 
             milestones.append(milestone)
 
-        # Determine overall milestone (most advanced detected)
         final_milestone = max(set(milestones), key=milestones.count) if milestones else "Planning"
         milestone_completion_map = {
             "Planning": 10,
@@ -224,7 +218,7 @@ def process_image(images, project_name):
     except Exception as e:
         return f"Error: {str(e)}", "Failure", "", "", "0%"
 
-# Gradio UI
+# UI
 with gr.Blocks(css="""
     .gradio-container {
         background-color: #f0f4f8;