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
Files changed (1) hide show
  1. 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: