DSatishchandra's picture
Update services/detection_service.py
96f138d verified
raw
history blame
1.62 kB
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)
# Perform inference using ultralytics API
results = model.predict(img_tensor, verbose=False)
# Process results
faults = []
for detection in results[0].boxes: # ultralytics API
class_id = int(detection.cls)
if class_id == 0:
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.xyxy[0][0].item(), detection.xyxy[0][1].item())})
elif class_id == 4:
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)
# Perform inference using ultralytics API
results = model.predict(img_tensor, verbose=False)
# Process results
faults = []
for detection in results[0].boxes:
class_id = int(detection.cls)
if class_id == 2:
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.xyxy[0][0].item(), detection.xyxy[0][1].item())})
return faults