Spaces:
No application file
No application file
File size: 1,647 Bytes
aaaf2f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import torch
import cv2
from PIL import Image
from torchvision import transforms
import numpy as np
# Dummy models – replace with your own
def load_model(path):
model = torch.load(path, map_location="cpu")
model.eval()
return model
thermal_model = load_model("models/thermal_fault_model.pt")
crack_model = load_model("models/crack_detector_model.pt")
dust_model = load_model("models/dust_detector_model.pt")
def detect_faults(image):
transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
img_tensor = transform(image).unsqueeze(0)
results = []
output_img = image.copy()
# Simulated prediction
if torch.rand(1).item() > 0.5:
results.append({
"type": "Thermal Fault",
"location": "Panel 12-B",
"confidence": 0.93
})
cv2.putText(output_img, "Thermal Fault", (50, 50),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
if torch.rand(1).item() > 0.5:
results.append({
"type": "Crack Detected",
"location": "Panel 9-C",
"confidence": 0.88
})
cv2.putText(output_img, "Crack", (150, 100),
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
if torch.rand(1).item() > 0.5:
results.append({
"type": "Dust Coverage",
"location": "Panel 3-A",
"confidence": 0.91
})
cv2.putText(output_img, "Dust", (200, 150),
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 2)
return results, output_img
|