Update app.py

app.py

@@ -1,4 +1,3 @@
-import gradio as gr
 from PIL import Image, ImageDraw, ImageFont
 import torch
 from torchvision import models, transforms
@@ -7,6 +6,9 @@ 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  # Add this import here
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -43,11 +45,9 @@ transform = transforms.Compose([
 ])
 
 # Define valid picklist values for Fault_Type__c and Severity__c
-# Replace these with your actual Salesforce picklist values
 VALID_FAULT_TYPES = ["Crack", "Rust", "Spalling", "Deformation", "Corrosion"]
 VALID_SEVERITIES = ["Minor", "Moderate", "Critical"]
 
-# Map confidence score to severity level ensuring valid picklist values
 def get_severity(score):
     if score >= 0.9:
         return "Critical"
@@ -56,13 +56,10 @@ def get_severity(score):
     else:
         return "Minor"
 
-# Map defect to a valid picklist fault type
 def map_defect_type():
-    # You can add logic here to decide which fault type to return
-    # For now, always return the first valid fault type
+    # Always return first valid fault type (customize if needed)
     return VALID_FAULT_TYPES[0]
 
-# Upload annotated image to Salesforce as ContentVersion record
 def upload_image_to_salesforce(image, filename="detected_image.jpg", record_id=None):
     try:
         buffered = BytesIO()
@@ -81,23 +78,58 @@ def upload_image_to_salesforce(image, filename="detected_image.jpg", record_id=N
         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
+# New function: generate PDF report bytes from defect list
+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
+
+# New function: upload PDF bytes to Salesforce ContentVersion
+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:
-        # 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)
-
-        # Use default font
         try:
             font = ImageFont.truetype("arial.ttf", 18)
         except:
@@ -106,27 +138,20 @@ def detect_defects(image):
         output = []
         for i in range(len(predictions[0]['boxes'])):
             score = predictions[0]['scores'][i].item()
-            if score < 0.3:  # confidence threshold
+            if score < 0.3:
                 continue
-
             box = predictions[0]['boxes'][i].tolist()
-
             defect_type = map_defect_type()
             severity = get_severity(score)
-
-            # Append defect info to output list
             output.append({
                 "type": defect_type,
                 "confidence": round(score, 2),
                 "severity": severity,
             })
-
-            # Draw rectangle and label
             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 defects found, create Salesforce record & upload annotated image
         if output:
             current_date = datetime.now().strftime("%Y-%m-%d")
             inspection_name = f"Inspection-{current_date}-{len(output):03d}"
@@ -144,18 +169,30 @@ def detect_defects(image):
 
                 record_id = inspection_record.get("id")
 
-                content_version_id = upload_image_to_salesforce(
+                content_version_id_img = 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}"
-                    })
+                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)}"})
 
@@ -167,7 +204,6 @@ def detect_defects(image):
         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"),