Update app.py

app.py

@@ -1,214 +1,214 @@
 import gradio as gr
-from PIL import Image
-import torch
-from transformers import ViTForImageClassification, ViTImageProcessor
+import requests
 import os
-import re
-from dotenv import load_dotenv
+from PIL import Image
+import numpy as np
+from transformers import pipeline
 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 time
+from dotenv import load_dotenv
 
-# Load environment variables
+# Load environment variables from .env file
 load_dotenv()
-HF_API_KEY = os.getenv("HF_API_KEY")
-SF_USERNAME = os.getenv("SF_USERNAME")
-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 (for Gradio UI)
-if not all([SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, SF_CONSUMER_KEY, SF_CONSUMER_SECRET]):
-    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:
-    sf = Salesforce(
-        username=SF_USERNAME,
-        password=SF_PASSWORD,
-        security_token=SF_SECURITY_TOKEN,
-        consumer_key=SF_CONSUMER_KEY,
-        consumer_secret=SF_CONSUMER_SECRET,
-        domain='login'  # Use 'test' for sandbox
-    )
-except Exception as e:
-    print(f"Salesforce connection failed: {str(e)}")
-    raise
-
-# Load pre-trained model and processor
-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)):
+# Function to validate photo size (< 20MB)
+def validate_photo_size(image_file):
+    max_size_mb = 20
+    if isinstance(image_file, Image.Image):
+        # Convert PIL Image to bytes for size check
+        img_byte_arr = io.BytesIO()
+        image_file.save(img_byte_arr, format='JPEG')
+        file_size_mb = img_byte_arr.tell() / (1024 * 1024)  # Convert bytes to MB
+        return file_size_mb <= max_size_mb, None
+    return False, "Invalid image format"
+
+# Function to process image with AI and predict milestone
+def predict_milestone(image):
     try:
-        # Validate payload
-        if "image" not in payload:
-            raise HTTPException(status_code=400, detail="Image field is required")
+        # Simulate AI processing time (ensure < 5 seconds)
+        start_time = time.time()
         
-        # 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
+        # Process image with Hugging Face model
+        model = pipeline("image-classification", model="microsoft/resnet-50")
+        predictions = model(image)
+        
+        # Placeholder logic: Map model output to construction milestones
+        milestone = predictions[0]["label"]  # Example: "positive" -> "Walls Erected"
+        confidence = predictions[0]["score"]
+        
+        # Map model output to construction milestones (customize this)
+        milestone_map = {
+            "positive": "Walls Erected",
+            "negative": "Foundation Completed",
+            # Add more mappings based on your model
         }
-
+        completion_map = {
+            "positive": 60.00,  # Example: Walls = 60% complete
+            "negative": 20.00,  # Example: Foundation = 20% complete
+        }
+        
+        predicted_milestone = milestone_map.get(milestone, "Unknown Milestone")
+        completion_percentage = completion_map.get(milestone, 0.00)
+        
+        processing_time = time.time() - start_time
+        if processing_time > 5:
+            return None, None, "AI took too long to process (> 5 seconds)."
+        
+        return predicted_milestone, completion_percentage, None
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+        return None, None, f"AI failed to process the image: {str(e)}"
 
-def process_image(image, project_name):
+# Function to upload image to Salesforce and get a URL
+def upload_image_to_salesforce(image, project_name):
     try:
-        # Validate inputs
-        if image is None:
-            return "Error: Please upload an image to proceed.", "Pending", "", ""
-        if not project_name:
-            return "Error: Please enter a project name to proceed.", "Pending", "", ""
-        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", "", ""
-
-        # Validate image size and type
-        image_size_mb = os.path.getsize(image) / (1024 * 1024)
-        if image_size_mb > 20:
-            return "Error: Image size exceeds 20MB.", "Failure", "", ""
-        if not image.lower().endswith(('.jpg', '.jpeg', '.png')):
-            return "Error: Only JPG/PNG images are supported.", "Failure", "", ""
-
-        # Preprocess image
-        img = Image.open(image).convert("RGB")
-        max_size = 1024
-        img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
-        inputs = processor(images=img, return_tensors="pt")
-
-        # Run inference
-        with torch.no_grad():
-            outputs = model(**inputs)
-            logits = outputs.logits
-            probabilities = torch.softmax(logits, dim=1)
+        # Placeholder: Simulate uploading image to Salesforce ContentVersion
+        image_url = f"https://your-salesforce-instance.com/file/{project_name}.jpg"  # Simulated URL
+        return image_url, None
+    except Exception as e:
+        return None, f"Failed to upload image to Salesforce: {str(e)}"
 
-        # Get top predictions
-        top_probs, top_indices = torch.topk(probabilities, 3, dim=1)
-        top_probs = top_probs[0].tolist()
-        top_indices = top_indices[0].tolist()
+# Function to update Salesforce Construction_Project__c object and fetch fields
+def update_salesforce_record(sf, project_name, milestone, percentage, image_url, status, comments):
+    try:
+        # Query to check if the project exists
+        query = f"SELECT Id FROM Construction_Project__c WHERE Name = '{project_name}'"
+        result = sf.query(query)
         
-        # Map indices to labels
-        predicted_idx = top_indices[0]
-        milestone = model.config.id2label.get(predicted_idx, "Unknown Milestone")
+        if result['totalSize'] == 0:
+            return None, f"No project found with Name: {project_name}"
         
-        # Format top predictions
-        prediction_details = "\n".join([f"{model.config.id2label.get(idx, 'Unknown Milestone')}: {prob:.2f}" for idx, prob in zip(top_indices, top_probs)])
-
-        # Update Salesforce record
-        record = {
-            "Name": project_name,
-            "Current_Milestone__c": milestone,
-            "Last_Updated_On__c": datetime.now().isoformat(),
-            "Upload_Status__c": "Success",
-            "Comments__c": f"AI Prediction: {milestone}",
-            "Version__c": 1
+        record_id = result['records'][0]['Id']
+        
+        # Update the record
+        sf.Construction_Project__c.update(record_id, {
+            'Current_Milestone__c': milestone,
+            'Completion_Percentage__c': percentage,
+            'Last_Updated_Image__c': image_url,
+            'Last_Updated_On__c': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
+            'Upload_Status__c': status,
+            'Comments__c': comments
+        })
+        
+        # Fetch the updated record to get the specified fields
+        updated_query = f"SELECT Current_Milestone__c, Last_Updated_Image__c, Last_Updated_On__c, Upload_Status__c FROM Construction_Project__c WHERE Id = '{record_id}'"
+        updated_result = sf.query(updated_query)
+        
+        if updated_result['totalSize'] == 0:
+            return None, "Failed to retrieve updated record."
+        
+        record = updated_result['records'][0]
+        fields_output = {
+            'Current_Milestone__c': record.get('Current_Milestone__c', 'N/A'),
+            'Last_Updated_Image__c': record.get('Last_Updated_Image__c', 'N/A'),
+            'Last_Updated_On__c': record.get('Last_Updated_On__c', 'N/A'),
+            'Upload_Status__c': record.get('Upload_Status__c', 'N/A')
         }
-
-        try:
-            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:
-                project_id = result["records"][0]["Id"]
-                current_version = result["records"][0].get("Version__c", 0)
-                record["Version__c"] = current_version + 1
-                sf.Construction_Project__c.update(project_id, record)
-            else:
-                sf.Construction_Project__c.create(record)
-        except Exception as e:
-            return f"Error: Failed to update Salesforce - {str(e)}", "Failure", "", prediction_details
-
-        return (
-            f"Success: Milestone: {milestone}",
-            "Success",
-            milestone,
-            prediction_details
+        return fields_output, None
+    except Exception as e:
+        return None, f"Failed to update Salesforce: {str(e)}"
+
+# Main Gradio function
+def process_construction_photo(project_name, image):
+    if not project_name or not image:
+        return None, "Please provide a project name and upload a photo."
+    
+    # Connect to Salesforce
+    try:
+        sf = Salesforce(
+            username=os.getenv('SALESFORCE_USERNAME'),
+            password=os.getenv('SALESFORCE_PASSWORD'),
+            security_token=os.getenv('SALESFORCE_SECURITY_TOKEN'),
+            domain=os.getenv('SALESFORCE_DOMAIN')
         )
-
     except Exception as e:
-        return f"Error: {str(e)}", "Failure", "", ""
-
-# 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")
-    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")
-    upload_status = gr.Textbox(label="Upload Status")
-    milestone = gr.Textbox(label="Detected Milestone")
-    prediction_details = gr.Textbox(label="Top Predictions")
-    progress = gr.Slider(0, 100, label="Processing Progress", interactive=False, value=0)
-
-    def update_progress():
-        return 50
-
-    def complete_progress():
-        return 100
-
-    submit_button.click(
-        fn=update_progress,
-        outputs=progress
-    ).then(
-        fn=process_image,
-        inputs=[image_input, project_name],
-        outputs=[output_text, upload_status, milestone, prediction_details]
-    ).then(
-        fn=complete_progress,
-        outputs=progress
+        return None, f"Failed to connect to Salesforce: {str(e)}"
+
+    # Validate photo size
+    is_valid, error = validate_photo_size(image)
+    if not is_valid:
+        return None, error or "Photo is too large! Please upload a photo smaller than 20MB."
+
+    # Process the image with AI
+    milestone, percentage, error = predict_milestone(image)
+    
+    if error:
+        fields, error_message = update_salesforce_record(
+            sf=sf,
+            project_name=project_name,
+            milestone=None,
+            percentage=0.00,
+            image_url=None,
+            status="Failure",
+            comments=error
+        )
+        error_text = f"AI Error: {error}"
+        if error_message:
+            error_text += f"\nSalesforce Error: {error_message}"
+        if fields:
+            error_text += "\nUpdated Salesforce Fields:\n"
+            for field, value in fields.items():
+                error_text += f"{field}: {value}\n"
+        return None, error_text
+    
+    # Upload image to Salesforce
+    image_url, upload_error = upload_image_to_salesforce(image, project_name)
+    
+    if upload_error:
+        fields, error_message = update_salesforce_record(
+            sf=sf,
+            project_name=project_name,
+            milestone=milestone,
+            percentage=percentage,
+            image_url=None,
+            status="Failure",
+            comments=upload_error
+        )
+        error_text = f"Upload Error: {upload_error}"
+        if error_message:
+            error_text += f"\nSalesforce Error: {error_message}"
+        if fields:
+            error_text += "\nUpdated Salesforce Fields:\n"
+            for field, value in fields.items():
+                error_text += f"{field}: {value}\n"
+        return None, error_text
+    
+    # Update Salesforce with success
+    fields, error_message = update_salesforce_record(
+        sf=sf,
+        project_name=project_name,
+        milestone=milestone,
+        percentage=percentage,
+        image_url=image_url,
+        status="Success",
+        comments="Photo processed successfully"
     )
-
-# Mount FastAPI app to Gradio
-demo.launch()
+    
+    if error_message:
+        return None, f"Salesforce Error: {error_message}"
+    
+    # Prepare output with AI results and Salesforce fields
+    result_text = f"Success! Milestone: {milestone}, Completion: {percentage}%\nProgress saved to Salesforce!\n\nSalesforce Fields:\n"
+    for field, value in fields.items():
+        result_text += f"{field}: {value}\n"
+    
+    return image, result_text
+
+# Gradio interface
+iface = gr.Interface(
+    fn=process_construction_photo,
+    inputs=[
+        gr.Textbox(label="Project Name (e.g., Sunshine Apartments)", placeholder="Sunshine Apartments"),
+        gr.Image(type="pil", label="Upload a Construction Photo")
+    ],
+    outputs=[
+        gr.Image(label="Uploaded Photo"),
+        gr.Textbox(label="Result")
+    ],
+    title="Construction Project Progress Tracker",
+    description="Upload a photo of your construction site, and the AI will tell you the progress!"
+)
+
+# Launch the Gradio app
+if __name__ == "__main__":
+    iface.launch()