Update services/detection_service.py

services/detection_service.py

@@ -2,37 +2,41 @@ import torch
 import numpy as np
 
 def detect_faults_solar(model, image):
+    # Preprocess image
     img_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
     img_tensor = img_tensor.unsqueeze(0)
 
-    with torch.no_grad():
-        results = model(img_tensor)
+    # Perform inference using ultralytics API
+    results = model.predict(img_tensor, verbose=False)
 
+    # Process results
     faults = []
-    for detection in results.xyxy[0]:
-        class_id = int(detection[5])
+    for detection in results[0].boxes:  # ultralytics API
+        class_id = int(detection.cls)
         if class_id == 0:
-            faults.append({"type": "Crack", "location": (detection[0].item(), detection[1].item())})
+            faults.append({"type": "Crack", "location": (detection.xyxy[0][0].item(), detection.xyxy[0][1].item())})
         elif class_id == 1:
-            faults.append({"type": "Damage", "location": (detection[0].item(), detection[1].item())})
+            faults.append({"type": "Damage", "location": (detection.xyxy[0][0].item(), detection.xyxy[0][1].item())})
         elif class_id == 4:
-            faults.append({"type": "Hotspot", "location": (detection[0].item(), detection[1].item())})
+            faults.append({"type": "Hotspot", "location": (detection.xyxy[0][0].item(), detection.xyxy[0][1].item())})
 
     return faults
 
 def detect_faults_windmill(model, image):
+    # Preprocess image
     img_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
     img_tensor = img_tensor.unsqueeze(0)
 
-    with torch.no_grad():
-        results = model(img_tensor)
+    # Perform inference using ultralytics API
+    results = model.predict(img_tensor, verbose=False)
 
+    # Process results
     faults = []
-    for detection in results.xyxy[0]:
-        class_id = int(detection[5])
+    for detection in results[0].boxes:
+        class_id = int(detection.cls)
         if class_id == 2:
-            faults.append({"type": "Blade Damage", "location": (detection[0].item(), detection[1].item())})
+            faults.append({"type": "Blade Damage", "location": (detection.xyxy[0][0].item(), detection.xyxy[0][1].item())})
         elif class_id == 3:
-            faults.append({"type": "Motor Fault", "location": (detection[0].item(), detection[1].item())})
+            faults.append({"type": "Motor Fault", "location": (detection.xyxy[0][0].item(), detection.xyxy[0][1].item())})
 
     return faults