637 lines
25 KiB
Python
637 lines
25 KiB
Python
|
import datetime
|
||
|
|
import os
|
||
|
|
import time
|
||
|
|
import ffmpeg
|
||
|
|
import torch
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
from multiprocessing import Process, Manager
|
||
|
|
from threading import Thread
|
||
|
|
from read_data import LoadImages, LoadStreams
|
||
|
|
import torch.backends.cudnn as cudnn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
import torchvision
|
||
|
|
from algorithm.yolov5.models.common import DetectMultiBackend
|
||
|
|
|
||
|
|
from algorithm.Remote_sense.nms_rotated import nms_rotated_ext
|
||
|
|
|
||
|
|
|
||
|
|
from PIL import Image, ImageDraw, ImageFont
|
||
|
|
|
||
|
|
pi = 3.141592
|
||
|
|
|
||
|
|
class Remote_Sense():
|
||
|
|
time_reference = datetime.datetime.now()
|
||
|
|
counter_frame = 0
|
||
|
|
processed_fps = 0
|
||
|
|
|
||
|
|
def __init__(self,video_path=None):
|
||
|
|
|
||
|
|
|
||
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
|
|
||
|
|
|
||
|
|
#self.model = torch.load('weight/remote_sensing/oriented.pt', map_location=self.device)['model'].float().fuse()
|
||
|
|
|
||
|
|
self.model = DetectMultiBackend(weights='weight/remote_sensing/oriented.pt', dnn=True, rotation = True)
|
||
|
|
# self.model.Detect.rotations = True
|
||
|
|
self.classes = self.model.names
|
||
|
|
|
||
|
|
self.frame = [None]
|
||
|
|
|
||
|
|
if video_path is not None:
|
||
|
|
self.video_name = video_path
|
||
|
|
else:
|
||
|
|
self.video_name = 'vid2.mp4' # A default video file
|
||
|
|
|
||
|
|
self.imgsz = 2048
|
||
|
|
|
||
|
|
self.dataset = LoadImages(path =self.video_name, img_size = self.imgsz)
|
||
|
|
|
||
|
|
self.names = self.model.names
|
||
|
|
|
||
|
|
def use_webcam(self, source):
|
||
|
|
# self.dataset.release() # Release any existing video capture
|
||
|
|
# self.cap = cv2.VideoCapture(0) # Open default webcam
|
||
|
|
# print('use_webcam')
|
||
|
|
source = source
|
||
|
|
cudnn.benchmark = True
|
||
|
|
self.dataset = LoadStreams(source, img_size=self.imgsz)
|
||
|
|
|
||
|
|
def class_to_label(self, x):
|
||
|
|
return self.classes[int(x)]
|
||
|
|
|
||
|
|
def get_frame(self):
|
||
|
|
|
||
|
|
colors = Colors()
|
||
|
|
|
||
|
|
for im0s in self.dataset:
|
||
|
|
# print(self.dataset.mode)
|
||
|
|
# print(self.dataset)
|
||
|
|
if self.dataset.mode == 'stream':
|
||
|
|
image = im0s[0].copy()
|
||
|
|
else:
|
||
|
|
image = im0s.copy()
|
||
|
|
img = image[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
||
|
|
|
||
|
|
img0 = img.copy()
|
||
|
|
|
||
|
|
img = torch.tensor(img0)
|
||
|
|
|
||
|
|
img = img.float() # uint8 to fp16/32
|
||
|
|
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||
|
|
if img.ndimension() == 3:
|
||
|
|
img = img.unsqueeze(0)
|
||
|
|
img = img.to(self.device)
|
||
|
|
self.model.to(self.device)
|
||
|
|
result = self.model(img)
|
||
|
|
|
||
|
|
pred = result[0]
|
||
|
|
|
||
|
|
pred = non_max_suppression_obb(pred, conf_thres=0.1, iou_thres=0.2, multi_label=True, max_det=1000)
|
||
|
|
|
||
|
|
# print(pred)
|
||
|
|
txt = ""
|
||
|
|
for i, det in enumerate(pred): # per image
|
||
|
|
pred_poly = rbox2poly(det[:, :5]) # (n, [x1 y1 x2 y2 x3 y3 x4 y4])
|
||
|
|
annotator = Annotator(image, line_width=3, example=str(self.names))
|
||
|
|
if len(det):
|
||
|
|
pred_poly = scale_polys(img.shape[2:], pred_poly, image.shape)
|
||
|
|
det = torch.cat((pred_poly, det[:, -2:]), dim=1) # (n, [poly conf cls])
|
||
|
|
|
||
|
|
# Print results
|
||
|
|
for c in det[:, -1].unique():
|
||
|
|
n = (det[:, -1] == c).sum() # detections per class
|
||
|
|
txt += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
|
||
|
|
|
||
|
|
for *poly, conf, cls in reversed(det):
|
||
|
|
c = int(cls)
|
||
|
|
label = None
|
||
|
|
# print(poly, label)
|
||
|
|
annotator.poly_label(poly, label, color=colors(c, True))
|
||
|
|
|
||
|
|
im0 = annotator.result()
|
||
|
|
|
||
|
|
|
||
|
|
# Draw the number of people on the frame and display it
|
||
|
|
ret, jpeg = cv2.imencode(".jpg", im0)
|
||
|
|
|
||
|
|
return jpeg.tobytes(), txt
|
||
|
|
|
||
|
|
|
||
|
|
class Colors:
|
||
|
|
# Ultralytics color palette https://ultralytics.com/
|
||
|
|
def __init__(self):
|
||
|
|
# hex = matplotlib.colors.TABLEAU_COLORS.values()
|
||
|
|
hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
|
||
|
|
'2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
|
||
|
|
self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
|
||
|
|
self.n = len(self.palette)
|
||
|
|
|
||
|
|
def __call__(self, i, bgr=False):
|
||
|
|
c = self.palette[int(i) % self.n]
|
||
|
|
return (c[2], c[1], c[0]) if bgr else c
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def hex2rgb(h): # rgb order (PIL)
|
||
|
|
return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
|
||
|
|
|
||
|
|
class Annotator:
|
||
|
|
|
||
|
|
def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
|
||
|
|
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images.'
|
||
|
|
self.pil = pil or not is_ascii(example) or is_chinese(example)
|
||
|
|
if self.pil: # use PIL
|
||
|
|
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
||
|
|
self.im_cv2 = im
|
||
|
|
self.draw = ImageDraw.Draw(self.im)
|
||
|
|
self.font = 'Arial.Unicode.ttf'
|
||
|
|
else: # use cv2
|
||
|
|
self.im = im
|
||
|
|
self.im_cv2 = im
|
||
|
|
self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) # line width
|
||
|
|
|
||
|
|
def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
|
||
|
|
# Add one xyxy box to image with label
|
||
|
|
if self.pil or not is_ascii(label):
|
||
|
|
self.draw.rectangle(box, width=self.lw, outline=color) # box
|
||
|
|
if label:
|
||
|
|
w, h = self.font.getsize(label) # text width, height
|
||
|
|
outside = box[1] - h >= 0 # label fits outside box
|
||
|
|
self.draw.rectangle([box[0],
|
||
|
|
box[1] - h if outside else box[1],
|
||
|
|
box[0] + w + 1,
|
||
|
|
box[1] + 1 if outside else box[1] + h + 1], fill=color)
|
||
|
|
# self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0
|
||
|
|
self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
|
||
|
|
else: # cv2
|
||
|
|
p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
|
||
|
|
cv2.rectangle(self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA)
|
||
|
|
if label:
|
||
|
|
tf = max(self.lw - 1, 1) # font thickness
|
||
|
|
w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0] # text width, height
|
||
|
|
outside = p1[1] - h - 3 >= 0 # label fits outside box
|
||
|
|
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
|
||
|
|
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
|
||
|
|
cv2.putText(self.im, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, self.lw / 3, txt_color,
|
||
|
|
thickness=tf, lineType=cv2.LINE_AA)
|
||
|
|
|
||
|
|
def poly_label(self, poly, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
|
||
|
|
if isinstance(poly, torch.Tensor):
|
||
|
|
poly = poly.cpu().numpy()
|
||
|
|
if isinstance(poly[0], torch.Tensor):
|
||
|
|
poly = [x.cpu().numpy() for x in poly]
|
||
|
|
polygon_list = np.array([(poly[0], poly[1]), (poly[2], poly[3]), \
|
||
|
|
(poly[4], poly[5]), (poly[6], poly[7])], np.int32)
|
||
|
|
cv2.drawContours(image=self.im_cv2, contours=[polygon_list], contourIdx=-1, color=color, thickness=self.lw)
|
||
|
|
if label:
|
||
|
|
tf = max(self.lw - 1, 1) # font thicknes
|
||
|
|
xmax, xmin, ymax, ymin = max(poly[0::2]), min(poly[0::2]), max(poly[1::2]), min(poly[1::2])
|
||
|
|
x_label, y_label = int((xmax + xmin)/2), int((ymax + ymin)/2)
|
||
|
|
w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0] # text width, height
|
||
|
|
cv2.rectangle(
|
||
|
|
self.im_cv2,
|
||
|
|
(x_label, y_label),
|
||
|
|
(x_label + w + 1, y_label + int(1.5*h)),
|
||
|
|
color, -1, cv2.LINE_AA
|
||
|
|
)
|
||
|
|
cv2.putText(self.im_cv2, label, (x_label, y_label + h), 0, self.lw / 3, txt_color, thickness=tf, lineType=cv2.LINE_AA)
|
||
|
|
self.im = self.im_cv2 if isinstance(self.im_cv2, Image.Image) else Image.fromarray(self.im_cv2)
|
||
|
|
|
||
|
|
def rectangle(self, xy, fill=None, outline=None, width=1):
|
||
|
|
# Add rectangle to image (PIL-only)
|
||
|
|
self.draw.rectangle(xy, fill, outline, width)
|
||
|
|
|
||
|
|
def text(self, xy, text, txt_color=(255, 255, 255)):
|
||
|
|
# Add text to image (PIL-only)
|
||
|
|
w, h = self.font.getsize(text) # text width, height
|
||
|
|
self.draw.text((xy[0], xy[1] - h + 1), text, fill=txt_color, font=self.font)
|
||
|
|
|
||
|
|
def result(self):
|
||
|
|
# Return annotated image as array
|
||
|
|
return np.asarray(self.im)
|
||
|
|
|
||
|
|
|
||
|
|
def time_synchronized():
|
||
|
|
# pytorch-accurate time
|
||
|
|
if torch.cuda.is_available():
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
return time.time()
|
||
|
|
|
||
|
|
def rbox2poly_single(rrect):
|
||
|
|
"""
|
||
|
|
rrect:[x_ctr,y_ctr,w,h,angle]
|
||
|
|
to
|
||
|
|
poly:[x0,y0,x1,y1,x2,y2,x3,y3]
|
||
|
|
"""
|
||
|
|
x_ctr, y_ctr, width, height, angle = rrect[:5]
|
||
|
|
tl_x, tl_y, br_x, br_y = -width/2, -height/2, width/2, height/2
|
||
|
|
rect = np.array([[tl_x, br_x, br_x, tl_x], [tl_y, tl_y, br_y, br_y]])
|
||
|
|
R = np.array([[np.cos(angle), -np.sin(angle)],
|
||
|
|
[np.sin(angle), np.cos(angle)]])
|
||
|
|
poly = R.dot(rect)
|
||
|
|
x0, x1, x2, x3 = poly[0, :4] + x_ctr
|
||
|
|
y0, y1, y2, y3 = poly[1, :4] + y_ctr
|
||
|
|
poly = np.array([x0, y0, x1, y1, x2, y2, x3, y3], dtype=np.float32)
|
||
|
|
poly = get_best_begin_point_single(poly)
|
||
|
|
return poly
|
||
|
|
|
||
|
|
def is_ascii(s=''):
|
||
|
|
# Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7)
|
||
|
|
s = str(s) # convert list, tuple, None, etc. to str
|
||
|
|
return len(s.encode().decode('ascii', 'ignore')) == len(s)
|
||
|
|
def is_chinese(s='人工智能'):
|
||
|
|
import re
|
||
|
|
# Is string composed of any Chinese characters?
|
||
|
|
return re.search('[\u4e00-\u9fff]', s)
|
||
|
|
|
||
|
|
def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
|
||
|
|
# Rescale boxes (xyxy) from img1_shape to img0_shape
|
||
|
|
if ratio_pad is None: # calculate from img0_shape
|
||
|
|
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
||
|
|
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
||
|
|
else:
|
||
|
|
gain = ratio_pad[0][0]
|
||
|
|
pad = ratio_pad[1]
|
||
|
|
|
||
|
|
boxes[..., [0, 2]] -= pad[0] # x padding
|
||
|
|
boxes[..., [1, 3]] -= pad[1] # y padding
|
||
|
|
boxes[..., :4] /= gain
|
||
|
|
clip_boxes(boxes, img0_shape)
|
||
|
|
return boxes
|
||
|
|
|
||
|
|
def scale_segments(img1_shape, segments, img0_shape, ratio_pad=None, normalize=False):
|
||
|
|
# Rescale coords (xyxy) from img1_shape to img0_shape
|
||
|
|
if ratio_pad is None: # calculate from img0_shape
|
||
|
|
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
||
|
|
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
||
|
|
else:
|
||
|
|
gain = ratio_pad[0][0]
|
||
|
|
pad = ratio_pad[1]
|
||
|
|
|
||
|
|
segments[:, 0] -= pad[0] # x padding
|
||
|
|
segments[:, 1] -= pad[1] # y padding
|
||
|
|
segments /= gain
|
||
|
|
clip_segments(segments, img0_shape)
|
||
|
|
if normalize:
|
||
|
|
segments[:, 0] /= img0_shape[1] # width
|
||
|
|
segments[:, 1] /= img0_shape[0] # height
|
||
|
|
return segments
|
||
|
|
|
||
|
|
def clip_boxes(boxes, shape):
|
||
|
|
# Clip boxes (xyxy) to image shape (height, width)
|
||
|
|
if isinstance(boxes, torch.Tensor): # faster individually
|
||
|
|
boxes[..., 0].clamp_(0, shape[1]) # x1
|
||
|
|
boxes[..., 1].clamp_(0, shape[0]) # y1
|
||
|
|
boxes[..., 2].clamp_(0, shape[1]) # x2
|
||
|
|
boxes[..., 3].clamp_(0, shape[0]) # y2
|
||
|
|
else: # np.array (faster grouped)
|
||
|
|
boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1]) # x1, x2
|
||
|
|
boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0]) # y1, y2
|
||
|
|
|
||
|
|
|
||
|
|
def clip_segments(segments, shape):
|
||
|
|
# Clip segments (xy1,xy2,...) to image shape (height, width)
|
||
|
|
if isinstance(segments, torch.Tensor): # faster individually
|
||
|
|
segments[:, 0].clamp_(0, shape[1]) # x
|
||
|
|
segments[:, 1].clamp_(0, shape[0]) # y
|
||
|
|
else: # np.array (faster grouped)
|
||
|
|
segments[:, 0] = segments[:, 0].clip(0, shape[1]) # x
|
||
|
|
segments[:, 1] = segments[:, 1].clip(0, shape[0]) # y
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def masks2segments(masks, strategy='largest'):
|
||
|
|
# Convert masks(n,160,160) into segments(n,xy)
|
||
|
|
segments = []
|
||
|
|
for x in masks.int().cpu().numpy().astype('uint8'):
|
||
|
|
c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
|
||
|
|
if c:
|
||
|
|
if strategy == 'concat': # concatenate all segments
|
||
|
|
c = np.concatenate([x.reshape(-1, 2) for x in c])
|
||
|
|
elif strategy == 'largest': # select largest segment
|
||
|
|
c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
|
||
|
|
else:
|
||
|
|
c = np.zeros((0, 2)) # no segments found
|
||
|
|
segments.append(c.astype('float32'))
|
||
|
|
return segments
|
||
|
|
|
||
|
|
def process_mask(protos, masks_in, bboxes, shape, upsample=False):
|
||
|
|
"""
|
||
|
|
Crop before upsample.
|
||
|
|
proto_out: [mask_dim, mask_h, mask_w]
|
||
|
|
out_masks: [n, mask_dim], n is number of masks after nms
|
||
|
|
bboxes: [n, 4], n is number of masks after nms
|
||
|
|
shape:input_image_size, (h, w)
|
||
|
|
|
||
|
|
return: h, w, n
|
||
|
|
"""
|
||
|
|
|
||
|
|
c, mh, mw = protos.shape # CHW
|
||
|
|
ih, iw = shape
|
||
|
|
# print(masks_in.shape, protos.shape)
|
||
|
|
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw) # CHW
|
||
|
|
|
||
|
|
downsampled_bboxes = bboxes.clone()
|
||
|
|
downsampled_bboxes[:, 0] *= mw / iw
|
||
|
|
downsampled_bboxes[:, 2] *= mw / iw
|
||
|
|
downsampled_bboxes[:, 3] *= mh / ih
|
||
|
|
downsampled_bboxes[:, 1] *= mh / ih
|
||
|
|
|
||
|
|
masks = crop_mask(masks, downsampled_bboxes) # CHW
|
||
|
|
if upsample:
|
||
|
|
masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW
|
||
|
|
return masks.gt_(0.5)
|
||
|
|
|
||
|
|
|
||
|
|
def crop_mask(masks, boxes):
|
||
|
|
"""
|
||
|
|
"Crop" predicted masks by zeroing out everything not in the predicted bbox.
|
||
|
|
Vectorized by Chong (thanks Chong).
|
||
|
|
|
||
|
|
Args:
|
||
|
|
- masks should be a size [h, w, n] tensor of masks
|
||
|
|
- boxes should be a size [n, 4] tensor of bbox coords in relative point form
|
||
|
|
"""
|
||
|
|
|
||
|
|
n, h, w = masks.shape
|
||
|
|
x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(1,1,n)
|
||
|
|
r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :] # rows shape(1,w,1)
|
||
|
|
c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None] # cols shape(h,1,1)
|
||
|
|
|
||
|
|
return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def scale_image(im1_shape, masks, im0_shape, ratio_pad=None):
|
||
|
|
"""
|
||
|
|
img1_shape: model input shape, [h, w]
|
||
|
|
img0_shape: origin pic shape, [h, w, 3]
|
||
|
|
masks: [h, w, num]
|
||
|
|
"""
|
||
|
|
# Rescale coordinates (xyxy) from im1_shape to im0_shape
|
||
|
|
if ratio_pad is None: # calculate from im0_shape
|
||
|
|
gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1]) # gain = old / new
|
||
|
|
pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2 # wh padding
|
||
|
|
else:
|
||
|
|
pad = ratio_pad[1]
|
||
|
|
top, left = int(pad[1]), int(pad[0]) # y, x
|
||
|
|
bottom, right = int(im1_shape[0] - pad[1]), int(im1_shape[1] - pad[0])
|
||
|
|
|
||
|
|
if len(masks.shape) < 2:
|
||
|
|
raise ValueError(f'"len of masks shape" should be 2 or 3, but got {len(masks.shape)}')
|
||
|
|
masks = masks[top:bottom, left:right]
|
||
|
|
# masks = masks.permute(2, 0, 1).contiguous()
|
||
|
|
# masks = F.interpolate(masks[None], im0_shape[:2], mode='bilinear', align_corners=False)[0]
|
||
|
|
# masks = masks.permute(1, 2, 0).contiguous()
|
||
|
|
|
||
|
|
masks = cv2.resize(masks, (im0_shape[1], im0_shape[0]))
|
||
|
|
|
||
|
|
if len(masks.shape) == 2:
|
||
|
|
masks = masks[:, :, None]
|
||
|
|
return masks
|
||
|
|
|
||
|
|
def xywh2xyxy(x):
|
||
|
|
# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
||
|
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||
|
|
y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
|
||
|
|
y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
|
||
|
|
y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
|
||
|
|
y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
|
||
|
|
return y
|
||
|
|
|
||
|
|
def non_max_suppression_obb(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
|
||
|
|
labels=(), max_det=1500):
|
||
|
|
"""Runs Non-Maximum Suppression (NMS) on inference results_obb
|
||
|
|
Args:
|
||
|
|
prediction (tensor): (b, n_all_anchors, [cx cy l s obj num_cls theta_cls])
|
||
|
|
agnostic (bool): True = NMS will be applied between elements of different categories
|
||
|
|
labels : () or
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
list of detections, len=batch_size, on (n,7) tensor per image [xylsθ, conf, cls] θ ∈ [-pi/2, pi/2)
|
||
|
|
"""
|
||
|
|
|
||
|
|
nc = prediction.shape[2] - 5 - 180 # number of classes
|
||
|
|
xc = prediction[..., 4] > conf_thres # candidates
|
||
|
|
class_index = nc + 5
|
||
|
|
|
||
|
|
# Checks
|
||
|
|
assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
|
||
|
|
assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
|
||
|
|
|
||
|
|
# Settings
|
||
|
|
max_wh = 4096 # min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
||
|
|
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
||
|
|
time_limit = 30.0 # seconds to quit after
|
||
|
|
# redundant = True # require redundant detections
|
||
|
|
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
|
||
|
|
|
||
|
|
t = time.time()
|
||
|
|
output = [torch.zeros((0, 7), device=prediction.device)] * prediction.shape[0]
|
||
|
|
for xi, x in enumerate(prediction): # image index, image inference
|
||
|
|
# Apply constraints
|
||
|
|
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
|
||
|
|
x = x[xc[xi]] # confidence, (tensor): (n_conf_thres, [cx cy l s obj num_cls theta_cls])
|
||
|
|
|
||
|
|
# Cat apriori labels if autolabelling
|
||
|
|
if labels and len(labels[xi]):
|
||
|
|
l = labels[xi]
|
||
|
|
v = torch.zeros((len(l), nc + 5), device=x.device)
|
||
|
|
v[:, :4] = l[:, 1:5] # box
|
||
|
|
v[:, 4] = 1.0 # conf
|
||
|
|
v[range(len(l)), l[:, 0].long() + 5] = 1.0 # cls
|
||
|
|
x = torch.cat((x, v), 0)
|
||
|
|
|
||
|
|
# If none remain process next image
|
||
|
|
if not x.shape[0]:
|
||
|
|
continue
|
||
|
|
|
||
|
|
# Compute conf
|
||
|
|
x[:, 5:class_index] *= x[:, 4:5] # conf = obj_conf * cls_conf
|
||
|
|
|
||
|
|
_, theta_pred = torch.max(x[:, class_index:], 1, keepdim=True) # [n_conf_thres, 1] θ ∈ int[0, 179]
|
||
|
|
theta_pred = (theta_pred - 90) / 180 * pi # [n_conf_thres, 1] θ ∈ [-pi/2, pi/2)
|
||
|
|
|
||
|
|
# Detections matrix nx7 (xyls, θ, conf, cls) θ ∈ [-pi/2, pi/2)
|
||
|
|
if multi_label:
|
||
|
|
i, j = (x[:, 5:class_index] > conf_thres).nonzero(as_tuple=False).T # ()
|
||
|
|
x = torch.cat((x[i, :4], theta_pred[i], x[i, j + 5, None], j[:, None].float()), 1)
|
||
|
|
else: # best class only
|
||
|
|
conf, j = x[:, 5:class_index].max(1, keepdim=True)
|
||
|
|
x = torch.cat((x[:, :4], theta_pred, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
||
|
|
|
||
|
|
# Filter by class
|
||
|
|
if classes is not None:
|
||
|
|
x = x[(x[:, 6:7] == torch.tensor(classes, device=x.device)).any(1)]
|
||
|
|
|
||
|
|
# Apply finite constraint
|
||
|
|
# if not torch.isfinite(x).all():
|
||
|
|
# x = x[torch.isfinite(x).all(1)]
|
||
|
|
|
||
|
|
# Check shape
|
||
|
|
n = x.shape[0] # number of boxes
|
||
|
|
if not n: # no boxes
|
||
|
|
continue
|
||
|
|
elif n > max_nms: # excess boxes
|
||
|
|
x = x[x[:, 5].argsort(descending=True)[:max_nms]] # sort by confidence
|
||
|
|
|
||
|
|
# Batched NMS
|
||
|
|
c = x[:, 6:7] * (0 if agnostic else max_wh) # classes
|
||
|
|
rboxes = x[:, :5].clone()
|
||
|
|
rboxes[:, :2] = rboxes[:, :2] + c # rboxes (offset by class)
|
||
|
|
scores = x[:, 5] # scores
|
||
|
|
_, i = obb_nms(rboxes, scores, iou_thres)
|
||
|
|
if i.shape[0] > max_det: # limit detections
|
||
|
|
i = i[:max_det]
|
||
|
|
|
||
|
|
output[xi] = x[i]
|
||
|
|
if (time.time() - t) > time_limit:
|
||
|
|
print(f'WARNING: NMS time limit {time_limit}s exceeded')
|
||
|
|
break # time limit exceeded
|
||
|
|
|
||
|
|
return output
|
||
|
|
|
||
|
|
def obb_nms(dets, scores, iou_thr, device_id=None):
|
||
|
|
"""
|
||
|
|
RIoU NMS - iou_thr.
|
||
|
|
Args:
|
||
|
|
dets (tensor/array): (num, [cx cy w h θ]) θ∈[-pi/2, pi/2)
|
||
|
|
scores (tensor/array): (num)
|
||
|
|
iou_thr (float): (1)
|
||
|
|
Returns:
|
||
|
|
dets (tensor): (n_nms, [cx cy w h θ])
|
||
|
|
inds (tensor): (n_nms), nms index of dets
|
||
|
|
"""
|
||
|
|
if isinstance(dets, torch.Tensor):
|
||
|
|
is_numpy = False
|
||
|
|
dets_th = dets
|
||
|
|
elif isinstance(dets, np.ndarray):
|
||
|
|
is_numpy = True
|
||
|
|
device = 'cpu' if device_id is None else f'cuda:{device_id}'
|
||
|
|
dets_th = torch.from_numpy(dets).to(device)
|
||
|
|
else:
|
||
|
|
raise TypeError('dets must be eithr a Tensor or numpy array, '
|
||
|
|
f'but got {type(dets)}')
|
||
|
|
|
||
|
|
if dets_th.numel() == 0: # len(dets)
|
||
|
|
inds = dets_th.new_zeros(0, dtype=torch.int64)
|
||
|
|
else:
|
||
|
|
# same bug will happen when bboxes is too small
|
||
|
|
too_small = dets_th[:, [2, 3]].min(1)[0] < 0.001 # [n]
|
||
|
|
if too_small.all(): # all the bboxes is too small
|
||
|
|
inds = dets_th.new_zeros(0, dtype=torch.int64)
|
||
|
|
else:
|
||
|
|
ori_inds = torch.arange(dets_th.size(0)) # 0 ~ n-1
|
||
|
|
ori_inds = ori_inds[~too_small]
|
||
|
|
dets_th = dets_th[~too_small] # (n_filter, 5)
|
||
|
|
scores = scores[~too_small]
|
||
|
|
|
||
|
|
inds = nms_rotated_ext.nms_rotated(dets_th, scores, iou_thr)
|
||
|
|
inds = ori_inds[inds]
|
||
|
|
|
||
|
|
if is_numpy:
|
||
|
|
inds = inds.cpu().numpy()
|
||
|
|
return dets[inds, :], inds
|
||
|
|
|
||
|
|
|
||
|
|
def poly_nms(dets, iou_thr, device_id=None):
|
||
|
|
if isinstance(dets, torch.Tensor):
|
||
|
|
is_numpy = False
|
||
|
|
dets_th = dets
|
||
|
|
elif isinstance(dets, np.ndarray):
|
||
|
|
is_numpy = True
|
||
|
|
device = 'cpu' if device_id is None else f'cuda:{device_id}'
|
||
|
|
dets_th = torch.from_numpy(dets).to(device)
|
||
|
|
else:
|
||
|
|
raise TypeError('dets must be eithr a Tensor or numpy array, '
|
||
|
|
f'but got {type(dets)}')
|
||
|
|
|
||
|
|
if dets_th.device == torch.device('cpu'):
|
||
|
|
raise NotImplementedError
|
||
|
|
inds = nms_rotated_ext.nms_poly(dets_th.float(), iou_thr)
|
||
|
|
|
||
|
|
if is_numpy:
|
||
|
|
inds = inds.cpu().numpy()
|
||
|
|
return dets[inds, :], inds
|
||
|
|
|
||
|
|
|
||
|
|
def box_iou(box1, box2, eps=1e-7):
|
||
|
|
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
|
||
|
|
"""
|
||
|
|
Return intersection-over-union (Jaccard index) of boxes.
|
||
|
|
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
|
||
|
|
Arguments:
|
||
|
|
box1 (Tensor[N, 4])
|
||
|
|
box2 (Tensor[M, 4])
|
||
|
|
Returns:
|
||
|
|
iou (Tensor[N, M]): the NxM matrix containing the pairwise
|
||
|
|
IoU values for every element in boxes1 and boxes2
|
||
|
|
"""
|
||
|
|
|
||
|
|
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
||
|
|
(a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
|
||
|
|
inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp(0).prod(2)
|
||
|
|
|
||
|
|
# IoU = inter / (area1 + area2 - inter)
|
||
|
|
return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps)
|
||
|
|
|
||
|
|
|
||
|
|
def rbox2poly(obboxes):
|
||
|
|
"""
|
||
|
|
Trans rbox format to poly format.
|
||
|
|
Args:
|
||
|
|
rboxes (array/tensor): (num_gts, [cx cy l s θ]) θ∈[-pi/2, pi/2)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
polys (array/tensor): (num_gts, [x1 y1 x2 y2 x3 y3 x4 y4])
|
||
|
|
"""
|
||
|
|
if isinstance(obboxes, torch.Tensor):
|
||
|
|
center, w, h, theta = obboxes[:, :2], obboxes[:, 2:3], obboxes[:, 3:4], obboxes[:, 4:5]
|
||
|
|
Cos, Sin = torch.cos(theta), torch.sin(theta)
|
||
|
|
|
||
|
|
vector1 = torch.cat(
|
||
|
|
(w/2 * Cos, -w/2 * Sin), dim=-1)
|
||
|
|
vector2 = torch.cat(
|
||
|
|
(-h/2 * Sin, -h/2 * Cos), dim=-1)
|
||
|
|
point1 = center + vector1 + vector2
|
||
|
|
point2 = center + vector1 - vector2
|
||
|
|
point3 = center - vector1 - vector2
|
||
|
|
point4 = center - vector1 + vector2
|
||
|
|
order = obboxes.shape[:-1]
|
||
|
|
return torch.cat(
|
||
|
|
(point1, point2, point3, point4), dim=-1).reshape(*order, 8)
|
||
|
|
else:
|
||
|
|
center, w, h, theta = np.split(obboxes, (2, 3, 4), axis=-1)
|
||
|
|
Cos, Sin = np.cos(theta), np.sin(theta)
|
||
|
|
|
||
|
|
vector1 = np.concatenate(
|
||
|
|
[w/2 * Cos, -w/2 * Sin], axis=-1)
|
||
|
|
vector2 = np.concatenate(
|
||
|
|
[-h/2 * Sin, -h/2 * Cos], axis=-1)
|
||
|
|
|
||
|
|
point1 = center + vector1 + vector2
|
||
|
|
point2 = center + vector1 - vector2
|
||
|
|
point3 = center - vector1 - vector2
|
||
|
|
point4 = center - vector1 + vector2
|
||
|
|
order = obboxes.shape[:-1]
|
||
|
|
return np.concatenate(
|
||
|
|
[point1, point2, point3, point4], axis=-1).reshape(*order, 8)
|
||
|
|
|
||
|
|
def scale_polys(img1_shape, polys, img0_shape, ratio_pad=None):
|
||
|
|
# ratio_pad: [(h_raw, w_raw), (hw_ratios, wh_paddings)]
|
||
|
|
# Rescale coords (xyxyxyxy) from img1_shape to img0_shape
|
||
|
|
if ratio_pad is None: # calculate from img0_shape
|
||
|
|
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = resized / raw
|
||
|
|
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
||
|
|
else:
|
||
|
|
gain = ratio_pad[0][0] # h_ratios
|
||
|
|
pad = ratio_pad[1] # wh_paddings
|
||
|
|
|
||
|
|
polys[:, [0, 2, 4, 6]] -= pad[0] # x padding
|
||
|
|
polys[:, [1, 3, 5, 7]] -= pad[1] # y padding
|
||
|
|
polys[:, :8] /= gain # Rescale poly shape to img0_shape
|
||
|
|
#clip_polys(polys, img0_shape)
|
||
|
|
return polys
|