Update app.py

app.py

@@ -8,6 +8,8 @@ from io import BytesIO
 import json
 from datetime import datetime
 import logging
+import requests
+from fpdf import FPDF
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -34,50 +36,33 @@ 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'
-]
+# Load model from Hugging Face (Assuming you have a YOLOv8 model hosted)
+MODEL_URL = "https://huggingface.co/your_huggingface_model_path"  # Replace with actual model URL
+HEADERS = {
+    "Authorization": "Bearer YOUR_HUGGINGFACE_API_KEY"  # Replace with your Hugging Face API key
+}
 
 # Image transformations
 transform = transforms.Compose([
     transforms.ToTensor(),
 ])
 
-# 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',
-}
+# Function to generate PDF report
+def generate_pdf_report(faults):
+    pdf = FPDF()
+    pdf.set_auto_page_break(auto=True, margin=15)
+    pdf.add_page()
+
+    pdf.set_font("Arial", size=12)
+    pdf.cell(200, 10, txt="Drone Fault Detection Report", ln=True, align='C')
+    
+    pdf.ln(10)
+    pdf.cell(200, 10, txt="Faults Detected:", ln=True)
 
-def map_defect_type(coco_label):
-    return COCO_TO_DEFECT_MAPPING.get(coco_label, "Crack")
+    for fault in faults:
+        pdf.cell(200, 10, txt=f"Defect: {fault['type']}, Severity: {fault['severity']}, Confidence: {fault['confidence']}", ln=True)
+
+    return pdf
 
 # Function to upload image to Salesforce as ContentVersion
 def upload_image_to_salesforce(image, filename="detected_image.jpg", record_id=None):
@@ -101,51 +86,51 @@ def upload_image_to_salesforce(image, filename="detected_image.jpg", record_id=N
 # Detect defects and integrate with Salesforce
 def detect_defects(image):
     if not image:
-        return None, {"error": "No image provided"}
+        return None, "No image provided"
 
     try:
-        # Perform detection
-        image_tensor = transform(image).unsqueeze(0)
-        with torch.no_grad():
-            predictions = model(image_tensor)
-
-        result_image = image.copy()
-        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
-                continue
+        # Prepare image for model inference
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        img_data = base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-            box = predictions[0]['boxes'][i].tolist()
-            label_idx = predictions[0]['labels'][i].item()
-            coco_label = COCO_INSTANCE_CATEGORY_NAMES[label_idx]
-            defect_type = map_defect_type(coco_label)
-            severity = get_severity(score)
+        # Send image to Hugging Face hosted model for detection (assuming YOLOv8 model)
+        response = requests.post(MODEL_URL, headers=HEADERS, json={"inputs": img_data})
+        response_data = response.json()
 
-            output.append({
+        # Assuming response is a list of detected objects with bounding boxes, labels, and scores
+        faults = []
+        for item in response_data:
+            defect_type = item["label"]
+            severity = "Critical" if item["score"] > 0.9 else "Moderate" if item["score"] > 0.7 else "Minor"
+            faults.append({
                 "type": defect_type,
-                "confidence": round(score, 2),
+                "confidence": round(item["score"], 2),
                 "severity": severity,
-                "coco_label": coco_label
+                "box": item["bbox"]
             })
 
+        result_image = image.copy()
+        draw = ImageDraw.Draw(result_image)
+
+        # Annotate image with bounding boxes and labels
+        for fault in faults:
+            box = fault["box"]
             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]), f"{fault['type']}: {fault['severity']}", fill="red")
 
         # Create Salesforce record if detections exist
-        if output:
+        if faults:
             try:
                 current_date = datetime.now().strftime("%Y-%m-%d")
-                inspection_name = f"Inspection-{current_date}-{len(output):03d}"
+                inspection_name = f"Inspection-{current_date}-{len(faults):03d}"
 
                 # 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
+                    "Fault_Type__c": faults[0]["type"],  # Mapping defect type
+                    "Severity__c": faults[0]["severity"],  # Mapping severity
+                    "Fault_Summary__c": "\n".join([f"{fault['type']} (Confidence: {fault['confidence']})" for fault in faults]),  # Summarizing the defects
                     "Status__c": "New",  # Default status
                     "Annotated_Image_URL__c": "",  # Placeholder for image URL
                     "Report_PDF__c": ""  # Placeholder for report PDF URL
@@ -158,20 +143,43 @@ def detect_defects(image):
                     record_id=record_id
                 )
 
-                if content_version_id:
+                # Generate and upload the PDF report
+                pdf = generate_pdf_report(faults)
+                pdf_output = "/tmp/report.pdf"
+                pdf.output(pdf_output)
+
+                # Upload the PDF report to Salesforce (ContentVersion)
+                with open(pdf_output, "rb") as f:
+                    report_data = f.read()
+
+                pdf_data = base64.b64encode(report_data).decode("utf-8")
+                content_version = sf.ContentVersion.create({
+                    "Title": f"Report_{record_id}.pdf",
+                    "PathOnClient": f"Report_{record_id}.pdf",
+                    "VersionData": pdf_data,
+                    "FirstPublishLocationId": record_id
+                })
+
+                if content_version:
                     sf.Drone_Structure_Inspection__c.update(record_id, {
+                        "Report_PDF__c": f"/sfc/servlet.shepherd/version/download/{content_version['id']}",
                         "Annotated_Image_URL__c": f"/sfc/servlet.shepherd/version/download/{content_version_id}"
                     })
 
-                output.append({"salesforce_record_id": record_id})
+                faults.append(f"Salesforce record ID: {record_id}")
             except Exception as e:
-                output.append({"error": f"Failed to create Salesforce record: {str(e)}"})
+                faults.append(f"Failed to create Salesforce record: {str(e)}")
+
+        # Return a more readable, text-based format instead of JSON
+        result_text = "\nDetected Faults and Severities:\n"
+        for fault in faults:
+            result_text += f"- {fault}\n"
 
-        return result_image, output
+        return result_image, result_text
 
     except Exception as e:
         logging.error(f"Processing failed: {str(e)}")
-        return None, {"error": f"Processing failed: {str(e)}"}
+        return None, f"Processing failed: {str(e)}"
 
 # Gradio Interface
 demo = gr.Interface(
@@ -179,10 +187,10 @@ demo = gr.Interface(
     inputs=gr.Image(type="pil", label="Upload Drone Image"),
     outputs=[
         gr.Image(label="Detection Result"),
-        gr.JSON(label="Detected Faults with Severity")
+        gr.Textbox(label="Detected Faults with Severity", lines=10)  # Updated to use text output
     ],
     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."
+    description="Detects objects using YOLOv8 or Faster R-CNN hosted on Hugging Face and stores results in Salesforce. Fine-tune the model for structural defects like cracks, rust, and spalling."
 )
 
 if __name__ == "__main__":