Update app.py

app.py

@@ -1,28 +1,138 @@
+import gradio as gr
+from PIL import Image, ImageDraw, ImageFont
+import torch
+from torchvision import models, transforms
+from simple_salesforce import Salesforce
+import base64
+from io import BytesIO
+import logging
+from datetime import datetime
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+
+# Salesforce Credentials (replace with your own or environment variables)
+SALESFORCE_USERNAME = "drone@sathkrutha.com"
+SALESFORCE_PASSWORD = "Komal1303@"
+SALESFORCE_SECURITY_TOKEN = "53AWRskW9EjWUsSL5LU6nFTy3"
+SALESFORCE_INSTANCE_URL = "https://sathikrutha-a-dev-ed.my.salesforce.com"
+
+# Salesforce Site or parent record ID where content will be linked
+SITE_RECORD_ID = "a003000000xxxxx"  # TODO: Replace with actual Site__c record ID
+
+# 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 the Faster R-CNN pretrained model (replace with your fine-tuned weights if any)
+model = models.detection.fasterrcnn_resnet50_fpn(weights="FasterRCNN_ResNet50_FPN_Weights.COCO_V1")
+model.eval()
+
+# COCO categories list for mapping labels (standard)
+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 transformation for the model input
+transform = transforms.Compose([
+    transforms.ToTensor(),
+])
+
+# Map confidence score to severity level
+def get_severity(score):
+    if score >= 0.9:
+        return "Critical"
+    elif score >= 0.7:
+        return "Moderate"
+    else:
+        return "Minor"
+
+# Temporary mapping from COCO labels to defect types (replace with your own)
+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")
+
+# Upload annotated image to Salesforce as ContentVersion record
+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 SITE_RECORD_ID
+        })
+        logging.info(f"Image uploaded to Salesforce 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)}")
+
+# Main defect detection and Salesforce integration function
 def detect_defects(image):
-    if not image:
-        return None, {"error": "No image provided"}
+    if image is None:
+        return None, "No image provided"
 
     try:
-        # Perform detection
+        # Convert image and prepare input tensor
         image_tensor = transform(image).unsqueeze(0)
+
         with torch.no_grad():
             predictions = model(image_tensor)
 
+        # Copy original image to draw bounding boxes and labels
         result_image = image.copy()
         draw = ImageDraw.Draw(result_image)
-        output = []
 
+        # Optional: Use a truetype font for nicer text, fallback if not available
+        try:
+            font = ImageFont.truetype("arial.ttf", 18)
+        except:
+            font = ImageFont.load_default()
+
+        output = []
         for i in range(len(predictions[0]['boxes'])):
             score = predictions[0]['scores'][i].item()
-            if score < 0.5:  # Increased threshold for better quality detections
+            if score < 0.3:  # confidence threshold
                 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)
 
+            # Append defect info to output list
             output.append({
                 "type": defect_type,
                 "confidence": round(score, 2),
@@ -30,34 +140,29 @@ def detect_defects(image):
                 "coco_label": coco_label
             })
 
+            # Draw rectangle and label
             draw.rectangle(box, outline="red", width=3)
-            draw.text((box[0], box[1]), f"{defect_type}: {severity}", fill="red")
-
-        # Apply NMS to filter overlapping detections
-        boxes = predictions[0]['boxes']
-        scores = predictions[0]['scores']
-        keep = apply_nms(boxes, scores, threshold=0.5)
-        boxes = boxes[keep]
-        output = [output[i] for i in keep]
+            text = f"{defect_type}: {severity}"
+            draw.text((box[0], box[1] - 20 if box[1] > 20 else box[1],), text, fill="red", font=font)
 
-        # Create Salesforce record if detections exist
+        # If defects found, create Salesforce record & upload annotated image
         if output:
-            try:
-                current_date = datetime.now().strftime("%Y-%m-%d")
-                inspection_name = f"Inspection-{current_date}-{len(output):03d}"
+            current_date = datetime.now().strftime("%Y-%m-%d")
+            inspection_name = f"Inspection-{current_date}-{len(output):03d}"
 
-                # Creating the Salesforce record with updated fields
+            try:
                 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
+                    "Fault_Type__c": output[0]["type"],
+                    "Severity__c": output[0]["severity"],
+                    "Fault_Summary__c": str(output),
+                    "Status__c": "New",
+                    "Annotated_Image_URL__c": "",
+                    "Report_PDF__c": ""
                 })
 
                 record_id = inspection_record.get("id")
+
                 content_version_id = upload_image_to_salesforce(
                     result_image,
                     filename=f"detected_defect_{record_id}.jpg",
@@ -73,8 +178,28 @@ def detect_defects(image):
             except Exception as e:
                 output.append({"error": f"Failed to create Salesforce record: {str(e)}"})
 
-        return result_image, output
+            return result_image, str(output)
+
+        return result_image, "No defects detected above confidence threshold."
 
     except Exception as e:
-        logging.error(f"Processing failed: {str(e)}")
-        return None, {"error": f"Processing failed: {str(e)}"}
+        logging.error(f"Detection failed: {str(e)}")
+        return None, f"Detection failed: {str(e)}"
+
+# Gradio interface definition
+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=(
+        "Upload drone-captured images to detect structural defects like cracks, rust, spalling, "
+        "and deformations using Faster R-CNN. Detected faults are stored in Salesforce with annotated images."
+    )
+)
+
+if __name__ == "__main__":
+    demo.launch()