Update app.py

app.py

@@ -1,23 +1,109 @@
-from PIL import ImageFont
+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 logging
+from datetime import datetime
 
-# Update transform
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+
+# Salesforce Credentials
+SALESFORCE_USERNAME = "drone@sathkrutha.com"
+SALESFORCE_PASSWORD = "Komal1303@"
+SALESFORCE_SECURITY_TOKEN = "53AWRskW9EjWUsSL5LU6nFTy3"
+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 = "a003000000xxxxx"  # TODO: Update with actual ID from Site__c
+
+# Connect to Salesforce
+try:
+    sf = Salesforce(
+        username=SALESFORCE_USERNAME,
+        password=SALESFORCE_PASSWORD,
+        security_token=SALESFORCE_SECURITY_TOKEN,
+        instance_url=SALESFORCE_INSTANCE_URL
+    )
+    logging.info("Salesforce connection established.")
+except Exception as e:
+    logging.error(f"Failed to connect to Salesforce: {str(e)}")
+    raise Exception(f"Failed to connect to Salesforce: {str(e)}")
+
+# Load Model
+model = models.detection.fasterrcnn_resnet50_fpn(weights="FasterRCNN_ResNet50_FPN_Weights.COCO_V1")
+model.eval()
+
+# 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',
+    'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
+    'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
+    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
+    'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon',
+    'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
+    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet',
+    'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
+    'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
+    'hair drier', 'toothbrush'
+]
+
+# Image transformations
 transform = transforms.Compose([
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std=[0.229, 0.224, 0.225])
 ])
 
-# Load font for drawing text
-try:
-    font = ImageFont.truetype("arial.ttf", 16)
-except IOError:
-    font = ImageFont.load_default()
+# Map model severity to Salesforce picklist values
+def get_severity(score):
+    if score >= 0.9:
+        return "Critical"
+    elif score >= 0.7:
+        return "Moderate"
+    else:
+        return "Minor"
+
+# Temporary mapping for COCO labels to structural defects
+COCO_TO_DEFECT_MAPPING = {
+    'car': 'Crack',
+    'person': 'Rust',
+    'bicycle': 'Deformation',
+    'truck': 'Corrosion',
+    'boat': 'Spalling',
+}
+
+def map_defect_type(coco_label):
+    return COCO_TO_DEFECT_MAPPING.get(coco_label, "Crack")
 
+# Function to upload image to Salesforce as ContentVersion
+def upload_image_to_salesforce(image, filename="detected_image.jpg", record_id=None):
+    try:
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        img_data = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        
+        content_version = sf.ContentVersion.create({
+            "Title": filename,
+            "PathOnClient": filename,
+            "VersionData": img_data,
+            "FirstPublishLocationId": record_id if record_id else None
+        })
+        logging.info(f"Image uploaded to Salesforce with ContentVersion ID: {content_version['id']}")
+        return content_version["id"]
+    except Exception as e:
+        logging.error(f"Failed to upload image to Salesforce: {str(e)}")
+        raise Exception(f"Failed to upload image to Salesforce: {str(e)}")
+
+# Detect defects and integrate with Salesforce
 def detect_defects(image):
     if not image:
         return None, {"error": "No image provided"}
 
     try:
+        # Perform detection
         image_tensor = transform(image).unsqueeze(0)
         with torch.no_grad():
             predictions = model(image_tensor)
@@ -26,20 +112,14 @@ def detect_defects(image):
         draw = ImageDraw.Draw(result_image)
         output = []
 
-        boxes = predictions[0]['boxes']
-        labels = predictions[0]['labels']
-        scores = predictions[0]['scores']
-
-        for i in range(len(boxes)):
-            score = scores[i].item()
-            if score < 0.5:
+        for i in range(len(predictions[0]['boxes'])):
+            score = predictions[0]['scores'][i].item()
+            if score < 0.5:  # Lowered threshold for testing
                 continue
 
-            box = boxes[i].tolist()
-            label_idx = labels[i].item()
+            box = predictions[0]['boxes'][i].tolist()
+            label_idx = predictions[0]['labels'][i].item()
             coco_label = COCO_INSTANCE_CATEGORY_NAMES[label_idx]
-
-            # Use your custom mapping if available
             defect_type = map_defect_type(coco_label)
             severity = get_severity(score)
 
@@ -50,17 +130,59 @@ def detect_defects(image):
                 "coco_label": coco_label
             })
 
-            # Draw bounding box and label
             draw.rectangle(box, outline="red", width=3)
-            draw.text((box[0], box[1] - 15), f"{defect_type}: {severity}", fill="red", font=font)
+            draw.text((box[0], box[1]), f"{defect_type}: {severity}", fill="red")
 
-        if not output:
-            output = [{"message": "No defects detected above threshold"}]
+        # 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}"
 
-        # Continue with Salesforce upload code as before...
+                # Creating the Salesforce record with updated fields
+                inspection_record = sf.Drone_Structure_Inspection__c.create({
+                    "Inspection_Date__c": current_date,
+                    "Fault_Type__c": output[0]["type"],  # Mapping defect type
+                    "Severity__c": output[0]["severity"],  # Mapping severity
+                    "Fault_Summary__c": str(output),  # Summarizing the defects
+                    "Status__c": "New",  # Default status
+                    "Annotated_Image_URL__c": "",  # Placeholder for image URL
+                    "Report_PDF__c": ""  # Placeholder for report PDF URL
+                })
+
+                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
+                )
+
+                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:
         logging.error(f"Processing failed: {str(e)}")
         return None, {"error": f"Processing failed: {str(e)}"}
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=detect_defects,
+    inputs=gr.Image(type="pil", label="Upload Drone Image"),
+    outputs=[
+        gr.Image(label="Detection Result"),
+        gr.Textbox(label="Detected Faults with Severity")
+    ],
+    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()