Nawinkumar15/Solar_Panel_Faults_Detection_

1from transformers import DetrImageProcessor, DetrForObjectDetection2import torch3from PIL import Image4 5class DetectionService:6    def __init__(self, model_name="facebook/detr-resnet-50"):7        self.processor = DetrImageProcessor.from_pretrained(model_name, revision="no_timm")8        self.model = DetrForObjectDetection.from_pretrained(model_name, revision="no_timm")9        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")10        self.model.to(self.device)11        self.model.eval()12        self.frame_counter = 013        self.frame_skip = 5  # Process every 5th frame for performance14 15    def detect_objects(self, image, confidence_threshold=0.9):16        """Detect objects in an image, skipping frames for performance."""17        self.frame_counter += 118        if self.frame_counter % self.frame_skip != 0:19            return []  # Skip detection for this frame20 21        inputs = self.processor(images=image, return_tensors="pt").to(self.device)22        with torch.no_grad():23            outputs = self.model(**inputs)24        target_sizes = torch.tensor([image.size[::-1]]).to(self.device)25        results = self.processor.post_process_object_detection(26            outputs, target_sizes=target_sizes, threshold=confidence_threshold27        )[0]28        detections = []29        for score, label, box in zip(30            results["scores"], results["labels"], results["boxes"]31        ):32            box = box.cpu().numpy().astype(int)33            detections.append({34                "score": score.item(),35                "label": self.model.config.id2label[label.item()],36                "box": {"xmin": box[0], "ymin": box[1], "xmax": box[2], "ymax": box[3]}37            })38        return detections