Update app.py

app.py

@@ -7,6 +7,10 @@ import re
 from dotenv import load_dotenv
 from simple_salesforce import Salesforce
 from datetime import datetime
+from fastapi import FastAPI, HTTPException, Security, Depends
+from fastapi.security import APIKeyHeader
+import base64
+import io
 
 # Load environment variables
 load_dotenv()
@@ -16,10 +20,11 @@ SF_PASSWORD = os.getenv("SF_PASSWORD")
 SF_SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")
 SF_CONSUMER_KEY = os.getenv("SF_CONSUMER_KEY")
 SF_CONSUMER_SECRET = os.getenv("SF_CONSUMER_SECRET")
+API_KEY = os.getenv("API_KEY", "your-api-key-here")  # Set in Space Secrets
 
-# Validate Salesforce credentials
+# Validate Salesforce credentials (for Gradio UI)
 if not all([SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, SF_CONSUMER_KEY, SF_CONSUMER_SECRET]):
-    raise ValueError("Missing Salesforce credentials. Please set SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, SF_CONSUMER_KEY, and SF_CONSUMER_SECRET in environment variables.")
+    raise ValueError("Missing Salesforce credentials. Set SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, SF_CONSUMER_KEY, and SF_CONSUMER_SECRET in environment variables.")
 
 # Initialize Salesforce connection
 try:
@@ -36,15 +41,71 @@ except Exception as e:
     raise
 
 # Load pre-trained model and processor
-model_name = "google/vit-base-patch16-224"  # Placeholder; replace with "nasreshsuguru/construction-milestone-detector" if trained
+model_name = "google/vit-base-patch16-224"
 processor = ViTImageProcessor.from_pretrained(model_name)
 model = ViTForImageClassification.from_pretrained(model_name)
 model.eval()
 
+# FastAPI app for API endpoint
+app = FastAPI()
+
+# API Key authentication
+api_key_header = APIKeyHeader(name="X-API-Key")
+async def verify_api_key(api_key: str = Security(api_key_header)):
+    if api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API Key")
+    return api_key
+
+@app.post("/predict-milestone")
+async def predict_milestone(payload: dict, api_key: str = Depends(verify_api_key)):
+    try:
+        # Validate payload
+        if "image" not in payload:
+            raise HTTPException(status_code=400, detail="Image field is required")
+        
+        # Decode base64 image
+        image_data = payload["image"]
+        if image_data.startswith("data:image"):
+            image_data = image_data.split(",")[1]  # Remove data URI prefix
+        img_bytes = base64.b64decode(image_data)
+        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+
+        # Validate image size (max 20MB per TSD)
+        img_bytes_size = len(img_bytes) / (1024 * 1024)
+        if img_bytes_size > 20:
+            raise HTTPException(status_code=400, detail="Image size exceeds 20MB")
+
+        # Validate image type
+        if not img.format.lower() in ["jpeg", "png"]:
+            raise HTTPException(status_code=400, detail="Only JPG/PNG images are supported")
+
+        # Preprocess image
+        max_size = 1024  # Optimize for performance
+        img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+        inputs = processor(images=img, return_tensors="pt")
+
+        # Run inference (within 5 seconds per TSD)
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+            probabilities = torch.softmax(logits, dim=1)
+
+        # Get top prediction
+        predicted_idx = torch.argmax(probabilities, dim=1).item()
+        milestone = model.config.id2label.get(predicted_idx, "Unknown Milestone")
+
+        # Mock percent_complete (since model doesn't provide it)
+        percent_complete = 75  # Adjust based on milestone or train a model to predict this
+
+        return {
+            "milestone": milestone,
+            "percent_complete": percent_complete
+        }
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+
 def process_image(image, project_name):
-    """
-    Process uploaded image, predict construction milestone, and update Salesforce.
-    """
     try:
         # Validate inputs
         if image is None:
@@ -63,7 +124,7 @@ def process_image(image, project_name):
 
         # Preprocess image
         img = Image.open(image).convert("RGB")
-        max_size = 1024  # Optimize for performance
+        max_size = 1024
         img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
         inputs = processor(images=img, return_tensors="pt")
 
@@ -96,7 +157,7 @@ def process_image(image, project_name):
         }
 
         try:
-            project_name = project_name.replace("'", "''")  # Basic escaping
+            project_name = project_name.replace("'", "''")
             query = f"SELECT Id, Version__c FROM Construction_Project__c WHERE Name = '{project_name}'"
             result = sf.query(query)
             if result["totalSize"] > 0:
@@ -119,7 +180,7 @@ def process_image(image, project_name):
     except Exception as e:
         return f"Error: {str(e)}", "Failure", "", ""
 
-# Gradio interface
+# Gradio interface (for manual testing)
 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>")
     project_name = gr.Textbox(label="Project Name", placeholder="Enter project name")
@@ -149,5 +210,5 @@ with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: A
         outputs=progress
     )
 
-if __name__ == "__main__":
-    demo.launch()  # Fixed: Properly indented under the if statement
+# Mount FastAPI app to Gradio
+demo.launch()