Spaces:
Runtime error
Runtime error
| from collections import Counter | |
| import numpy as np | |
| from PIL import Image | |
| from io import BytesIO | |
| from ultralytics import YOLO | |
| class DengueDetector: | |
| def __init__(self, model_path="./models/DetectsmallTest1.pt"): | |
| self.model = YOLO(model_path) | |
| self.names = self.model.names | |
| def calculate_intensity(self, objects): | |
| weights = {"piscina": 9, "caixa_agua": 4, "carro": 1} | |
| score = sum(weights.get(obj["class"], 0) for obj in objects) | |
| return score | |
| def detect_image(self, image_bytes): | |
| # Carregar imagem da memória | |
| img = Image.open(BytesIO(image_bytes)).convert("RGB") | |
| img_np = np.array(img) # YOLO aceita np.array diretamente | |
| height, width = img_np.shape[:2] | |
| # Detectar objetos | |
| results = self.model(img_np) | |
| result = results[0] | |
| boxes = result.boxes | |
| class_ids = boxes.cls.tolist() | |
| confidences = boxes.conf.tolist() | |
| class_names = [self.names[int(cls)] for cls in class_ids] | |
| counts = Counter(class_names) | |
| # Construir lista de detecções | |
| detections = [] | |
| for i in range(len(boxes)): | |
| x1, y1, x2, y2 = map(float, boxes.xyxy[i]) | |
| conf = float(confidences[i]) | |
| cls_id = int(class_ids[i]) | |
| detections.append({ | |
| "class": self.names[cls_id], | |
| "confidence": round(conf, 4), | |
| "box": { | |
| "x1": x1, "y1": y1, "x2": x2, "y2": y2, | |
| "original_width": width, "original_height": height | |
| } | |
| }) | |
| intensity_score = self.calculate_intensity(detections) | |
| return { | |
| "total": len(class_ids), | |
| "contagem": counts, | |
| "objetos": detections, | |
| "intensity_score": intensity_score | |
| } |