Update app.py

app.py

@@ -1,9 +1,4 @@
-import streamlit as st
-import requests
-import os
-from PIL import Image
-import numpy as np```python
-import streamlit as st
+import gradio as gr
 import requests
 import os
 from PIL import Image
@@ -17,33 +12,16 @@ from dotenv import load_dotenv
 # Load environment variables from .env file
 load_dotenv()
 
-# Set up page title and layout
-st.title("Construction Project Progress Tracker")
-st.write("Upload a photo of your construction site, and the AI will tell you the progress!")
-
-# Salesforce connection using .env variables
-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:
-    st.error(f"Failed to connect to Salesforce: {str(e)}")
-    st.stop()
-
-# Hugging Face model (placeholder; replace with your custom-trained model)
-# Using a demo image classification model for now
-model = pipeline("image-classification", model="microsoft/resnet-50")
-
 # Function to validate photo size (< 20MB)
 def validate_photo_size(image_file):
     max_size_mb = 20
-    image_file.seek(0, os.SEEK_END)
-    file_size_mb = image_file.tell() / (1024 * 1024)  # Convert bytes to MB
-    image_file.seek(0)  # Reset file pointer
-    return file_size_mb <= max_size_mb
+    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):
@@ -52,10 +30,10 @@ def predict_milestone(image):
         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
-        # Replace with actual milestone mapping based on your trained model
         milestone = predictions[0]["label"]  # Example: "positive" -> "Walls Erected"
         confidence = predictions[0]["score"]
         
@@ -82,18 +60,16 @@ def predict_milestone(image):
         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_file, project_id):
+def upload_image_to_salesforce(image, project_id):
     try:
         # Placeholder: Simulate uploading image to Salesforce ContentVersion
-        # Replace with actual Salesforce file upload API call
-        # Example: Upload to ContentVersion and get ContentDocumentLink
         image_url = f"https://your-salesforce-instance.com/file/{project_id}.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
-def update_salesforce_record(project_id, milestone, percentage, image_url, status, comments):
+def update_salesforce_record(sf, project_id, 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_id}'"
@@ -117,266 +93,88 @@ def update_salesforce_record(project_id, milestone, percentage, image_url, statu
     except Exception as e:
         return f"Failed to update Salesforce: {str(e)}"
 
-# Streamlit UI
-with st.form(key="photo_upload_form"):
-    project_id = st.text_input("Enter Project ID (e.g., Sunshine Apartments)", "Sunshine Apartments")
-    uploaded_file = st.file_uploader("Upload a Construction Photo", type=["jpg", "jpeg", "png"])
-    submit_button = st.form_submit_button(label="Upload and Analyze")
-
-if submit_button and uploaded_file is not None:
-    # Validate photo size
-    if not validate_photo_size(uploaded_file):
-        st.error("Photo is too large! Please upload a photo smaller than 20MB.")
-    else:
-        # Display the uploaded image
-        image = Image.open(uploaded_file)
-        st.image(image, caption="Uploaded Construction Photo", use_column_width=True)
-
-        # Process the image with AI
-        milestone, percentage, error = predict_milestone(image)
-        
-        if error:
-            st.error(error)
-            # Update Salesforce with failure status
-            error_message = update_salesforce_record(
-                project_id=project_id,
-                milestone=None,
-                percentage=0.00,
-                image_url=None,
-                status="Failure",
-                comments=error
-            )
-            if error_message:
-                st.error(error_message)
-        else:
-            # Upload image to Salesforce
-            image_url, upload_error = upload_image_to_salesforce(uploaded_file, project_id)
-            
-            if upload_error:
-                st.error(upload_error)
-                # Update Salesforce with failure status
-                error_message = update_salesforce_record(
-                    project_id=project_id,
-                    milestone=milestone,
-                    percentage=percentage,
-                    image_url=None,
-                    status="Failure",
-                    comments=upload_error
-                )
-                if error_message:
-                    st.error(error_message)
-            else:
-                # Update Salesforce with success
-                st.success(f"AI Result: Milestone = {milestone}, Completion = {percentage}%")
-                error_message = update_salesforce_record(
-                    project_id=project_id,
-                    milestone=milestone,
-                    percentage=percentage,
-                    image_url=image_url,
-                    status="Success",
-                    comments="Photo processed successfully"
-                )
-                if error_message:
-                    st.error(error_message)
-                else:
-                    st.success("Progress saved to Salesforce!")
-else:
-    st.info("Please upload a photo to analyze.")
-```
-
----
-
-#### **2. requirements.txt**
-
-This lists all the Python libraries needed to run the app, including `torch` for Hugging Face and `python-dotenv` for the `.env` file.
-
-<xaiArtifact artifact_id="9acd8b9b-5c22-4fef-a4cd-74849c7b0075" artifact_version_id="afecf7ad-ee42-4754-a130-3b34b8d1a974" title="requirements.txt" contentType="text/plain">
-streamlit==1.39.0
-transformers==4.44.2
-simple-salesforce==1.12.6
-pillow==10.4.0
-numpy==1.26.4
-requests==2.32.3
-torch==2.4.1
-python-dotenv==1.0.1
-from transformers import pipeline
-from simple_salesforce import Salesforceimport streamlit as st
-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()
-
-# Set up page title and layout
-st.title("Construction Project Progress Tracker")
-st.write("Upload a photo of your construction site, and the AI will tell you the progress!")
-
-# Salesforce connection using .env variables
-sf = Salesforce(
-    username=os.getenv('SALESFORCE_USERNAME'),
-    password=os.getenv('SALESFORCE_PASSWORD'),
-    security_token=os.getenv('SALESFORCE_SECURITY_TOKEN'),
-    domain=os.getenv('SALESFORCE_DOMAIN')
-)
-
-# Rest of the app.py code remains the same...
-import io
-import time
-
-# Set up page title and layout
-st.title("Construction Project Progress Tracker")
-st.write("Upload a photo of your construction site, and the AI will tell you the progress!")
-
-# Salesforce connection (replace with your credentials)
-sf = Salesforce(
-    username='your_salesforce_username',
-    password='your_salesforce_password',
-    security_token='your_security_token',
-    domain='login'  # Use 'test' for sandbox
-)
-
-# Hugging Face model (replace with your actual model for construction milestone detection)
-# For demo, we use a placeholder image classification model
-model = pipeline("image-classification", model="microsoft/resnet-50")
-
-# Function to validate photo size (< 20MB)
-def validate_photo_size(image_file):
-    max_size_mb = 20
-    image_file.seek(0, os.SEEK_END)
-    file_size_mb = image_file.tell() / (1024 * 1024)  # Convert bytes to MB
-    image_file.seek(0)  # Reset file pointer
-    return file_size_mb <= max_size_mb
-
-# Function to process image with AI and predict milestone
-def predict_milestone(image):
+# Main Gradio function
+def process_construction_photo(project_id, image):
+    if not project_id or not image:
+        return None, "Please provide a project ID and upload a photo."
+    
+    # Connect to Salesforce
     try:
-        # Simulate AI processing time (ensure < 5 seconds)
-        start_time = time.time()
-        
-        # Process image with Hugging Face model
-        predictions = model(image)
-        
-        # Placeholder logic: Map model output to construction milestones
-        # Replace with actual milestone mapping based on your trained model
-        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
+        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, None, f"AI failed to process the image: {str(e)}"
+        return None, f"Failed to connect to Salesforce: {str(e)}"
 
-# Function to upload image to Salesforce and get a URL
-def upload_image_to_salesforce(image_file, project_id):
-    try:
-        # Placeholder: Simulate uploading image to Salesforce ContentVersion
-        # Replace with actual Salesforce file upload API call
-        image_url = "https://your-salesforce-file-url.com/example-image.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
-def update_salesforce_record(project_id, milestone, percentage, image_url, status, comments):
-    try:
-        sf.Construction_Project__c.update(project_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
-        })
-        return None
-    except Exception as e:
-        return f"Failed to update Salesforce: {str(e)}"
-
-# Streamlit UI
-with st.form(key="photo_upload_form"):
-    project_id = st.text_input("Enter Project ID (e.g., Sunshine Apartments)", "Sunshine Apartments")
-    uploaded_file = st.file_uploader("Upload a Construction Photo", type=["jpg", "jpeg", "png"])
-    submit_button = st.form_submit_button(label="Upload and Analyze")
-
-if submit_button and uploaded_file is not None:
     # Validate photo size
-    if not validate_photo_size(uploaded_file):
-        st.error("Photo is too large! Please upload a photo smaller than 20MB.")
-    else:
-        # Display the uploaded image
-        image = Image.open(uploaded_file)
-        st.image(image, caption="Uploaded Construction Photo", use_column_width=True)
+    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:
+        error_message = update_salesforce_record(
+            sf=sf,
+            project_id=project_id,
+            milestone=None,
+            percentage=0.00,
+            image_url=None,
+            status="Failure",
+            comments=error
+        )
+        return None, f"AI Error: {error}\nSalesforce Error: {error_message}" if error_message else f"AI Error: {error}"
+    
+    # Upload image to Salesforce
+    image_url, upload_error = upload_image_to_salesforce(image, project_id)
+    
+    if upload_error:
+        error_message = update_salesforce_record(
+            sf=sf,
+            project_id=project_id,
+            milestone=milestone,
+            percentage=percentage,
+            image_url=None,
+            status="Failure",
+            comments=upload_error
+        )
+        return None, f"Upload Error: {upload_error}\nSalesforce Error: {error_message}" if error_message else f"Upload Error: {upload_error}"
+    
+    # Update Salesforce with success
+    error_message = update_salesforce_record(
+        sf=sf,
+        project_id=project_id,
+        milestone=milestone,
+        percentage=percentage,
+        image_url=image_url,
+        status="Success",
+        comments="Photo processed successfully"
+    )
+    
+    if error_message:
+        return None, f"Salesforce Error: {error_message}"
+    
+    return image, f"Success! Milestone: {milestone}, Completion: {percentage}%\nProgress saved to Salesforce!"
+
+# Gradio interface
+iface = gr.Interface(
+    fn=process_construction_photo,
+    inputs=[
+        gr.Textbox(label="Project ID (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!"
+)
 
-        # Process the image with AI
-        milestone, percentage, error = predict_milestone(image)
-        
-        if error:
-            st.error(error)
-            # Update Salesforce with failure status
-            error_message = update_salesforce_record(
-                project_id=project_id,
-                milestone=None,
-                percentage=0.00,
-                image_url=None,
-                status="Failure",
-                comments=error
-            )
-            if error_message:
-                st.error(error_message)
-        else:
-            # Upload image to Salesforce
-            image_url, upload_error = upload_image_to_salesforce(uploaded_file, project_id)
-            
-            if upload_error:
-                st.error(upload_error)
-                # Update Salesforce with failure status
-                error_message = update_salesforce_record(
-                    project_id=project_id,
-                    milestone=milestone,
-                    percentage=percentage,
-                    image_url=None,
-                    status="Failure",
-                    comments=upload_error
-                )
-                if error_message:
-                    st.error(error_message)
-            else:
-                # Update Salesforce with success
-                st.success(f"AI Result: Milestone = {milestone}, Completion = {percentage}%")
-                error_message = update_salesforce_record(
-                    project_id=project_id,
-                    milestone=milestone,
-                    percentage=percentage,
-                    image_url=image_url,
-                    status="Success",
-                    comments="Photo processed successfully"
-                )
-                if error_message:
-                    st.error(error_message)
-                else:
-                    st.success("Progress saved to Salesforce!")
-else:
-    st.info("Please upload a photo to analyze.")
+# Launch the Gradio app
+if __name__ == "__main__":
+    iface.launch()