Update app.py

app.py

@@ -1,71 +1,189 @@
+import gradio as gr
+from PIL import Image, ImageDraw
+import torch
+from torchvision import models, transforms
+from si mple_salesforce import Salesforce
+import base64
+from io import BytesIO
+import json
+from datetime import datetime
 import logging
-import requests
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 
-MODEL_URL = "https://huggingface.co/your_huggingface_model_path"  # Replace with actual model URL
-HEADERS = {
-    "Authorization": "Bearer YOUR_HUGGINGFACE_API_KEY"  # Replace with your Hugging Face API key
+# 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(),
+])
+
+# 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 detect_defects(image):
-    if not image:
-        return None, "No image provided"
+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:
-        # Prepare image for model inference
         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)}")
 
-        # Send image to Hugging Face hosted model for detection (assuming YOLOv8 model)
-        response = requests.post(MODEL_URL, headers=HEADERS, json={"inputs": img_data})
-
-        # Debugging the raw response
-        logging.info(f"Response Status Code: {response.status_code}")
-        logging.info(f"Response Text: {response.text}")  # Log the raw response
-
-        # Check for valid response
-        if response.status_code != 200:
-            logging.error(f"Error in API request: {response.status_code}")
-            return None, f"API Error: {response.status_code}"
-
-        # Attempt to parse the JSON response
-        try:
-            response_data = response.json()
-        except ValueError:
-            logging.error("API response is not valid JSON")
-            return None, "API response is not valid JSON"
-
-        # Assuming response is a list of detected objects with bounding boxes, labels, and scores
-        faults = []
-        for item in response_data:
-            defect_type = item["label"]
-            severity = "Critical" if item["score"] > 0.9 else "Moderate" if item["score"] > 0.7 else "Minor"
-            faults.append({
-                "type": defect_type,
-                "confidence": round(item["score"], 2),
-                "severity": severity,
-                "box": item["bbox"]
-            })
+# 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)
 
         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.5:  # Lowered threshold for testing
+                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
+            })
 
-        # Annotate image with bounding boxes and labels
-        for fault in faults:
-            box = fault["box"]
             draw.rectangle(box, outline="red", width=3)
-            draw.text((box[0], box[1]), f"{fault['type']}: {fault['severity']}", fill="red")
+            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}"
+
+                # 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": json.dumps(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
+                )
 
-        # Return a more readable, text-based format instead of JSON
-        result_text = "\nDetected Faults and Severities:\n"
-        for fault in faults:
-            result_text += f"- {fault['type']} (Confidence: {fault['confidence']}, Severity: {fault['severity']})\n"
+                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}"
+                    })
 
-        return result_image, result_text
+                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, 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.JSON(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()