Marcus Vinicius Zerbini Canhaço
commited on
Commit
·
9baead7
1
Parent(s):
b74c010
feat: atualização do detector com otimizações para GPU T4
Browse files- src/domain/detectors/gpu.py +59 -0
src/domain/detectors/gpu.py
CHANGED
@@ -363,6 +363,65 @@ class WeaponDetectorGPU(BaseDetector):
|
|
363 |
logger.error(f"Erro ao obter uso de memória GPU: {str(e)}")
|
364 |
return 0
|
365 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
366 |
def _should_clear_cache(self):
|
367 |
"""Determina se o cache deve ser limpo baseado no uso de memória."""
|
368 |
try:
|
|
|
363 |
logger.error(f"Erro ao obter uso de memória GPU: {str(e)}")
|
364 |
return 0
|
365 |
|
366 |
+
def _apply_nms(self, detections: list, iou_threshold: float = 0.5) -> list:
|
367 |
+
"""Aplica Non-Maximum Suppression nas detecções usando operações em GPU."""
|
368 |
+
try:
|
369 |
+
if not detections:
|
370 |
+
return []
|
371 |
+
|
372 |
+
# Converter detecções para tensores na GPU
|
373 |
+
boxes = torch.tensor([[d["box"][0], d["box"][1], d["box"][2], d["box"][3]] for d in detections], device=self.device)
|
374 |
+
scores = torch.tensor([d["confidence"] for d in detections], device=self.device)
|
375 |
+
labels = [d["label"] for d in detections]
|
376 |
+
|
377 |
+
# Calcular áreas dos boxes
|
378 |
+
area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
|
379 |
+
|
380 |
+
# Ordenar por score
|
381 |
+
_, order = scores.sort(descending=True)
|
382 |
+
|
383 |
+
keep = []
|
384 |
+
while order.numel() > 0:
|
385 |
+
if order.numel() == 1:
|
386 |
+
keep.append(order.item())
|
387 |
+
break
|
388 |
+
i = order[0]
|
389 |
+
keep.append(i.item())
|
390 |
+
|
391 |
+
# Calcular IoU com os boxes restantes
|
392 |
+
xx1 = torch.max(boxes[i, 0], boxes[order[1:], 0])
|
393 |
+
yy1 = torch.max(boxes[i, 1], boxes[order[1:], 1])
|
394 |
+
xx2 = torch.min(boxes[i, 2], boxes[order[1:], 2])
|
395 |
+
yy2 = torch.min(boxes[i, 3], boxes[order[1:], 3])
|
396 |
+
|
397 |
+
w = torch.clamp(xx2 - xx1, min=0)
|
398 |
+
h = torch.clamp(yy2 - yy1, min=0)
|
399 |
+
inter = w * h
|
400 |
+
|
401 |
+
# Calcular IoU
|
402 |
+
ovr = inter / (area[i] + area[order[1:]] - inter)
|
403 |
+
|
404 |
+
# Encontrar boxes com IoU menor que o threshold
|
405 |
+
ids = (ovr <= iou_threshold).nonzero().squeeze()
|
406 |
+
if ids.numel() == 0:
|
407 |
+
break
|
408 |
+
order = order[ids + 1]
|
409 |
+
|
410 |
+
# Construir lista de detecções filtradas
|
411 |
+
filtered_detections = []
|
412 |
+
for idx in keep:
|
413 |
+
filtered_detections.append({
|
414 |
+
"confidence": scores[idx].item(),
|
415 |
+
"box": boxes[idx].tolist(),
|
416 |
+
"label": labels[idx]
|
417 |
+
})
|
418 |
+
|
419 |
+
return filtered_detections
|
420 |
+
|
421 |
+
except Exception as e:
|
422 |
+
logger.error(f"Erro ao aplicar NMS na GPU: {str(e)}")
|
423 |
+
return []
|
424 |
+
|
425 |
def _should_clear_cache(self):
|
426 |
"""Determina se o cache deve ser limpo baseado no uso de memória."""
|
427 |
try:
|