Update app.py

app.py

@@ -1,110 +1,23 @@
-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 json
-from datetime import datetime
-import logging
+from PIL import ImageFont
 
-# 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
+# Update transform
 transform = transforms.Compose([
     transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
 ])
 
-# 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)}")
+# Load font for drawing text
+try:
+    font = ImageFont.truetype("arial.ttf", 16)
+except IOError:
+    font = ImageFont.load_default()
 
-# 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)
@@ -113,14 +26,20 @@ def detect_defects(image):
         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
+        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:
                 continue
 
-            box = predictions[0]['boxes'][i].tolist()
-            label_idx = predictions[0]['labels'][i].item()
+            box = boxes[i].tolist()
+            label_idx = 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)
 
@@ -131,59 +50,17 @@ 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]), f"{defect_type}: {severity}", fill="red")
+            draw.text((box[0], box[1] - 15), f"{defect_type}: {severity}", fill="red", font=font)
 
-        # 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}"
+        if not output:
+            output = [{"message": "No defects detected above threshold"}]
 
-                # 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
-                )
-
-                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)}"})
+        # Continue with Salesforce upload code as before...
 
         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.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()