Update app.py

app.py

@@ -1,114 +1,80 @@
-from flask import Flask, request, jsonify
-from ultralytics import YOLO
-from PIL import Image
-import io
-import os
-from simple_salesforce import Salesforce
+def detect_defects(image):
+    if not image:
+        return None, {"error": "No image provided"}
 
-app = Flask(__name__)
+    try:
+        # Perform detection
+        image_tensor = transform(image).unsqueeze(0)
+        with torch.no_grad():
+            predictions = model(image_tensor)
 
-# Salesforce configuration
-SALESFORCE_USERNAME = "drone@sathkrutha.com"
-SALESFORCE_PASSWORD = "Komal1303@"
-SALESFORCE_SECURITY_TOKEN = "53AWRskW9EjWUsSL5LU6nFTy3"
+        result_image = image.copy()
+        draw = ImageDraw.Draw(result_image)
+        output = []
 
-# Initialize Salesforce client
-sf = Salesforce(
-    username=SALESFORCE_USERNAME,
-    password=SALESFORCE_PASSWORD,
-    security_token=SALESFORCE_SECURITY_TOKEN
-)
+        for i in range(len(predictions[0]['boxes'])):
+            score = predictions[0]['scores'][i].item()
+            if score < 0.5:  # Increased threshold for better quality detections
+                continue
 
-# 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"
+            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)
 
-# Define fault types and severities for mapping
-FAULT_TYPES = ["Crack", "Rust", "Deformation", "Corrosion", "Spalling"]
-SEVERITIES = ["Minor", "Moderate", "Critical"]
+            output.append({
+                "type": defect_type,
+                "confidence": round(score, 2),
+                "severity": severity,
+                "coco_label": coco_label
+            })
 
-@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
+            draw.rectangle(box, outline="red", width=3)
+            draw.text((box[0], box[1]), f"{defect_type}: {severity}", fill="red")
 
-        # 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
+        # Apply NMS to filter overlapping detections
+        boxes = predictions[0]['boxes']
+        scores = predictions[0]['scores']
+        keep = apply_nms(boxes, scores, threshold=0.5)
+        boxes = boxes[keep]
+        output = [output[i] for i in keep]
+
+        # Create Salesforce record if detections exist
+        if output:
+            try:
+                current_date = datetime.now().strftime("%Y-%m-%d")
+                inspection_name = f"Inspection-{current_date}-{len(output):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": str(output),  # Summarizing the defects
+                    "Status__c": "New",  # Default status
+                    "Annotated_Image_URL__c": "",  # Placeholder for image URL
+                    "Report_PDF__c": ""  # Placeholder for report PDF URL
+                })
 
-if __name__ == "__main__":
-    app.run(debug=True, host="0.0.0.0", port=5000)
+                record_id = inspection_record.get("id")
+                content_version_id = 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}"
+                    })
+
+                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, output
+
+    except Exception as e:
+        logging.error(f"Processing failed: {str(e)}")
+        return None, {"error": f"Processing failed: {str(e)}"}