Spaces:
Runtime error
Runtime error
| 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) |