Update app.py

app.py

@@ -1,40 +1,114 @@
-import gradio as gr
-from PIL import Image, ImageDraw
-import random
-
-# Simulated defect labels
-DEFECT_LABELS = ['crack', 'spalling', 'rust', 'deformation']
-
-# Dummy detection function
-def simulate_defect_detection(image):
-    image = image.convert("RGB")
-    draw = ImageDraw.Draw(image)
-    output_lines = []
-
-    # Randomly simulate 2–4 detections
-    for _ in range(random.randint(2, 4)):
-        label = random.choice(DEFECT_LABELS)
-        conf = round(random.uniform(0.7, 0.99), 2)
-        x1, y1 = random.randint(20, 250), random.randint(20, 250)
-        x2, y2 = x1 + random.randint(50, 120), y1 + random.randint(50, 120)
-        draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
-        draw.text((x1, y1 - 12), f"{label} ({conf})", fill="red")
-        output_lines.append(f"{label}: {conf}")
-
-    return image, "\n".join(output_lines)
-
-# Gradio UI
-demo = gr.Interface(
-    fn=simulate_defect_detection,
-    inputs=gr.Image(type="pil", label="Upload Drone Image"),
-    outputs=[
-        gr.Image(type="pil", label="Simulated Structural Defects"),
-        gr.Textbox(label="Predicted Defects with Confidence")
-    ],
-    title="🧱 Structural Defect Detection (Demo Mode)",
-    description="Simulates detection of cracks, rust, spalling, and deformation in drone-captured images.",
-    allow_flagging="never"
+from flask import Flask, request, jsonify
+from ultralytics import YOLO
+from PIL import Image
+import io
+import os
+from simple_salesforce import Salesforce
+
+app = Flask(__name__)
+
+# Salesforce configuration
+SALESFORCE_USERNAME = "drone@sathkrutha.com"
+SALESFORCE_PASSWORD = "Komal1303@"
+SALESFORCE_SECURITY_TOKEN = "53AWRskW9EjWUsSL5LU6nFTy3"
+
+# Initialize Salesforce client
+sf = Salesforce(
+    username=SALESFORCE_USERNAME,
+    password=SALESFORCE_PASSWORD,
+    security_token=SALESFORCE_SECURITY_TOKEN
 )
 
+# Load YOLOv8 model (assumes model is in Space's root or a subdirectory)
+model = YOLO("yolov8n.pt")  # Replace with path to your fine-tuned model, e.g., "models/drone-inspection.pt"
+
+# Define fault types and severities for mapping
+FAULT_TYPES = ["Crack", "Rust", "Deformation", "Corrosion", "Spalling"]
+SEVERITIES = ["Minor", "Moderate", "Critical"]
+
+@app.route("/process-image", methods=["POST"])
+def process_image():
+    try:
+        # Validate request
+        if "image" not in request.files:
+            return jsonify({"error": "No image provided"}), 400
+        
+        image_file = request.files["image"]
+        site_id = request.form.get("site_id")
+        inspection_date = request.form.get("inspection_date")
+        
+        if not site_id or not inspection_date:
+            return jsonify({"error": "Missing site_id or inspection_date"}), 400
+
+        # Load and preprocess image
+        image = Image.open(image_file).convert("RGB")
+        
+        # Run YOLOv8 inference
+        results = model(image)
+        
+        # Process predictions
+        fault_type = "None"
+        severity = "None"
+        confidence = 0.0
+        
+        # Extract top prediction (assuming YOLOv8 output format)
+        if results and results[0].boxes:
+            top_box = results[0].boxes[0]  # Get highest confidence detection
+            class_id = int(top_box.cls)
+            confidence = float(top_box.conf)
+            
+            # Map class_id to fault_type and severity
+            # Assume class IDs are structured: 0-14 (5 fault types x 3 severities)
+            fault_idx = class_id // len(SEVERITIES)
+            severity_idx = class_id % len(SEVERITIES)
+            if fault_idx < len(FAULT_TYPES):
+                fault_type = FAULT_TYPES[fault_idx]
+                severity = SEVERITIES[severity_idx]
+        
+        # Generate placeholder URLs (replace with actual logic for annotated images/reports)
+        annotated_image_url = "https://example.com/annotated_image.png"
+        report_pdf_url = "https://example.com/report.pdf"
+        
+        # Create Salesforce record
+        inspection_data = {
+            "Site__c": site_id,
+            "Inspection_Date__c": inspection_date,
+            "Fault_Type__c": fault_type,
+            "Severity__c": severity,
+            "Annotated_Image_URL__c": annotated_image_url,
+            "Report_PDF__c": report_pdf_url,
+            "Fault_Summary__c": f"Detected {fault_type} with {severity} severity (Confidence: {confidence:.2f})",
+            "Status__c": "New"
+        }
+        
+        # Upload image to Salesforce as ContentVersion
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format="PNG")
+        img_byte_arr = img_byte_arr.getvalue()
+        
+        content_version = {
+            "Title": f"Drone_Image_{site_id}_{inspection_date}",
+            "PathOnClient": image_file.filename,
+            "VersionData": img_byte_arr
+        }
+        content_result = sf.ContentVersion.create(content_version)
+        
+        # Link ContentVersion to inspection record
+        inspection_data["Drone_Image__c"] = content_result["id"]
+        
+        # Create Drone_Structure_Inspection__c record
+        result = sf.Drone_Structure_Inspection__c.create(inspection_data)
+        
+        return jsonify({
+            "message": "Inspection processed successfully",
+            "record_id": result["id"],
+            "fault_type": fault_type,
+            "severity": severity,
+            "confidence": confidence
+        }), 200
+    
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
 if __name__ == "__main__":
-    demo.launch()
+    app.run(debug=True, host="0.0.0.0", port=5000)