Update app.py

app.py

@@ -1,90 +1,123 @@
 import gradio as gr
-import datetime
+from PIL import Image
 import os
-import shutil
-import uuid
-from fastapi import FastAPI
-from gradio.routes import mount_gradio_app
+from dotenv import load_dotenv
 from simple_salesforce import Salesforce
+from datetime import datetime
+import random
+import shutil
+
+# Load environment variables
+load_dotenv()
+SF_USERNAME = os.getenv("SF_USERNAME")
+SF_PASSWORD = os.getenv("SF_PASSWORD")
+SF_SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")
+PUBLIC_URL_BASE = os.getenv("PUBLIC_URL_BASE", "https://your-public-host.com/public_uploads")
 
-# Salesforce login
-sf = Salesforce(
-    username='your_username',
-    password='your_password',
-    security_token='your_token'
-)
-
-# Hugging Face public folder for Spaces
-PUBLIC_DIR = "/home/user/app/public_uploads"
-os.makedirs(PUBLIC_DIR, exist_ok=True)
-
-# FastAPI instance
-app = FastAPI()
-
-# Dummy model for milestone detection
-def predict_milestone(image_path):
-    return {
-        "milestone": "Foundation",
-        "confidence": 0.9,
-        "completion_percentage": 80
-    }
-
-def process_image(image):
-    # Save uploaded image with unique name
-    timestamp = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
-    filename = f"{timestamp}_{os.path.basename(image.name)}"
-    local_path = os.path.join(PUBLIC_DIR, filename)
-    shutil.copy(image.name, local_path)
-
-    # Publicly accessible URL on Hugging Face Spaces
-    public_url = f"https://{os.environ['HF_SPACE_ID']}.hf.space/file/public_uploads/{filename}"
-
-    # Model prediction
-    result = predict_milestone(local_path)
-    milestone = result["milestone"]
-    confidence = result["confidence"]
-    completion = result["completion_percentage"]
-
-    # Save to Salesforce
-    now_str = datetime.datetime.now().strftime('%d/%m/%Y, %I:%M %p')
-    sf.Construction__c.create({
-        'Project_Name__c': 'Construction',
-        'Current_Milestone__c': milestone,
-        'Completion_Percentage__c': f"{completion}%",
-        'Last_Updated_Image__c': public_url,
-        'Last_Updated_On__c': now_str,
-        'Upload_Status__c': 'Success',
-        'Comments__c': f"AI Prediction: {milestone} with {confidence*100:.1f}% confidence"
-    })
-
-    return (
-        f"Success: Milestone: {milestone}, Completion: {completion}%",
-        "Success",
-        milestone,
-        f"Confidence Score: {confidence}",
-        completion
+# 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,
+        password=SF_PASSWORD,
+        security_token=SF_SECURITY_TOKEN,
+        domain='login'
     )
+except Exception as e:
+    print(f"Salesforce connection failed: {str(e)}")
+    raise
 
-# Gradio UI
-def build_ui():
-    with gr.Blocks() as demo:
-        gr.Markdown("## Process Image")
-        with gr.Row():
-            image_input = gr.Image(type="filepath", label="Upload Image")
-        with gr.Row():
-            result = gr.Textbox(label="Result")
-        with gr.Row():
-            upload_status = gr.Textbox(label="Upload Status")
-            milestone = gr.Textbox(label="Detected Milestone")
-        with gr.Row():
-            confidence = gr.Textbox(label="Confidence Score")
-            percentage = gr.Slider(minimum=0, maximum=100, label="Completion Percentage")
-        image_input.change(
-            fn=process_image,
-            inputs=[image_input],
-            outputs=[result, upload_status, milestone, confidence, percentage]
+# Valid milestones
+VALID_MILESTONES = ["Foundation", "Walls Erected", "Planning", "Completed"]
+
+# Mock AI prediction
+def mock_ai_model(image):
+    img = image.convert("RGB")
+    max_size = 1024
+    img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+    milestone = random.choice(VALID_MILESTONES)
+    completion_percent = random.choice([10, 30, 50, 80, 100])
+    confidence_score = round(random.uniform(0.85, 0.95), 2)
+    return milestone, completion_percent, confidence_score
+
+# Gradio function
+def process_image(image, project_name):
+    try:
+        if image is None:
+            return "Error: Please upload an image to proceed.", "Pending", "", "", 0
+
+        img = Image.open(image)
+        image_size_mb = os.path.getsize(image) / (1024 * 1024)
+        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
+
+        # Save image to public folder
+        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_image_path = os.path.join(upload_dir, image_filename)
+        shutil.copy(image, saved_image_path)
+
+        # Publicly accessible URL (replace base with deployed static site URL)
+        image_url = f"{PUBLIC_URL_BASE}/{image_filename}"
+
+        # AI Prediction
+        milestone, percent_complete, confidence_score = mock_ai_model(img)
+
+        # Format timestamp
+        formatted_timestamp = datetime.now().strftime("%d/%m/%Y, %I:%M %p")
+
+        # Salesforce record
+        record = {
+            "Name__c": project_name,
+            "Current_Milestone__c": milestone,
+            "Completion_Percentage__c": percent_complete,
+            "Last_Updated_On__c": formatted_timestamp,
+            "Upload_Status__c": "Success",
+            "Comments__c": f"AI Prediction: {milestone} with {confidence_score * 100:.1f}% confidence",
+            "Last_Updated_Image__c": image_url
+        }
+
+        try:
+            sf.Construction__c.create(record)
+        except Exception as e:
+            return f"Error: Failed to update Salesforce - {str(e)}", "Failure", "", "", 0
+
+        return (
+            f"Success: Milestone: {milestone}, Completion: {percent_complete}%",
+            "Success",
+            milestone,
+            f"Confidence Score: {confidence_score}",
+            percent_complete
         )
-    return demo
 
-demo_app = build_ui()
-mount_gradio_app(app, demo_app, path="/")
+    except Exception as e:
+        return f"Error: {str(e)}", "Failure", "", "", 0
+
+# Gradio UI
+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>")
+    with gr.Row():
+        image_input = gr.Image(type="filepath", label="Upload Construction Site Photo (JPG/PNG, ≤ 20MB)")
+        project_name_input = gr.Textbox(label="Project Name (Required)", placeholder="e.g. Project_12345")
+
+    submit_button = gr.Button("Process Image")
+    output_text = gr.Textbox(label="Result")
+    upload_status = gr.Textbox(label="Upload Status")
+    milestone = gr.Textbox(label="Detected Milestone")
+    confidence = gr.Textbox(label="Confidence Score")
+    progress = gr.Slider(0, 100, label="Completion Percentage", interactive=False, value=0)
+
+    submit_button.click(
+        fn=process_image,
+        inputs=[image_input, project_name_input],
+        outputs=[output_text, upload_status, milestone, confidence, progress]
+    )
+
+demo.launch(share=False)