Update app.py

app.py

@@ -2,12 +2,44 @@ import gradio as gr
 from PIL import Image, ImageDraw
 import torch
 from torchvision import models, transforms
+from simple_salesforce import Salesforce
+import base64
+from io import BytesIO
+import os
 
-# ✅ Directly load a pretrained model (you would need to fine-tune this model on structural defects)
+# Salesforce Credentials (replace with your values)
+SALESFORCE_USERNAME = "drone@sathkrutha.com"
+SALESFORCE_PASSWORD = "Komal1303@"
+SALESFORCE_SECURITY_TOKEN = "53AWRskW9EjWUsSL5LU6nFTy3"
+SALESFORCE_INSTANCE_URL = "https://login.salesforce.com"  # or your sandbox URL
+
+# Connect to Salesforce
+try:
+    sf = Salesforce(
+        username=SALESFORCE_USERNAME,
+        password=SALESFORCE_PASSWORD,
+        security_token=SALESFORCE_SECURITY_TOKEN,
+        instance_url=SALESFORCE_INSTANCE_URL
+    )
+    print("Connected to Salesforce successfully!")
+except Exception as e:
+    print(f"Failed to connect to Salesforce: {e}")
+    sf = None
+
+# Load Model (Option 1: Pretrained Faster R-CNN from torchvision)
 model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
 model.eval()
 
-# Define labels (this could be adjusted if you have specific defect labels)
+# Option 2: Load a fine-tuned model from Hugging Face (Placeholder)
+# Example: If you have a fine-tuned Faster R-CNN on Hugging Face for structural defects
+"""
+from transformers import AutoModelForObjectDetection
+model = AutoModelForObjectDetection.from_pretrained("your_huggingface_username/defect-detection-model")
+model.eval()
+# Note: Ensure the model is compatible with torchvision's detection pipeline or adapt the inference logic.
+"""
+
+# Define labels (customize for structural defects if fine-tuned)
 COCO_INSTANCE_CATEGORY_NAMES = [
     '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
     'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter',
@@ -27,7 +59,7 @@ transform = transforms.Compose([
     transforms.ToTensor(),
 ])
 
-# Default severity labels (can be customized based on detection confidence)
+# Severity labels
 def get_severity(score):
     if score >= 0.9:
         return "High"
@@ -36,25 +68,47 @@ def get_severity(score):
     else:
         return "Low"
 
+# Function to upload image to Salesforce as a ContentVersion (file)
+def upload_image_to_salesforce(image, filename="detected_image.jpg"):
+    if not sf:
+        return None
+    try:
+        # Convert PIL image to base64
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        img_data = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        
+        # Create ContentVersion (Salesforce File)
+        content_version = sf.ContentVersion.create({
+            "Title": filename,
+            "PathOnClient": filename,
+            "VersionData": img_data,
+            "FirstPublishLocationId": None  # Optionally link to a record
+        })
+        return content_version["id"]
+    except Exception as e:
+        print(f"Failed to upload image to Salesforce: {e}")
+        return None
+
+# Detect defects and integrate with Salesforce
 def detect_defects(image):
+    # Perform detection
     image_tensor = transform(image).unsqueeze(0)
-
     with torch.no_grad():
         predictions = model(image_tensor)
 
     result_image = image.copy()
     draw = ImageDraw.Draw(result_image)
-
     output = []
+
     for i in range(len(predictions[0]['boxes'])):
         score = predictions[0]['scores'][i].item()
-        if score < 0.7:  # filter low-confidence predictions
+        if score < 0.7:  # Filter low-confidence predictions
             continue
 
         box = predictions[0]['boxes'][i].tolist()
         label_idx = predictions[0]['labels'][i].item()
         label = COCO_INSTANCE_CATEGORY_NAMES[label_idx]
-
         severity = get_severity(score)
 
         output.append({
@@ -66,6 +120,23 @@ def detect_defects(image):
         draw.rectangle(box, outline="red", width=3)
         draw.text((box[0], box[1]), f"{label}: {severity}", fill="red")
 
+    # Upload image to Salesforce
+    content_version_id = upload_image_to_salesforce(result_image, filename="detected_defect.jpg")
+
+    # Store detection results in Salesforce
+    if sf and output:
+        try:
+            for defect in output:
+                sf.Defect__c.create({
+                    "Defect_Type__c": defect["type"],
+                    "Confidence__c": defect["confidence"],
+                    "Severity__c": defect["severity"],
+                    "Image_URL__c": f"ContentVersion ID: {content_version_id}" if content_version_id else ""
+                })
+            print("Defect records created in Salesforce!")
+        except Exception as e:
+            print(f"Failed to create defect records: {e}")
+
     return result_image, output
 
 # Gradio Interface
@@ -76,9 +147,9 @@ demo = gr.Interface(
         gr.Image(label="Detection Result"),
         gr.JSON(label="Detected Faults with Severity")
     ],
-    title="Structural Defect Detection (Demo)",
-    description="Detects general objects using Faster R-CNN pretrained on COCO. This model needs fine-tuning for structural defects like cracks, rust, and spalling."
+    title="Structural Defect Detection with Salesforce Integration",
+    description="Detects objects using Faster R-CNN and stores results in Salesforce. Fine-tune the model for structural defects like cracks, rust, and spalling."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()