detection_helpers.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. import cv2
  2. import torch
  3. from numpy import random
  4. from models.experimental import attempt_load
  5. from utils.datasets import letterbox, np
  6. from utils.general import check_img_size, non_max_suppression, apply_classifier,scale_coords, xyxy2xywh
  7. from utils.plots import plot_one_box
  8. from utils.torch_utils import select_device, load_classifier,TracedModel
  9. class Detector:
  10. def __init__(self, conf_thres:float = 0.25, iou_thresh:float = 0.45, agnostic_nms:bool = False, save_conf:bool = False, classes:list = None):
  11. # device_ = torch.device("cuda:3" if torch.cuda.is_available() else 'cpu')
  12. '''
  13. args:
  14. conf_thres: Thresholf for Classification
  15. iou_thres: Thresholf for IOU box to consider
  16. agnostic_nms: whether to use Class-Agnostic NMS
  17. save_conf: whether to save confidences in 'save_txt' labels afters inference
  18. classes: Filter by class from COCO. can be in the format [0] or [0,1,2] etc
  19. '''
  20. self.device = select_device("cuda:3" if torch.cuda.is_available() else 'cpu')
  21. self.conf_thres = conf_thres
  22. self.iou_thres = iou_thresh
  23. self.classes = classes
  24. self.agnostic_nms = agnostic_nms
  25. self.save_conf = save_conf
  26. print(torch.cuda.device_count())
  27. def load_model(self, weights:str, img_size:int = 640, trace:bool = True, classify:bool = False):
  28. '''
  29. weights: Path to the model
  30. img_size: Input image size of the model
  31. trace: Whether to trace the model or not
  32. classify: whether to load the second stage classifier model or not
  33. '''
  34. self.half = self.device.type != 'cpu' # half precision only supported on CUDA
  35. self.model = attempt_load(weights, map_location=self.device) # load FP32 model
  36. self.stride = int(self.model.stride.max()) # model stride
  37. self.imgsz = check_img_size(img_size, s=self.stride) # check img_size
  38. if trace:
  39. self.model = TracedModel(self.model, self.device, img_size)
  40. if self.half:
  41. self.model.half() # to FP1
  42. # Run inference for CUDA just once
  43. if self.device.type != 'cpu':
  44. self.model(torch.zeros(1, 3, self.imgsz, self.imgsz).to(self.device).type_as(next(self.model.parameters()))) # run once
  45. # Second-stage classifier
  46. self.classify = classify
  47. if classify:
  48. self.modelc = load_classifier(name='resnet101', n=2) # initialize
  49. self.modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=self.device)['model']).to(self.device).eval()
  50. # Get names and colors of Colors for BB creation
  51. self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
  52. self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]
  53. @torch.no_grad()
  54. def detect(self, source, plot_bb:bool =True):
  55. '''
  56. source: Path to image file, video file, link or text etc
  57. plot_bb: whether to plot the bounding box around image or return the prediction
  58. '''
  59. img, im0 = self.load_image(source)
  60. img = torch.from_numpy(img).to(self.device)
  61. img = img.half() if self.half else img.float() # uint8 to fp16/32
  62. img /= 255.0 # 0 - 255 to 0.0 - 1.0
  63. if img.ndimension() == 3: # Single batch -> single image
  64. img = img.unsqueeze(0)
  65. # Inference
  66. pred = self.model(img, augment=False)[0] # We don not need any augment during inference time
  67. # Apply NMS
  68. pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, classes=self.classes, agnostic=self.agnostic_nms)
  69. # Apply Classifier
  70. if self.classify:
  71. pred = apply_classifier(pred, self.modelc, img, im0) # I thnk we need to add a new axis to im0
  72. # Post - Process detections
  73. det = pred[0]# detections per image but as we have just 1 image, it is the 0th index
  74. if len(det):
  75. # Rescale boxes from img_size to im0 size
  76. det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
  77. # Write results
  78. for *xyxy, conf, cls in reversed(det):
  79. if plot_bb: # Add bbox to image # save_img
  80. label = f'{self.names[int(cls)]} {conf:.2f}'
  81. plot_one_box(xyxy, im0, label=label, color=self.colors[int(cls)], line_thickness=1)
  82. return im0 if plot_bb else det.detach().cpu().numpy()
  83. return im0 if plot_bb else None # just in case there's no detection, return the original image. For tracking purpose plot_bb has to be False always
  84. def load_image(self, img0):
  85. '''
  86. Load and pre process the image
  87. args: img0: Path of image or numpy image in 'BGR" format
  88. '''
  89. if isinstance(img0, str): img0 = cv2.imread(img0) # BGR
  90. assert img0 is not None, 'Image Not Found '
  91. # Padded resize
  92. img = letterbox(img0, self.imgsz, stride=self.stride)[0]
  93. # Convert
  94. img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
  95. img = np.ascontiguousarray(img)
  96. return img, img0
  97. def save_txt(self, det, im0_shape, txt_path):
  98. '''
  99. Save the results of an image in a .txt file
  100. args:
  101. det: detecttions from the model
  102. im0_shape: Shape of Original image
  103. txt_path: File of the text path
  104. '''
  105. gn = torch.tensor(im0_shape)[[1, 0, 1, 0]] # normalization gain whwh
  106. for *xyxy, conf, cls in reversed(det):
  107. xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
  108. line = (cls, *xywh, conf) if self.save_conf else (cls, *xywh) # label format
  109. with open(txt_path + '.txt', 'a') as f:
  110. f.write(('%g ' * len(line)).rstrip() % line + '\n')