Update app.py

app.py

@@ -1,14 +1,9 @@
 import gradio as gr
 from PIL import Image
 import os
-import re  # Added for regular expression validation
 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
 import random  # For mock predictions
 
 # Load environment variables
@@ -17,11 +12,9 @@ SF_USERNAME = os.getenv("SF_USERNAME")
 SF_PASSWORD = os.getenv("SF_PASSWORD")
 SF_SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")
 
-API_KEY = os.getenv("API_KEY", "your-api-key-here")
-
 # Validate Salesforce credentials
-if not all([SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN ]):
-    raise ValueError("Missing Salesforce credentials. Set SF_USERNAME, SF_PASSWORD SF_SECURITY_TOKEN  in environment variables.")
+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:
@@ -29,23 +22,12 @@ try:
         username=SF_USERNAME,
         password=SF_PASSWORD,
         security_token=SF_SECURITY_TOKEN,
-        
         domain='login'  # Use 'test' for sandbox
     )
 except Exception as e:
     print(f"Salesforce connection failed: {str(e)}")
     raise
 
-# 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
-
 # Mock AI model for milestone detection (since we can't train a real model here)
 def mock_ai_model(image):
     # Preprocessing: Resize, normalize (simulated)
@@ -71,69 +53,32 @@ def mock_ai_model(image):
 
     return milestone, completion_percent, confidence_score
 
-@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))
-
-        # Validate image size (max 20MB)
-        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")
-
-        # Run mock AI model
-        milestone, percent_complete, confidence_score = mock_ai_model(img)
-
-        return {
-            "milestone": milestone,
-            "percent_complete": percent_complete,
-            "confidence_score": confidence_score
-        }
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
-
 # Function for Gradio UI to process the image
-def process_image(image, project_name):
+def process_image(image):
     try:
         # Validate inputs
         if image is None:
             return "Error: Please upload an image to proceed.", "Pending", "", "", 0
-        if not project_name:
-            return "Error: Please enter a project name to proceed.", "Pending", "", "", 0
-        # Updated validation to allow letters, numbers, spaces, and hyphens
-        if not re.match(r'^[a-zA-Z0-9\s-]+$', project_name):
-            return "Error: Project name must be alphanumeric (letters, numbers, spaces, or hyphens).", "Pending", "", "", 0
 
         # Open and validate image
         img = Image.open(image)
 
         # Validate image size and type
         image_size_mb = os.path.getsize(image) / (1024 * 1024)
-        if image_size_mb_tb > 20:
+        if image_size_mb > 20:  # Fixed the typo: image_size_mb_tb to image_size_mb
             return "Error: Image size exceeds 20MB.", "Failure", "", "", 0
-        if not image.lower().endswith(('.jpg', '.jpeg', '.png')):
+        if not str(image).lower().endswith(('.jpg', '.jpeg', '.png')):
             return "Error: Only JPG/PNG images are supported.", "Failure", "", "", 0
 
         # Run mock AI model
         milestone, percent_complete, confidence_score = mock_ai_model(img)
 
-        # Update Salesforce record
+        # Generate a unique identifier for the Salesforce record
+        unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
+
+        # Store result in Salesforce
         record = {
-            "Name__c": project_name,  # Changed from "Name" to "Name__c"
+            "Name__c": unique_id,  # Use a timestamp-based unique ID instead of project name
             "Current_Milestone__c": milestone,
             "Completion_Percentage__c": percent_complete,
             "Last_Updated_On__c": datetime.now().isoformat(),
@@ -142,15 +87,7 @@ def process_image(image, project_name):
         }
 
         try:
-            # Escape single quotes in project_name to prevent SQL injection
-            escaped_project_name = project_name.replace("'", "''")
-            query = f"SELECT Id FROM Construction__c WHERE Name__c = '{escaped_project_name}'"
-            result = sf.query(query)
-            if result["totalSize"] > 0:
-                project_id = result["records"][0]["Id"]
-                sf.Construction__c.update(project_id, record)
-            else:
-                sf.Construction__c.create(record)
+            sf.Construction__c.create(record)  # Always create a new record
         except Exception as e:
             return f"Error: Failed to update Salesforce - {str(e)}", "Failure", "", "", 0
 
@@ -168,7 +105,6 @@ def process_image(image, project_name):
 # Gradio interface for 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 (e.g., My House or Site-123)")
     image_input = gr.Image(type="filepath", label="Upload Construction Site Photo (JPG/PNG, ≤ 20MB)")
     submit_button = gr.Button("Process Image")
     output_text = gr.Textbox(label="Result")
@@ -179,9 +115,9 @@ with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: A
 
     submit_button.click(
         fn=process_image,
-        inputs=[image_input, project_name],
+        inputs=[image_input],
         outputs=[output_text, upload_status, milestone, confidence, progress]
     )
 
 # Launch the Gradio app
-demo.launch()
+demo.launch(share=True)  # Note: share=True may not work in all environments; you may need to run this locally