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 from reportlab.lib.pagesizes import letter from reportlab.pdfgen import canvas import gradio as gr # 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 model = models.detection.fasterrcnn_resnet50_fpn(weights="FasterRCNN_ResNet50_FPN_Weights.COCO_V1") model.eval() # Image transformation for the model input transform = transforms.Compose([ transforms.ToTensor(), ]) VALID_FAULT_TYPES = ["Crack", "Rust", "Spalling", "Deformation", "Corrosion"] VALID_SEVERITIES = ["Minor", "Moderate", "Critical"] def get_severity(score): if score >= 0.9: return "Critical" elif score >= 0.7: return "Moderate" else: return "Minor" def map_defect_type(): return VALID_FAULT_TYPES[0] 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)}") def create_pdf_report(defect_list): buffer = BytesIO() c = canvas.Canvas(buffer, pagesize=letter) width, height = letter c.setFont("Helvetica-Bold", 14) c.drawString(30, height - 50, "Structural Defect Detection Report") c.setFont("Helvetica", 12) y = height - 80 for i, defect in enumerate(defect_list, 1): text = f"{i}. Type: {defect['type']}, Confidence: {defect['confidence']}, Severity: {defect['severity']}" c.drawString(30, y, text) y -= 20 if y < 50: c.showPage() c.setFont("Helvetica", 12) y = height - 50 c.save() pdf = buffer.getvalue() buffer.close() return pdf def upload_pdf_to_salesforce(pdf_bytes, filename="report.pdf", record_id=None): try: pdf_data = base64.b64encode(pdf_bytes).decode("utf-8") content_version = sf.ContentVersion.create({ "Title": filename, "PathOnClient": filename, "VersionData": pdf_data, "FirstPublishLocationId": record_id if record_id else SITE_RECORD_ID }) logging.info(f"PDF uploaded to Salesforce ContentVersion ID: {content_version['id']}") return content_version["id"] except Exception as e: logging.error(f"Failed to upload PDF to Salesforce: {str(e)}") raise Exception(f"Failed to upload PDF to Salesforce: {str(e)}") def detect_defects(image): if image is None: return None, "No image provided" try: image_tensor = transform(image).unsqueeze(0) with torch.no_grad(): predictions = model(image_tensor) result_image = image.copy() draw = ImageDraw.Draw(result_image) 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.3: continue box = predictions[0]['boxes'][i].tolist() defect_type = map_defect_type() severity = get_severity(score) output.append({ "type": defect_type, "confidence": round(score, 2), "severity": severity, }) draw.rectangle(box, outline="red", width=3) text = f"{defect_type}: {severity}" draw.text((box[0], box[1] - 20 if box[1] > 20 else box[1],), text, fill="red", font=font) if output: # Fixed date formatting to ensure proper YYYY-MM-DD format current_date = datetime.now().strftime("%Y-%m-%d") inspection_name = f"Inspection-{current_date}-{len(output):03d}" try: inspection_record = sf.Drone_Structure_Inspection__c.create({ "Inspection_Date__c": current_date, "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_img = upload_image_to_salesforce( result_image, filename=f"detected_defect_{record_id}.jpg", record_id=record_id ) pdf_bytes = create_pdf_report(output) content_version_id_pdf = upload_pdf_to_salesforce( pdf_bytes, filename=f"defect_report_{record_id}.pdf", record_id=record_id ) update_data = {} if content_version_id_img: update_data["Annotated_Image_URL__c"] = f"/sfc/servlet.shepherd/version/download/{content_version_id_img}" if content_version_id_pdf: update_data["Report_PDF__c"] = f"/sfc/servlet.shepherd/version/download/{content_version_id_pdf}" if update_data: sf.Drone_Structure_Inspection__c.update(record_id, update_data) 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, str(output) return result_image, "No defects detected above confidence threshold." except Exception as e: logging.error(f"Detection failed: {str(e)}") return None, f"Detection failed: {str(e)}" # Use gr.Blocks for more control over the UI with gr.Blocks() as demo: gr.Markdown( """ # Structural Defect Detection with Salesforce Integration 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. """ ) with gr.Row(): image_input = gr.Image(type="pil", label="Upload Drone Image") image_output = gr.Image(label="Detection Result") output_text = gr.Textbox(label="Detected Faults with Severity") with gr.Row(): clear_btn = gr.Button("Clear") submit_btn = gr.Button("Submit", variant="primary") submit_btn.click( fn=detect_defects, inputs=image_input, outputs=[image_output, output_text] ) clear_btn.click( fn=lambda: (None, ""), inputs=None, outputs=[image_input, output_text] ) if __name__ == "__main__": demo.launch(share=False)