Update app.py

app.py

@@ -12,7 +12,10 @@ from datetime import datetime
 SALESFORCE_USERNAME = "drone@sathkrutha.com"
 SALESFORCE_PASSWORD = "Komal1303@"
 SALESFORCE_SECURITY_TOKEN = "53AWRskW9EjWUsSL5LU6nFTy3"
-SALESFORCE_INSTANCE_URL = "https://sathikrutha-a-dev-ed.my.salesforce.com"  # Updated to your Developer Edition org
+SALESFORCE_INSTANCE_URL = "https://sathikrutha-a-dev-ed.my.salesforce.com"
+
+# Replace with a valid Site__c record ID from your Salesforce org
+SITE_RECORD_ID = "a00xxxxxxxxxxxx"  # TODO: Update with actual ID
 
 # Connect to Salesforce
 try:
@@ -22,16 +25,14 @@ try:
         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
+    raise Exception(f"Failed to connect to Salesforce: {str(e)}")
 
-# Load Model (Updated to use weights instead of deprecated pretrained parameter)
+# Load Model
 model = models.detection.fasterrcnn_resnet50_fpn(weights="FasterRCNN_ResNet50_FPN_Weights.COCO_V1")
 model.eval()
 
-# Define labels (COCO labels; ideally, fine-tune for structural defects)
+# Define labels (COCO labels; 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',
@@ -54,117 +55,112 @@ transform = transforms.Compose([
 # Map model severity to Salesforce picklist values
 def get_severity(score):
     if score >= 0.9:
-        return "Critical"  # Maps "High" to "Critical"
+        return "Critical"
     elif score >= 0.7:
-        return "Moderate"  # Maps "Medium" to "Moderate"
+        return "Moderate"
     else:
-        return "Minor"     # Maps "Low" to "Minor"
+        return "Minor"
 
-# Temporary mapping for COCO labels to structural defects (for demo purposes)
-# In a production scenario, fine-tune the model to detect actual defects
+# Temporary mapping for COCO labels to structural defects
 COCO_TO_DEFECT_MAPPING = {
-    'car': 'Crack',        # Example mapping; adjust based on your needs
+    'car': 'Crack',
     '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
+    return COCO_TO_DEFECT_MAPPING.get(coco_label, "Crack")
 
-# Function to upload image to Salesforce as a ContentVersion (file)
+# Function to upload image to Salesforce as ContentVersion
 def upload_image_to_salesforce(image, filename="detected_image.jpg", record_id=None):
-    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": record_id if record_id else None  # Link to the record
+            "FirstPublishLocationId": record_id if record_id else None
         })
         return content_version["id"]
     except Exception as e:
-        print(f"Failed to upload image to Salesforce: {e}")
-        return None
+        raise Exception(f"Failed to upload image to Salesforce: {str(e)}")
 
 # 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
-            continue
-
-        box = predictions[0]['boxes'][i].tolist()
-        label_idx = predictions[0]['labels'][i].item()
-        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": defect_type,
-            "confidence": round(score, 2),
-            "severity": severity,
-            "coco_label": coco_label  # Include original label for reference
-        })
+    if not image:
+        return None, {"error": "No image provided"}
 
-        draw.rectangle(box, outline="red", width=3)
-        draw.text((box[0], box[1]), f"{defect_type}: {severity}", fill="red")
-
-    # Create a single record for the inspection
-    if sf and output:
-        try:
-            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
+    try:
+        # 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:
+                continue
+
+            box = predictions[0]['boxes'][i].tolist()
+            label_idx = predictions[0]['labels'][i].item()
+            coco_label = COCO_INSTANCE_CATEGORY_NAMES[label_idx]
+            defect_type = map_defect_type(coco_label)
+            severity = get_severity(score)
+
+            output.append({
+                "type": defect_type,
+                "confidence": round(score, 2),
+                "severity": severity,
+                "coco_label": coco_label
             })
-            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}"
+
+            draw.rectangle(box, outline="red", width=3)
+            draw.text((box[0], box[1]), f"{defect_type}: {severity}", fill="red")
+
+        # Create Salesforce record if detections exist
+        if output:
+            try:
+                current_date = datetime.now().strftime("%Y-%m-%d")
+                inspection_name = f"Inspection-{current_date}-{len(output):03d}"
+
+                inspection_record = sf.Drone_Structure_Inspection__c.create({
+                    "Site__c": SITE_RECORD_ID,
+                    "Fault_Type__c": output[0]["type"],
+                    "Severity__c": output[0]["severity"],
+                    "Fault_Summary__c": json.dumps(output),
+                    "Inspection_Date__c": current_date,
+                    "Status__c": "New"
                 })
 
-        except Exception as e:
-            print(f"Failed to create inspection record: {e}")
+                record_id = inspection_record.get("id")
+                content_version_id = upload_image_to_salesforce(
+                    result_image,
+                    filename=f"detected_defect_{record_id}.jpg",
+                    record_id=record_id
+                )
 
-    return result_image, output
+                if content_version_id:
+                    sf.Drone_Structure_Inspection__c.update(record_id, {
+                        "Annotated_Image_URL__c": f"/sfc/servlet.shepherd/version/download/{content_version_id}"
+                    })
+
+                output.append({"salesforce_record_id": record_id})
+            except Exception as e:
+                output.append({"error": f"Failed to create Salesforce record: {str(e)}"})
+
+        return result_image, output
+
+    except Exception as e:
+        return None, {"error": f"Processing failed: {str(e)}"}
 
 # Gradio Interface
 demo = gr.Interface(