neerajkalyank's picture
Update app.py
4b5d09a verified
import gradio as gr
from PIL import Image
import os
from dotenv import load_dotenv
from simple_salesforce import Salesforce
from datetime import datetime
import hashlib
import shutil
import base64
import pytz
# Load environment variables
load_dotenv()
SF_USERNAME = os.getenv("SF_USERNAME")
SF_PASSWORD = os.getenv("SF_PASSWORD")
SF_SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")
# Validate Salesforce credentials
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:
sf = Salesforce(
username=SF_USERNAME,
password=SF_PASSWORD,
security_token=SF_SECURITY_TOKEN,
domain='login'
)
except Exception as e:
print(f"Salesforce connection failed: {str(e)}")
raise
# Valid milestones
VALID_MILESTONES = ["Foundation", "Walls Erected", "Planning", "Completed"]
# Adjust the timezone to your local timezone (replace 'Asia/Kolkata' with your timezone if needed)
local_timezone = pytz.timezone("Asia/Kolkata")
# Deterministic AI prediction with fixed confidence and percent
def mock_ai_model(image):
img = image.convert("RGB")
max_size = 1024
img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
img_bytes = img.tobytes()
img_hash = int(hashlib.sha256(img_bytes).hexdigest(), 16)
milestone_index = img_hash % len(VALID_MILESTONES)
milestone = VALID_MILESTONES[milestone_index]
milestone_completion_map = {
"Planning": 10,
"Foundation": 30,
"Walls Erected": 50,
"Completed": 100,
}
completion_percent = milestone_completion_map.get(milestone, 0)
confidence_raw = 0.85 + ((img_hash % 1000) / 1000) * (0.95 - 0.85)
confidence_score = round(confidence_raw, 2)
return milestone, completion_percent, confidence_score
# Image processing and Salesforce upload
def process_image(image, project_name):
try:
if image is None:
return "Error: Please upload an image to proceed.", "Pending", "", "", 0
img = Image.open(image)
image_size_mb = os.path.getsize(image) / (1024 * 1024)
if image_size_mb > 20:
return "Error: Image size exceeds 20MB.", "Failure", "", "", 0
if not str(image).lower().endswith(('.jpg', '.jpeg', '.png')):
return "Error: Only JPG/PNG images are supported.", "Failure", "", "", 0
# Save image to public folder temporarily before uploading to Salesforce
upload_dir = "public_uploads"
os.makedirs(upload_dir, exist_ok=True)
unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
image_filename = f"{unique_id}_{os.path.basename(image)}"
saved_image_path = os.path.join(upload_dir, image_filename)
shutil.copy(image, saved_image_path)
# Convert image to base64 before uploading to Salesforce
with open(saved_image_path, 'rb') as image_file:
image_data = base64.b64encode(image_file.read()).decode('utf-8')
# Create the ContentVersion record in Salesforce
content_version = {
'Title': image_filename,
'PathOnClient': saved_image_path,
'VersionData': image_data
}
# Upload the file to Salesforce
try:
content_version_result = sf.ContentVersion.create(content_version)
content_version_id = content_version_result['id']
# Generate the public URL for the uploaded file in Salesforce
file_url = f"https://sathkruthatechsolutionspri8-dev-ed.develop.lightning.force.com/{content_version_id}"
except Exception as e:
return f"Error: Failed to upload image to Salesforce - {str(e)}", "Failure", "", "", 0
# AI-based milestone and completion prediction
milestone, percent_complete, confidence_score = mock_ai_model(img)
# Detailed Completion Breakdown (based on detected milestone)
completion_details = {
"Planning": {
"completed": [
"Project outline and goals set, initial designs done."
],
"not_completed": [
"Detailed plans, permits, and contractor hiring pending."
]
},
"Foundation": {
"completed": [
"Foundation work is complete, concrete is poured."
],
"not_completed": [
"Plumbing, electrical groundwork not yet done."
]
},
"Walls Erected": {
"completed": [
"All structural walls are up."
],
"not_completed": [
"Roofing, windows, and internal walls are not yet finished."
]
},
"Completed": {
"completed": [
"All phases of the project are finished, including final touches."
],
"not_completed": [
"There should be no more pending work."
]
}
}
# Get the detailed completion information based on the detected milestone
completed_work = "\n".join([f"{idx+1}. {task}" for idx, task in enumerate(completion_details[milestone]["completed"])])
not_completed_work = "\n".join([f"{idx+1}. {task}" for idx, task in enumerate(completion_details[milestone]["not_completed"])])
# Adjust the current time to local timezone
local_time = datetime.now(local_timezone).strftime("%Y-%m-%dT%H:%M:%S%z") # Corrected ISO 8601 format with timezone
# Create the Salesforce record with the image URL and AI prediction
record = {
"Name__c": project_name,
"Current_Milestone__c": milestone,
"Completion_Percentage__c": percent_complete,
"Last_Updated_On__c": local_time, # Corrected format for Salesforce
"Upload_Status__c": "Success",
"Comments__c": f"{milestone} with {confidence_score*100}% confidence", # Removed "AI Prediction:"
"Last_Updated_Image__c": file_url
}
# Upload the record to Salesforce
try:
sf.Construction__c.create(record)
except Exception as e:
return f"Error: Failed to update Salesforce - {str(e)}", "Failure", "", "", 0
# Return the detailed result in the desired format
result = f"""
Completed:
{completed_work}
Not Completed:
{not_completed_work}
Confidence Score: {confidence_score * 100}%
"""
return result, "Success", milestone, f"Confidence Score: {confidence_score}", percent_complete
except Exception as e:
return f"Error: {str(e)}", "Failure", "", "", 0
# Gradio UI with added styling
with gr.Blocks(css="""
.gradio-container {
background-color: #f0f4f8;
font-family: Arial;
}
.title {
color: #2c3e50;
font-size: 24px;
text-align: center;
font-weight: bold;
}
.gradio-row {
text-align: center;
}
.gradio-container .output {
text-align: center;
}
.gradio-container .input {
text-align: center;
}
.gradio-container .button {
display: block;
margin: 0 auto;
}
""") as demo:
gr.Markdown("<h1 class='title'></h1>")
with gr.Row():
image_input = gr.Image(type="filepath", label="Upload Construction Site Photo (JPG/PNG, ≤ 20MB)")
project_name_input = gr.Textbox(label="Project Name (Required)", placeholder="e.g. Project_12345")
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")
confidence = gr.Textbox(label="Confidence Score")
progress = gr.Slider(0, 100, label="Completion Percentage", interactive=False, value=0)
submit_button.click(
fn=process_image,
inputs=[image_input, project_name_input],
outputs=[output_text, upload_status, milestone, confidence, progress]
)
demo.launch(share=True)