app.py

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)