Update app.py

app.py

@@ -1,153 +1,77 @@
 import gradio as gr
-import requests
 from PIL import Image
-import io
-import os
-import re
-from dotenv import load_dotenv  # Fixed: Added this import
-from simple_salesforce import Salesforce
-from datetime import datetime
-
-# Load environment variables
-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")
-
-# Validate Salesforce credentials
-if not all([SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, SF_CONSUMER_KEY, SF_CONSUMER_SECRET]):
-    raise ValueError("Missing Salesforce credentials. Please 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
-
-# Hugging Face Inference API endpoint (replace with your model)
-HF_MODEL_URL = "https://api-inference.huggingface.co/models/nasreshsuguru/construction-milestone-detector"
-
-def process_image(image, project_name):
+from transformers import ViTForImageClassification, ViTImageProcessor
+import torch
+
+# Load pre-trained model and processor from Hugging Face
+model_name = "google/vit-base-patch16-224"
+processor = ViTImageProcessor.from_pretrained(model_name)
+model = ViTForImageClassification.from_pretrained(model_name)
+model.eval()
+
+def process_image(image):
+    """
+    Process uploaded image and predict construction milestone and completion percentage.
+    """
     try:
-        # Validate inputs
+        # Validate image
         if image is None:
-            return "Error: Please upload an image to proceed.", "Pending", "", 0.0, 0.0, ""
-        if not project_name:
-            return "Error: Please enter a project name to proceed.", "Pending", "", 0.0, 0.0, ""
-        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", "", 0.0, 0.0, ""
-
-        # 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", "", 0.0, 0.0, ""
-        if not image.lower().endswith(('.jpg', '.jpeg', '.png')):
-            return "Error: Only JPG/PNG images are supported.", "Failure", "", 0.0, 0.0, ""
+            return "Error: No image uploaded.", "", 0.0, 0.0
 
         # Preprocess image
-        img = Image.open(image)
-        max_size = 1024  # Optimize for performance
-        img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
-        img = img.resize((224, 224))
-        img_byte_arr = io.BytesIO()
-        img.save(img_byte_arr, format="PNG")
-        img_byte_arr = img_byte_arr.getvalue()
-
-        # Call Hugging Face API
-        headers = {"Authorization": f"Bearer {HF_API_KEY}"}
-        response = requests.post(HF_MODEL_URL, headers=headers, data=img_byte_arr, timeout=10)
-
-        if response.status_code != 200:
-            return f"Error: Hugging Face API failed with status {response.status_code}", "Failure", "", 0.0, 0.0, ""
-
-        # Parse AI output
-        result = response.json()
-        top_predictions = sorted(result, key=lambda x: x["score"], reverse=True)[:3]
-        milestone = top_predictions[0]["label"]
-        confidence = top_predictions[0]["score"]
-        percent_complete = min(max(int(confidence * 100), 0), 100)
-        prediction_details = "\n".join([f"{pred['label']}: {pred['score']:.2f}" for pred in top_predictions])
-
-        # Update Salesforce record
-        record = {
-            "Name": project_name,
-            "Current_Milestone__c": milestone,
-            "Completion_Percentage__c": percent_complete,
-            "Last_Updated_On__c": datetime.now().isoformat(),
-            "Upload_Status__c": "Success",
-            "Comments__c": f"AI Confidence: {confidence:.2f}",
-            "Version__c": 1
-        }
-
-        try:
-            project_name = project_name.replace("'", "''")  # Basic escaping
-            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", "", 0.0, 0.0, prediction_details
+        img = Image.open(image).convert("RGB")
+        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)
+
+        # Mocked output parsing (replace with your model's label mapping)
+        # Assuming model outputs classes like 'Foundation Completed', etc.
+        predicted_idx = torch.argmax(probabilities, dim=1).item()
+        confidence = probabilities[0][predicted_idx].item()
+        
+        # Mocked milestone and completion (adjust based on your model)
+        milestone = model.config.id2label.get(predicted_idx, "Unknown Milestone")
+        completion_percent = min(max(int(confidence * 100), 0), 100)  # Mocked logic
 
         return (
-            f"Success: Milestone: {milestone}, Completion: {percent_complete}%",
-            "Success",
+            f"Success: Milestone: {milestone}, Completion: {completion_percent}%",
             milestone,
-            percent_complete,
-            confidence,
-            prediction_details
+            completion_percent,
+            confidence
         )
 
     except Exception as e:
-        return f"Error: {str(e)}", "Failure", "", 0.0, 0.0, ""
+        return f"Error: {str(e)}", "", 0.0, 0.0
 
 # Gradio interface
-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")
+with gr.Blocks() as demo:
+    gr.Markdown("# Construction Milestone Detector")
     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")
     completion_percent = gr.Slider(0, 100, label="Completion Percentage (%)", interactive=False)
     confidence_score = gr.Slider(0, 1, label="Confidence Score", interactive=False)
-    prediction_details = gr.Textbox(label="Top Predictions")
-    progress = gr.Slider(0, 100, label="Processing Progress", interactive=False, value=0)
+    progress = gr.Textbox(label="Processing Status", value="Ready")
 
     def update_progress():
-        return 50
-
-    def complete_progress():
-        return 100
+        return "Processing..."
 
     submit_button.click(
         fn=update_progress,
         outputs=progress
     ).then(
         fn=process_image,
-        inputs=[image_input, project_name],
-        outputs=[output_text, upload_status, milestone, completion_percent, confidence_score, prediction_details]
+        inputs=image_input,
+        outputs=[output_text, milestone, completion_percent, confidence_score]
     ).then(
-        fn=complete_progress,
+        fn=lambda: "Ready",
         outputs=progress
     )
 
 if __name__ == "__main__":
-    demo.launch()
-    
+    demo.launch()