Update app.py

app.py

@@ -9,9 +9,10 @@ import shutil
 import base64
 import pytz
 import logging
+import tempfile
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 # Load environment variables
@@ -21,7 +22,8 @@ SF_PASSWORD = os.getenv("SF_PASSWORD")
 SF_SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")
 
 if not all([SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN]):
-    raise ValueError("Missing Salesforce credentials.")
+    logger.error("Missing Salesforce credentials.")
+    raise ValueError("Missing Salesforce credentials in .env file.")
 
 # Connect to Salesforce
 try:
@@ -29,135 +31,106 @@ try:
         username=SF_USERNAME,
         password=SF_PASSWORD,
         security_token=SF_SECURITY_TOKEN,
-        domain="login"
+        domain='login'
     )
-    logger.info("Connected to Salesforce successfully.")
+    logger.info("Successfully connected to Salesforce.")
 except Exception as e:
     logger.error(f"Salesforce connection failed: {str(e)}")
     raise
 
 VALID_MILESTONES = ["Planning", "Foundation", "Walls Erected", "Completed"]
-local_timezone = pytz.timezone("Asia/Kolkata")
-
-# Simulated content-based milestone detection (placeholder for ML model)
+MILESTONE_COMPLETION_MAP = {
+    "Planning": 25,
+    "Foundation": 50,
+    "Walls Erected": 75,
+    "Completed": 100
+}
+LOCAL_TIMEZONE = pytz.timezone("Asia/Kolkata")
+ALLOWED_EXTENSIONS = {'.jpg', '.jpeg', '.png'}
+MAX_FILE_SIZE_MB = 20
+
+# Milestone detection based on image brightness
 def detect_milestone_from_image(image_path):
-    """
-    Placeholder for content-based milestone detection.
-    In a real implementation, use a trained ML model to analyze image features.
-    """
     try:
-        img = Image.open(image_path).convert("RGB")
+        img = Image.open(image_path).convert("L")
         img_array = np.array(img)
-        # Example: Use color distribution as a proxy for milestone
-        red_mean = np.mean(img_array[:, :, 0])
-        green_mean = np.mean(img_array[:, :, 1])
-        blue_mean = np.mean(img_array[:, :, 2])
+        brightness = np.mean(img_array)
+        logger.debug(f"Image {image_path} brightness: {brightness}")
 
-        # Simulated logic: Adjust these thresholds based on real data or ML model
-        if red_mean > 150 and green_mean < 100:  # High red (e.g., soil)
+        if brightness < 80:
             return "Planning"
-        elif green_mean > 120:  # High green (e.g., vegetation or scaffolding)
+        elif 80 <= brightness < 130:
             return "Foundation"
-        elif blue_mean > 130:  # High blue (e.g., sky visible with walls)
+        elif 130 <= brightness < 180:
             return "Walls Erected"
-        else:  # Balanced colors (e.g., finished building)
+        else:
             return "Completed"
     except Exception as e:
-        logger.error(f"Error in milestone detection: {str(e)}")
-        return "Planning"  # Default to earliest milestone on error
+        logger.error(f"Failed to detect milestone for {image_path}: {str(e)}")
+        return "Planning"  # Default milestone on error
 
-def validate_project_name(project_name):
-    """Check if project name is unique in Salesforce."""
-    try:
-        query = f"SELECT Id FROM Construction__c WHERE Name__c = '{project_name}'"
-        result = sf.query(query)
-        if result["totalSize"] > 0:
-            return False, "Error: Project name already exists in Salesforce."
-        return True, ""
-    except Exception as e:
-        logger.error(f"Error validating project name: {str(e)}")
-        return False, f"Error: Failed to validate project name - {str(e)}"
-
-def calculate_completion_percentage(milestone, image_count):
-    """Calculate completion percentage based on milestone and image count."""
-    milestone_base = {
-        "Planning": 10,
-        "Foundation": 30,
-        "Walls Erected": 50,
-        "Completed": 90
-    }
-    base_percentage = milestone_base.get(milestone, 10)
-    # Add incremental progress based on image count (max 10% boost)
-    increment = min(image_count * 2, 10)
-    total_percentage = min(base_percentage + increment, 100)
-    return total_percentage
-
-def process_image(images, project_name, confirmed_milestone=None):
+def process_image(images, project_name):
     try:
-        if not images or len(images) == 0:
-            return "Error: Please upload at least one image.", "Pending", "", "", 0, []
+        if not project_name or not project_name.strip():
+            return "Error: Project name is required.", "Failure", "", "", "0%"
 
-        if not project_name:
-            return "Error: Project name is required.", "Failure", "", "", 0, []
-
-        # Validate project name uniqueness
-        is_valid, validation_error = validate_project_name(project_name)
-        if not is_valid:
-            return validation_error, "Failure", "", "", 0, []
+        if not images or len(images) == 0:
+            return "Error: Please upload at least one image.", "Failure", "", "", "0%"
 
         milestones = []
         file_urls = []
-        for image in images:
-            img = Image.open(image)
-            image_size_mb = os.path.getsize(image) / (1024 * 1024)
-            if image_size_mb > 20:
-                return f"Error: Image {image} exceeds 20MB.", "Failure", "", "", 0, []
-            if not str(image).lower().endswith(('.jpg', '.jpeg', '.png')):
-                return "Error: Only JPG or PNG images are supported.", "Failure", "", "", 0, []
-
-            # Save image to local directory
-            upload_dir = "public_uploads"
-            os.makedirs(upload_dir, exist_ok=True)
-            unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
-            image_filename = f"{unique_id}_{os.path.basename(image)}"
-            saved_path = os.path.join(upload_dir, image_filename)
-            shutil.copy2(image, saved_path)
-
-            # Upload image to Salesforce
-            with open(saved_path, 'rb') as image_file:
-                image_data = base64.b64encode(image_file.read()).decode('8')
-
-            try:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            for image in images:
+                # Validate file extension
+                ext = os.path.splitext(image.name)[1].lower()
+                if ext not in ALLOWED_EXTENSIONS:
+                    return f"Error: {os.path.basename(image.name)} is not a supported format (JPG/PNG only).", "Failure", "", "", "0%"
+
+                # Validate file size
+                image_size_mb = os.path.getsize(image.name) / (1024 * 1024)
+                if image_size_mb > MAX_FILE_SIZE_MB:
+                    return f"Error: {os.path.basename(image.name)} exceeds {MAX_FILE_SIZE_MB}MB.", "Failure", "", "", "0%"
+
+                # Save image to temporary directory
+                unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
+                image_filename = f"{unique_id}_{os.path.basename(image.name)}"
+                temp_image_path = os.path.join(temp_dir, image_filename)
+                shutil.copy(image.name, temp_image_path)
+
+                # Upload to Salesforce
+                with open(temp_image_path, 'rb') as image_file:
+                    image_data = base64.b64encode(image_file.read()).decode('utf-8')
+
                 content_version = {
                     'Title': image_filename,
                     'PathOnClient': image_filename,
                     '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}"
-                file_urls.append(file_url)
-            except Exception as e:
-                logger.error(f"Failed to upload image to Salesforce: {str(e)}")
-                return f"Error: Failed to upload image - {str(e)}", "Failure", "", "", 0
-
-            milestone = detect_milestone_from_image(saved_path)
-            milestones.append(milestone)
+                try:
+                    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}"
+                    file_urls.append(file_url)
+                    logger.info(f"Uploaded {image_filename} to Salesforce: {file_url}")
+                except Exception as e:
+                    logger.error(f"Failed to upload {image_filename} to Salesforce: {str(e)}")
+                    return f"Error: Failed to upload {image_filename} to Salesforce.", "Failure", "", "", "0%"
+
+                # Detect milestone
+                milestone = detect_milestone_from_image(temp_image_path)
+                milestones.append(milestone)
 
         # Determine final milestone
-        final_milestone = confirmed_milestone or max(set(milestones), key=milestones.count)
-        if final_milestone not in VALID_MILESTONES:
-            return "Error: Invalid milestone detected.", "Failure", "", "", 0, []
-
-        # Calculate completion percentage
-        percent_complete = calculate_completion_percentage(final_milestone, len(images))
+        final_milestone = max(set(milestones), key=milestones.count) if milestones else "Planning"
+        percent_complete = MILESTONE_COMPLETION_MAP.get(final_milestone, 0)
 
-        # Prepare Salesforce record
-        now = datetime.now(local_timezone)
-        local_time = now.strftime("%Y-%m-%dT%H:%M:%S") + now.strftime("%z")[:-2] + ":" + now.strftime("%z")[-2:]
+        # Get current time in local timezone
+        now = datetime.now(LOCAL_TIMEZONE)
+        local_time = now.strftime("%Y-%m-%dT%H:%M:%S%z")[:-2] + ":" + now.strftime("%z")[-2:]
 
+        # Create Salesforce record
         record = {
-            "Name__c": project_name,
+            "Name__c": project_name.strip(),
             "Current_Milestone__c": final_milestone,
             "Completion_Percentage__c": percent_complete,
             "Last_Updated_On__c": local_time,
@@ -168,49 +141,57 @@ def process_image(images, project_name, confirmed_milestone=None):
 
         try:
             sf.Construction__c.create(record)
-            logger.info(f"Created Salesforce record for project: {project_name}")
+            logger.info(f"Created Salesforce record for project {project_name}: {final_milestone}")
         except Exception as e:
             logger.error(f"Failed to create Salesforce record: {str(e)}")
-            return f"Error: Failed to update Salesforce - {str(e)}", "Failure", "", "", 0, []
+            return f"Error: Failed to update Salesforce record.", "Failure", "", "", "0%"
 
         return (
             f"<h3 style='color:green;'>Milestone: {final_milestone}</h3>",
             "Success",
             final_milestone,
             f"{percent_complete}%",
-            percent_complete,
-            file_urls
+            file_urls[-1] if file_urls else ""
         )
 
     except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}")
-        return f"Error: {str(e)}", "Failure", "", "", 0, []
+        logger.error(f"Unexpected error in process_image: {str(e)}")
+        return f"Error: An unexpected error occurred - {str(e)}", "Failure", "", "", "0%"
 
 # Gradio UI
-with gr.Blocks(css=".title {text-align: center; font-size: 24px;}") as demo:
+with gr.Blocks(theme=gr.themes.Soft(), css=".title {text-align: center; font-size: 2em; margin-bottom: 20px;}") as demo:
     gr.Markdown("<h1 class='title'>Construction Progress Analyzer</h1>")
     with gr.Row():
         with gr.Column(scale=2):
-            image_input = gr.Files(type="filepath", label="Upload Construction Site Photos (JPG/PNG, ≤ 20MB)")
-            project_name_input = gr.Textbox(label="Project Name (Required)", placeholder="e.g. Project_12345")
+            image_input = gr.Files(
+                file_types=[".jpg", ".jpeg", ".png"],
+                label=f"Upload Construction Site Photos (JPG/PNG, ≤ {MAX_FILE_SIZE_MB}MB)",
+                file_count="multiple"
+            )
         with gr.Column(scale=1):
-            milestone_dropdown = gr.Dropdown(
-                choices=VALID_MILESTONES,
-                label="Confirm Milestone (Optional)",
-                info="Select to override detected milestone"
+            project_name_input = gr.Textbox(
+                label="Project Name (Required)",
+                placeholder="e.g., Highrise_001",
+                lines=1
             )
 
-    submit_button = gr.Button("Process Image")
+    submit_button = gr.Button("Process Images", variant="primary")
     output_html = gr.HTML(label="Result")
-    upload_status = gr.Textbox(label="Upload Status", interactive=False)
-    milestone_output = gr.Textbox(label="Detected Milestone", interactive=False)
-    progress = gr.Textbox(label="Completion Percentage", interactive=False)
+    with gr.Row():
+        upload_status = gr.Textbox(label="Upload Status", interactive=False)
+        milestone = gr.Textbox(label="Detected Milestone", interactive=False)
+        progress = gr.Textbox(label="Completion Percentage", interactive=False)
+    file_url = gr.Textbox(label="Last Image URL", interactive=False, visible=False)
 
     submit_button.click(
         fn=process_image,
-        inputs=[image_input, project_name_input, milestone_dropdown],
-        outputs=[output_html, upload_status, milestone_output, progress]
+        inputs=[image_input, project_name_input],
+        outputs=[output_html, upload_status, milestone, progress, file_url]
     )
 
-if __name__ == "__main__":
-    demo.launch(share=True)
+# Launch Gradio app
+try:
+    demo.launch(share=False, server_name="0.0.0.0", server_port=7860)
+except Exception as e:
+    logger.error(f"Failed to launch Gradio app: {str(e)}")
+    raise