solar_wind_fault_detection / services /detection_service.py
DSatishchandra's picture
Create detection_service.py
c8c5ca6 verified
raw
history blame
1.18 kB
import torch
from models.solar_model import load_solar_model
from models.windmill_model import load_windmill_model
def detect_faults_solar(model, image):
results = model(image)
faults = []
for detection in results.xywh[0]: # Assuming the results are in this format
class_id = int(detection[5].item()) # Get the class ID of the detected object
if class_id == 0: # Cracks
faults.append({"type": "Crack", "location": (detection[0], detection[1])})
elif class_id == 1: # Damage
faults.append({"type": "Damage", "location": (detection[0], detection[1])})
return faults
def detect_faults_windmill(model, image):
results = model(image)
faults = []
for detection in results.xywh[0]: # Assuming the results are in this format
class_id = int(detection[5].item()) # Get the class ID of the detected object
if class_id == 2: # Blade damage
faults.append({"type": "Blade Damage", "location": (detection[0], detection[1])})
elif class_id == 3: # Motor fault
faults.append({"type": "Motor Fault", "location": (detection[0], detection[1])})
return faults