Sanjayraju30's picture
Create utils/detector.py
84dba68 verified
raw
history blame
1.07 kB
from transformers import DetrImageProcessor, DetrForObjectDetection
from PIL import Image
import torch
import numpy as np
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
def detect_faults(image):
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
intrusion_detected = any(label == 1 for label in results["labels"].tolist())
# Simulated thermal detection (average red channel > 200 = overheat)
red_mean = np.array(image)[:, :, 0].mean()
overheating = red_mean > 200
# Simulated shade (brightness < 100 on average = dusty/shaded)
brightness = np.array(image).mean()
dusty = brightness < 100
return {
"Intrusion Detected": intrusion_detected,
"Overheating Panel": overheating,
"Dust/Shade Fault": dusty
}