app.py

import gradio as gr
import requests
import os
from PIL import Image
import numpy as np
from transformers import pipeline
from simple_salesforce import Salesforce
import io
import time
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# 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:
        # Simulate AI processing time (ensure < 5 seconds)
        start_time = time.time()
        
        # 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:
        return None, None, f"AI failed to process the image: {str(e)}"

# Function to upload image to Salesforce and get a URL
def upload_image_to_salesforce(image, project_name):
    try:
        # 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)}"

# 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)
        
        if result['totalSize'] == 0:
            return None, f"No project found with Name: {project_name}"
        
        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')
        }
        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 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"
    )
    
    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()