Update app.py

app.py

@@ -1,122 +1,90 @@
 import gradio as gr
-from PIL import Image
+import datetime
 import os
-from dotenv import load_dotenv
-from simple_salesforce import Salesforce
-from datetime import datetime
-import random
 import shutil
+import uuid
+from fastapi import FastAPI
+from gradio.routes import mount_gradio_app
+from simple_salesforce import Salesforce
 
-# 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")
-
-# 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'
+# 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
     )
-except Exception as e:
-    print(f"Salesforce connection failed: {str(e)}")
-    raise
-
-# 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 for URL generation
-        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)
-
-        # Create public URL assuming you're serving /public_uploads/ via static web server (e.g., on localhost or external host)
-        public_url_base = os.getenv("PUBLIC_URL_BASE", "http://localhost:7860/public_uploads")
-        image_url = f"{public_url_base}/{image_filename}"
-
-        # Predict
-        milestone, percent_complete, confidence_score = mock_ai_model(img)
-
-        # Construct Salesforce record
-        record = {
-            "Name__c": project_name,
-            "Current_Milestone__c": milestone,
-            "Completion_Percentage__c": percent_complete,
-            "Last_Updated_On__c": datetime.now().isoformat(),
-            "Upload_Status__c": "Success",
-            "Comments__c": f"AI Prediction: {milestone} with {confidence_score*100}% 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
-        )
-
-    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)
-
+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]
+        )
+    return demo
 
+demo_app = build_ui()
+mount_gradio_app(app, demo_app, path="/")