Update app.py

app.py

@@ -5,13 +5,14 @@ from torchvision import models, transforms
 from simple_salesforce import Salesforce
 import base64
 from io import BytesIO
-import os
+import json
+from datetime import datetime
 
-# Salesforce Credentials (replace with your values)
+# Salesforce Credentials
 SALESFORCE_USERNAME = "drone@sathkrutha.com"
 SALESFORCE_PASSWORD = "Komal1303@"
 SALESFORCE_SECURITY_TOKEN = "53AWRskW9EjWUsSL5LU6nFTy3"
-SALESFORCE_INSTANCE_URL = "https://login.salesforce.com"  # or your sandbox URL
+SALESFORCE_INSTANCE_URL = "https://sathikrutha-a-dev-ed.my.salesforce.com"  # Updated to your Developer Edition org
 
 # Connect to Salesforce
 try:
@@ -26,20 +27,11 @@ 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)
+# Load Model (Updated to use weights instead of deprecated pretrained parameter)
+model = models.detection.fasterrcnn_resnet50_fpn(weights="FasterRCNN_ResNet50_FPN_Weights.COCO_V1")
 model.eval()
 
-# 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)
+# Define labels (COCO labels; ideally, fine-tune for structural defects)
 COCO_INSTANCE_CATEGORY_NAMES = [
     '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
     'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter',
@@ -59,17 +51,32 @@ transform = transforms.Compose([
     transforms.ToTensor(),
 ])
 
-# Severity labels
+# Map model severity to Salesforce picklist values
 def get_severity(score):
     if score >= 0.9:
-        return "High"
+        return "Critical"  # Maps "High" to "Critical"
     elif score >= 0.7:
-        return "Medium"
+        return "Moderate"  # Maps "Medium" to "Moderate"
     else:
-        return "Low"
+        return "Minor"     # Maps "Low" to "Minor"
+
+# Temporary mapping for COCO labels to structural defects (for demo purposes)
+# In a production scenario, fine-tune the model to detect actual defects
+COCO_TO_DEFECT_MAPPING = {
+    'car': 'Crack',        # Example mapping; adjust based on your needs
+    'person': 'Rust',
+    'bicycle': 'Deformation',
+    'truck': 'Corrosion',
+    'boat': 'Spalling',
+    # Add more mappings as needed for other COCO labels
+}
+
+# Function to map COCO labels to Salesforce picklist values
+def map_defect_type(coco_label):
+    return COCO_TO_DEFECT_MAPPING.get(coco_label, "Crack")  # Default to "Crack" if no mapping
 
 # Function to upload image to Salesforce as a ContentVersion (file)
-def upload_image_to_salesforce(image, filename="detected_image.jpg"):
+def upload_image_to_salesforce(image, filename="detected_image.jpg", record_id=None):
     if not sf:
         return None
     try:
@@ -83,7 +90,7 @@ def upload_image_to_salesforce(image, filename="detected_image.jpg"):
             "Title": filename,
             "PathOnClient": filename,
             "VersionData": img_data,
-            "FirstPublishLocationId": None  # Optionally link to a record
+            "FirstPublishLocationId": record_id if record_id else None  # Link to the record
         })
         return content_version["id"]
     except Exception as e:
@@ -108,34 +115,54 @@ def detect_defects(image):
 
         box = predictions[0]['boxes'][i].tolist()
         label_idx = predictions[0]['labels'][i].item()
-        label = COCO_INSTANCE_CATEGORY_NAMES[label_idx]
+        coco_label = COCO_INSTANCE_CATEGORY_NAMES[label_idx]
+        defect_type = map_defect_type(coco_label)  # Map COCO label to defect type
         severity = get_severity(score)
 
         output.append({
-            "type": label,
+            "type": defect_type,
             "confidence": round(score, 2),
-            "severity": severity
+            "severity": severity,
+            "coco_label": coco_label  # Include original label for reference
         })
 
         draw.rectangle(box, outline="red", width=3)
-        draw.text((box[0], box[1]), f"{label}: {severity}", fill="red")
+        draw.text((box[0], box[1]), f"{defect_type}: {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
+    # Create a single record for the inspection
     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 ""
+            current_date = datetime.now().strftime("%Y-%m-%d")
+            inspection_name = f"Inspection-{current_date}-{len(output):03d}"
+
+            # Create the inspection record
+            inspection_record = sf.Drone_Structure_Inspection__c.create({
+                "Name": inspection_name,
+                "Fault_Type__c": output[0]["type"],  # Use the mapped defect type
+                "Severity__c": output[0]["severity"],  # Use the mapped severity
+                "Fault_Summary__c": json.dumps(output),  # Store all defects as JSON
+                "Inspection_Date__c": current_date,
+                "Status__c": "New"  # Default status
+                # Note: Site__c and Report_PDF__c are not set as they require additional input
+            })
+            print("Inspection record created in Salesforce!")
+
+            # Upload the annotated image and link it to the record
+            record_id = inspection_record.get("id")
+            content_version_id = upload_image_to_salesforce(
+                result_image,
+                filename="detected_defect.jpg",
+                record_id=record_id
+            )
+
+            # Update the record with the ContentVersion ID
+            if content_version_id:
+                sf.Drone_Structure_Inspection__c.update(record_id, {
+                    "Annotated_Image_URL__c": f"ContentVersion ID: {content_version_id}"
                 })
-            print("Defect records created in Salesforce!")
+
         except Exception as e:
-            print(f"Failed to create defect records: {e}")
+            print(f"Failed to create inspection record: {e}")
 
     return result_image, output