Rekham1110's picture
Update app.py
8614e2f verified
raw
history blame
8.11 kB
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()