Create detection_service.py
Browse files
services/detection_service.py
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from models.solar_model import load_solar_model
|
| 3 |
+
from models.windmill_model import load_windmill_model
|
| 4 |
+
|
| 5 |
+
def detect_faults_solar(model, image):
|
| 6 |
+
results = model(image)
|
| 7 |
+
faults = []
|
| 8 |
+
|
| 9 |
+
for detection in results.xywh[0]: # Assuming the results are in this format
|
| 10 |
+
class_id = int(detection[5].item()) # Get the class ID of the detected object
|
| 11 |
+
if class_id == 0: # Cracks
|
| 12 |
+
faults.append({"type": "Crack", "location": (detection[0], detection[1])})
|
| 13 |
+
elif class_id == 1: # Damage
|
| 14 |
+
faults.append({"type": "Damage", "location": (detection[0], detection[1])})
|
| 15 |
+
|
| 16 |
+
return faults
|
| 17 |
+
|
| 18 |
+
def detect_faults_windmill(model, image):
|
| 19 |
+
results = model(image)
|
| 20 |
+
faults = []
|
| 21 |
+
|
| 22 |
+
for detection in results.xywh[0]: # Assuming the results are in this format
|
| 23 |
+
class_id = int(detection[5].item()) # Get the class ID of the detected object
|
| 24 |
+
if class_id == 2: # Blade damage
|
| 25 |
+
faults.append({"type": "Blade Damage", "location": (detection[0], detection[1])})
|
| 26 |
+
elif class_id == 3: # Motor fault
|
| 27 |
+
faults.append({"type": "Motor Fault", "location": (detection[0], detection[1])})
|
| 28 |
+
|
| 29 |
+
return faults
|