# The implementation is adopted from FairMOT, # made publicly available under the MIT License at https://github.com/ifzhang/FairMOT import cv2 import numpy as np import torch from .image import transform_preds def xyxy2xywh(x): # Convert bounding box format from [x1, y1, x2, y2] to [x, y, w, h] y = torch.zeros(x.shape) if x.dtype is torch.float32 else np.zeros(x.shape) y[:, 0] = (x[:, 0] + x[:, 2]) / 2 y[:, 1] = (x[:, 1] + x[:, 3]) / 2 y[:, 2] = x[:, 2] - x[:, 0] y[:, 3] = x[:, 3] - x[:, 1] return y def xywh2xyxy(x): # Convert bounding box format from [x, y, w, h] to [x1, y1, x2, y2] y = torch.zeros(x.shape) if x.dtype is torch.float32 else np.zeros(x.shape) y[:, 0] = (x[:, 0] - x[:, 2] / 2) y[:, 1] = (x[:, 1] - x[:, 3] / 2) y[:, 2] = (x[:, 0] + x[:, 2] / 2) y[:, 3] = (x[:, 1] + x[:, 3] / 2) return y def bbox_iou(box1, box2, x1y1x2y2=False): """ Returns the IoU of two bounding boxes """ N, M = len(box1), len(box2) box1 = torch.from_numpy(np.stack(box1)) box2 = torch.from_numpy(np.stack(box2)) if x1y1x2y2: # Get the coordinates of bounding boxes b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3] b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3] else: # Transform from center and width to exact coordinates b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2 b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2 b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2 b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2 # get the coordinates of the intersection rectangle inter_rect_x1 = torch.max(b1_x1.unsqueeze(1), b2_x1) inter_rect_y1 = torch.max(b1_y1.unsqueeze(1), b2_y1) inter_rect_x2 = torch.min(b1_x2.unsqueeze(1), b2_x2) inter_rect_y2 = torch.min(b1_y2.unsqueeze(1), b2_y2) # Intersection area inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, 0) * torch.clamp( inter_rect_y2 - inter_rect_y1, 0) # Union Area b1_area = ((b1_x2 - b1_x1) * (b1_y2 - b1_y1)) b1_area = ((b1_x2 - b1_x1) * (b1_y2 - b1_y1)).view(-1, 1).expand(N, M) b2_area = ((b2_x2 - b2_x1) * (b2_y2 - b2_y1)).view(1, -1).expand(N, M) return inter_area / (b1_area + b2_area - inter_area + 1e-16) class LoadVideo: # for inference def __init__(self, path, img_size=(1088, 608)): self.cap = cv2.VideoCapture(path) self.frame_rate = int(round(self.cap.get(cv2.CAP_PROP_FPS))) self.vw = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) self.vh = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) self.vn = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) self.width = img_size[0] self.height = img_size[1] self.count = 0 self.w, self.h = 1920, 1080 print('Length of the video: {:d} frames'.format(self.vn)) def get_size(self, vw, vh, dw, dh): wa, ha = float(dw) / vw, float(dh) / vh a = min(wa, ha) return int(vw * a), int(vh * a) def __iter__(self): self.count = -1 return self def __next__(self): self.count += 1 if self.count == len(self): raise StopIteration # Read image res, img0 = self.cap.read() # BGR assert img0 is not None, 'Failed to load frame {:d}'.format(self.count) img0 = cv2.resize(img0, (self.w, self.h)) # Padded resize shape = [self.height, self.width] img, ratio, pad = letterbox(img0, shape) # Normalize RGB img = img[:, :, ::-1].transpose(2, 0, 1) img = np.ascontiguousarray(img, dtype=np.float32) img /= 255.0 return self.count, img, img0 def __len__(self): return self.vn # number of files def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True): # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 shape = img.shape[:2] # current shape [height, width] if isinstance(new_shape, int): new_shape = (new_shape, new_shape) # Scale ratio (new / old) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: # only scale down, do not scale up (for better test mAP) r = min(r, 1.0) # Compute padding ratio = r, r # width, height ratios new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[ 1] # wh padding if auto: # minimum rectangle dw, dh = np.mod(dw, 64), np.mod(dh, 64) # wh padding elif scaleFill: # stretch dw, dh = 0.0, 0.0 new_unpad = (new_shape[1], new_shape[0]) ratio = new_shape[1] / shape[1], new_shape[0] / shape[ 0] # width, height ratios dw /= 2 # divide padding into 2 sides dh /= 2 if shape[::-1] != new_unpad: # resize img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) img = cv2.copyMakeBorder( img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border return img, ratio, (dw, dh) def _gather_feat(feat, ind, mask=None): dim = feat.size(2) ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim) feat = feat.gather(1, ind) if mask is not None: mask = mask.unsqueeze(2).expand_as(feat) feat = feat[mask] feat = feat.view(-1, dim) return feat def _tranpose_and_gather_feat(feat, ind): feat = feat.permute(0, 2, 3, 1).contiguous() feat = feat.view(feat.size(0), -1, feat.size(3)) feat = _gather_feat(feat, ind) return feat class cfg_opt: K = 500 arch = 'yolo' conf_thres = 0.4 down_ratio = 4 head_conv = 256 heads = {'hm': 1, 'wh': 4, 'id': 64, 'reg': 2} img_size = (1088, 608) ltrb = True mean = [0.408, 0.447, 0.47] min_box_area = 100 num_classes = 1 reg_offset = True reid_dim = 64 std = [0.289, 0.274, 0.278] track_buffer = 30 def ctdet_post_process(dets, c, s, h, w, num_classes): ret = [] for i in range(dets.shape[0]): top_preds = {} dets[i, :, :2] = transform_preds(dets[i, :, 0:2], c[i], s[i], (w, h)) dets[i, :, 2:4] = transform_preds(dets[i, :, 2:4], c[i], s[i], (w, h)) classes = dets[i, :, -1] for j in range(num_classes): inds = (classes == j) det4 = dets[i, inds, :4].astype(np.float32) det5 = dets[i, inds, 4:5].astype(np.float32) top_preds[j + 1] = np.concatenate([det4, det5], axis=1).tolist() ret.append(top_preds) return ret