mirror of https://github.com/WongKinYiu/yolov7.git
1034 lines
42 KiB
Python
1034 lines
42 KiB
Python
# YOLOR general utils
|
|
|
|
import glob
|
|
import logging
|
|
import math
|
|
import os
|
|
import platform
|
|
import random
|
|
import re
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
import torchvision
|
|
import yaml
|
|
|
|
from utils.google_utils import gsutil_getsize
|
|
from utils.metrics import fitness
|
|
from utils.torch_utils import init_torch_seeds
|
|
|
|
from utils.torch_utils import is_parallel
|
|
from torch.nn import functional as F
|
|
from detectron2.structures.masks import BitMasks
|
|
from detectron2.structures import Boxes
|
|
from detectron2.layers.roi_align import ROIAlign
|
|
from detectron2.utils.memory import retry_if_cuda_oom
|
|
from detectron2.layers import paste_masks_in_image
|
|
|
|
# Settings
|
|
torch.set_printoptions(linewidth=320, precision=5, profile='long')
|
|
np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5
|
|
pd.options.display.max_columns = 10
|
|
cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
|
|
os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads
|
|
|
|
|
|
def set_logging(rank=-1):
|
|
logging.basicConfig(
|
|
format="%(message)s",
|
|
level=logging.INFO if rank in [-1, 0] else logging.WARN)
|
|
|
|
|
|
def init_seeds(seed=0):
|
|
# Initialize random number generator (RNG) seeds
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
init_torch_seeds(seed)
|
|
|
|
|
|
def merge_bases(rois, coeffs, attn_r, num_b, location_to_inds=None):
|
|
# merge predictions
|
|
# N = coeffs.size(0)
|
|
if location_to_inds is not None:
|
|
rois = rois[location_to_inds]
|
|
N, B, H, W = rois.size()
|
|
if coeffs.dim() != 4:
|
|
coeffs = coeffs.view(N, num_b, attn_r, attn_r)
|
|
# NA = coeffs.shape[1] // B
|
|
coeffs = F.interpolate(coeffs, (H, W),
|
|
mode="bilinear").softmax(dim=1)
|
|
# coeffs = coeffs.view(N, -1, B, H, W)
|
|
# rois = rois[:, None, ...].repeat(1, NA, 1, 1, 1)
|
|
# masks_preds, _ = (rois * coeffs).sum(dim=2) # c.max(dim=1)
|
|
masks_preds = (rois * coeffs).sum(dim=1)
|
|
return masks_preds
|
|
|
|
|
|
def get_latest_run(search_dir='.'):
|
|
# Return path to most recent 'last.pt' in /runs (i.e. to --resume from)
|
|
last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True)
|
|
return max(last_list, key=os.path.getctime) if last_list else ''
|
|
|
|
|
|
def isdocker():
|
|
# Is environment a Docker container
|
|
return Path('/workspace').exists() # or Path('/.dockerenv').exists()
|
|
|
|
|
|
def emojis(str=''):
|
|
# Return platform-dependent emoji-safe version of string
|
|
return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
|
|
|
|
|
|
def check_online():
|
|
# Check internet connectivity
|
|
import socket
|
|
try:
|
|
socket.create_connection(("1.1.1.1", 443), 5) # check host accesability
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def check_git_status():
|
|
# Recommend 'git pull' if code is out of date
|
|
print(colorstr('github: '), end='')
|
|
try:
|
|
assert Path('.git').exists(), 'skipping check (not a git repository)'
|
|
assert not isdocker(), 'skipping check (Docker image)'
|
|
assert check_online(), 'skipping check (offline)'
|
|
|
|
cmd = 'git fetch && git config --get remote.origin.url'
|
|
url = subprocess.check_output(cmd, shell=True).decode().strip().rstrip('.git') # github repo url
|
|
branch = subprocess.check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out
|
|
n = int(subprocess.check_output(f'git rev-list {branch}..origin/master --count', shell=True)) # commits behind
|
|
if n > 0:
|
|
s = f"⚠️ WARNING: code is out of date by {n} commit{'s' * (n > 1)}. " \
|
|
f"Use 'git pull' to update or 'git clone {url}' to download latest."
|
|
else:
|
|
s = f'up to date with {url} ✅'
|
|
print(emojis(s)) # emoji-safe
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
|
|
def check_requirements(requirements='requirements.txt', exclude=()):
|
|
# Check installed dependencies meet requirements (pass *.txt file or list of packages)
|
|
import pkg_resources as pkg
|
|
prefix = colorstr('red', 'bold', 'requirements:')
|
|
if isinstance(requirements, (str, Path)): # requirements.txt file
|
|
file = Path(requirements)
|
|
if not file.exists():
|
|
print(f"{prefix} {file.resolve()} not found, check failed.")
|
|
return
|
|
requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(file.open()) if x.name not in exclude]
|
|
else: # list or tuple of packages
|
|
requirements = [x for x in requirements if x not in exclude]
|
|
|
|
n = 0 # number of packages updates
|
|
for r in requirements:
|
|
try:
|
|
pkg.require(r)
|
|
except Exception as e: # DistributionNotFound or VersionConflict if requirements not met
|
|
n += 1
|
|
print(f"{prefix} {e.req} not found and is required by YOLOR, attempting auto-update...")
|
|
print(subprocess.check_output(f"pip install '{e.req}'", shell=True).decode())
|
|
|
|
if n: # if packages updated
|
|
source = file.resolve() if 'file' in locals() else requirements
|
|
s = f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" \
|
|
f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
|
|
print(emojis(s)) # emoji-safe
|
|
|
|
|
|
def check_img_size(img_size, s=32):
|
|
# Verify img_size is a multiple of stride s
|
|
new_size = make_divisible(img_size, int(s)) # ceil gs-multiple
|
|
if new_size != img_size:
|
|
print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size))
|
|
return new_size
|
|
|
|
|
|
def check_imshow():
|
|
# Check if environment supports image displays
|
|
try:
|
|
assert not isdocker(), 'cv2.imshow() is disabled in Docker environments'
|
|
cv2.imshow('test', np.zeros((1, 1, 3)))
|
|
cv2.waitKey(1)
|
|
cv2.destroyAllWindows()
|
|
cv2.waitKey(1)
|
|
return True
|
|
except Exception as e:
|
|
print(f'WARNING: Environment does not support cv2.imshow() or PIL Image.show() image displays\n{e}')
|
|
return False
|
|
|
|
|
|
def check_file(file):
|
|
# Search for file if not found
|
|
if Path(file).is_file() or file == '':
|
|
return file
|
|
else:
|
|
files = glob.glob('./**/' + file, recursive=True) # find file
|
|
assert len(files), f'File Not Found: {file}' # assert file was found
|
|
assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
|
|
return files[0] # return file
|
|
|
|
|
|
def check_dataset(dict):
|
|
# Download dataset if not found locally
|
|
val, s = dict.get('val'), dict.get('download')
|
|
if val and len(val):
|
|
val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
|
|
if not all(x.exists() for x in val):
|
|
print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
|
|
if s and len(s): # download script
|
|
print('Downloading %s ...' % s)
|
|
if s.startswith('http') and s.endswith('.zip'): # URL
|
|
f = Path(s).name # filename
|
|
torch.hub.download_url_to_file(s, f)
|
|
r = os.system('unzip -q %s -d ../ && rm %s' % (f, f)) # unzip
|
|
else: # bash script
|
|
r = os.system(s)
|
|
print('Dataset autodownload %s\n' % ('success' if r == 0 else 'failure')) # analyze return value
|
|
else:
|
|
raise Exception('Dataset not found.')
|
|
|
|
|
|
def make_divisible(x, divisor):
|
|
# Returns x evenly divisible by divisor
|
|
return math.ceil(x / divisor) * divisor
|
|
|
|
|
|
def clean_str(s):
|
|
# Cleans a string by replacing special characters with underscore _
|
|
return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)
|
|
|
|
|
|
def one_cycle(y1=0.0, y2=1.0, steps=100):
|
|
# lambda function for sinusoidal ramp from y1 to y2
|
|
return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1
|
|
|
|
|
|
def colorstr(*input):
|
|
# Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world')
|
|
*args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string
|
|
colors = {'black': '\033[30m', # basic colors
|
|
'red': '\033[31m',
|
|
'green': '\033[32m',
|
|
'yellow': '\033[33m',
|
|
'blue': '\033[34m',
|
|
'magenta': '\033[35m',
|
|
'cyan': '\033[36m',
|
|
'white': '\033[37m',
|
|
'bright_black': '\033[90m', # bright colors
|
|
'bright_red': '\033[91m',
|
|
'bright_green': '\033[92m',
|
|
'bright_yellow': '\033[93m',
|
|
'bright_blue': '\033[94m',
|
|
'bright_magenta': '\033[95m',
|
|
'bright_cyan': '\033[96m',
|
|
'bright_white': '\033[97m',
|
|
'end': '\033[0m', # misc
|
|
'bold': '\033[1m',
|
|
'underline': '\033[4m'}
|
|
return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
|
|
|
|
|
|
def labels_to_class_weights(labels, nc=80):
|
|
# Get class weights (inverse frequency) from training labels
|
|
if labels[0] is None: # no labels loaded
|
|
return torch.Tensor()
|
|
|
|
labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO
|
|
classes = labels[:, 0].astype(np.int) # labels = [class xywh]
|
|
weights = np.bincount(classes, minlength=nc) # occurrences per class
|
|
|
|
# Prepend gridpoint count (for uCE training)
|
|
# gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum() # gridpoints per image
|
|
# weights = np.hstack([gpi * len(labels) - weights.sum() * 9, weights * 9]) ** 0.5 # prepend gridpoints to start
|
|
|
|
weights[weights == 0] = 1 # replace empty bins with 1
|
|
weights = 1 / weights # number of targets per class
|
|
weights /= weights.sum() # normalize
|
|
return torch.from_numpy(weights)
|
|
|
|
|
|
def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
|
|
# Produces image weights based on class_weights and image contents
|
|
class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels])
|
|
image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)
|
|
# index = random.choices(range(n), weights=image_weights, k=1) # weight image sample
|
|
return image_weights
|
|
|
|
|
|
def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper)
|
|
# https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
|
|
# a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
|
|
# b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
|
|
# x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco
|
|
# x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet
|
|
x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
|
|
35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
|
|
64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
|
|
return x
|
|
|
|
|
|
def xyxy2xywh(x):
|
|
# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] 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 # x center
|
|
y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center
|
|
y[:, 2] = x[:, 2] - x[:, 0] # width
|
|
y[:, 3] = x[:, 3] - x[:, 1] # height
|
|
return y
|
|
|
|
|
|
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 xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
|
|
# Convert nx4 boxes from [x, y, w, h] normalized 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] = w * (x[:, 0] - x[:, 2] / 2) + padw # top left x
|
|
y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh # top left y
|
|
y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw # bottom right x
|
|
y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh # bottom right y
|
|
return y
|
|
|
|
|
|
def xyn2xy(x, w=640, h=640, padw=0, padh=0):
|
|
# Convert normalized segments into pixel segments, shape (n,2)
|
|
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
|
y[:, 0] = w * x[:, 0] + padw # top left x
|
|
y[:, 1] = h * x[:, 1] + padh # top left y
|
|
return y
|
|
|
|
|
|
def segment2box(segment, width=640, height=640):
|
|
# Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
|
|
x, y = segment.T # segment xy
|
|
inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height)
|
|
x, y, = x[inside], y[inside]
|
|
return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4)) # xyxy
|
|
|
|
|
|
def segments2boxes(segments):
|
|
# Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
|
|
boxes = []
|
|
for s in segments:
|
|
x, y = s.T # segment xy
|
|
boxes.append([x.min(), y.min(), x.max(), y.max()]) # cls, xyxy
|
|
return xyxy2xywh(np.array(boxes)) # cls, xywh
|
|
|
|
|
|
def resample_segments(segments, n=1000):
|
|
# Up-sample an (n,2) segment
|
|
for i, s in enumerate(segments):
|
|
x = np.linspace(0, len(s) - 1, n)
|
|
xp = np.arange(len(s))
|
|
segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T # segment xy
|
|
return segments
|
|
|
|
|
|
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
|
|
# 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]
|
|
|
|
coords[:, [0, 2]] -= pad[0] # x padding
|
|
coords[:, [1, 3]] -= pad[1] # y padding
|
|
coords[:, :4] /= gain
|
|
clip_coords(coords, img0_shape)
|
|
return coords
|
|
|
|
|
|
def clip_coords(boxes, img_shape):
|
|
# Clip bounding xyxy bounding boxes to image shape (height, width)
|
|
boxes[:, 0].clamp_(0, img_shape[1]) # x1
|
|
boxes[:, 1].clamp_(0, img_shape[0]) # y1
|
|
boxes[:, 2].clamp_(0, img_shape[1]) # x2
|
|
boxes[:, 3].clamp_(0, img_shape[0]) # y2
|
|
|
|
|
|
def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
|
|
# Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
|
|
box2 = box2.T
|
|
|
|
# Get the coordinates of bounding boxes
|
|
if x1y1x2y2: # x1, y1, x2, y2 = box1
|
|
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 xywh to xyxy
|
|
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
|
|
|
|
# Intersection area
|
|
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
|
|
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
|
|
|
|
# Union Area
|
|
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
|
|
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
|
|
union = w1 * h1 + w2 * h2 - inter + eps
|
|
|
|
iou = inter / union
|
|
|
|
if GIoU or DIoU or CIoU:
|
|
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
|
|
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
|
|
if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
|
|
c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared
|
|
rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
|
|
(b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared
|
|
if DIoU:
|
|
return iou - rho2 / c2 # DIoU
|
|
elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
|
|
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / (h2 + eps)) - torch.atan(w1 / (h1 + eps)), 2)
|
|
with torch.no_grad():
|
|
alpha = v / (v - iou + (1 + eps))
|
|
return iou - (rho2 / c2 + v * alpha) # CIoU
|
|
else: # GIoU https://arxiv.org/pdf/1902.09630.pdf
|
|
c_area = cw * ch + eps # convex area
|
|
return iou - (c_area - union) / c_area # GIoU
|
|
else:
|
|
return iou # IoU
|
|
|
|
|
|
|
|
|
|
def bbox_alpha_iou(box1, box2, x1y1x2y2=False, GIoU=False, DIoU=False, CIoU=False, alpha=2, eps=1e-9):
|
|
# Returns tsqrt_he IoU of box1 to box2. box1 is 4, box2 is nx4
|
|
box2 = box2.T
|
|
|
|
# Get the coordinates of bounding boxes
|
|
if x1y1x2y2: # x1, y1, x2, y2 = box1
|
|
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 xywh to xyxy
|
|
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
|
|
|
|
# Intersection area
|
|
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
|
|
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
|
|
|
|
# Union Area
|
|
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
|
|
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
|
|
union = w1 * h1 + w2 * h2 - inter + eps
|
|
|
|
# change iou into pow(iou+eps)
|
|
# iou = inter / union
|
|
iou = torch.pow(inter/union + eps, alpha)
|
|
# beta = 2 * alpha
|
|
if GIoU or DIoU or CIoU:
|
|
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
|
|
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
|
|
if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
|
|
c2 = (cw ** 2 + ch ** 2) ** alpha + eps # convex diagonal
|
|
rho_x = torch.abs(b2_x1 + b2_x2 - b1_x1 - b1_x2)
|
|
rho_y = torch.abs(b2_y1 + b2_y2 - b1_y1 - b1_y2)
|
|
rho2 = ((rho_x ** 2 + rho_y ** 2) / 4) ** alpha # center distance
|
|
if DIoU:
|
|
return iou - rho2 / c2 # DIoU
|
|
elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
|
|
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
|
|
with torch.no_grad():
|
|
alpha_ciou = v / ((1 + eps) - inter / union + v)
|
|
# return iou - (rho2 / c2 + v * alpha_ciou) # CIoU
|
|
return iou - (rho2 / c2 + torch.pow(v * alpha_ciou + eps, alpha)) # CIoU
|
|
else: # GIoU https://arxiv.org/pdf/1902.09630.pdf
|
|
# c_area = cw * ch + eps # convex area
|
|
# return iou - (c_area - union) / c_area # GIoU
|
|
c_area = torch.max(cw * ch + eps, union) # convex area
|
|
return iou - torch.pow((c_area - union) / c_area + eps, alpha) # GIoU
|
|
else:
|
|
return iou # torch.log(iou+eps) or iou
|
|
|
|
|
|
def box_iou(box1, box2):
|
|
# 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
|
|
"""
|
|
|
|
def box_area(box):
|
|
# box = 4xn
|
|
return (box[2] - box[0]) * (box[3] - box[1])
|
|
|
|
area1 = box_area(box1.T)
|
|
area2 = box_area(box2.T)
|
|
|
|
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
|
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
|
return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
|
|
|
|
|
|
def wh_iou(wh1, wh2):
|
|
# Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
|
|
wh1 = wh1[:, None] # [N,1,2]
|
|
wh2 = wh2[None] # [1,M,2]
|
|
inter = torch.min(wh1, wh2).prod(2) # [N,M]
|
|
return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
|
|
|
|
|
|
def box_giou(box1, box2):
|
|
"""
|
|
Return generalized intersection-over-union (Jaccard index) between two sets of boxes.
|
|
Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
|
|
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
|
|
Args:
|
|
boxes1 (Tensor[N, 4]): first set of boxes
|
|
boxes2 (Tensor[M, 4]): second set of boxes
|
|
Returns:
|
|
Tensor[N, M]: the NxM matrix containing the pairwise generalized IoU values
|
|
for every element in boxes1 and boxes2
|
|
"""
|
|
|
|
def box_area(box):
|
|
# box = 4xn
|
|
return (box[2] - box[0]) * (box[3] - box[1])
|
|
|
|
area1 = box_area(box1.T)
|
|
area2 = box_area(box2.T)
|
|
|
|
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
|
union = (area1[:, None] + area2 - inter)
|
|
|
|
iou = inter / union
|
|
|
|
lti = torch.min(box1[:, None, :2], box2[:, :2])
|
|
rbi = torch.max(box1[:, None, 2:], box2[:, 2:])
|
|
|
|
whi = (rbi - lti).clamp(min=0) # [N,M,2]
|
|
areai = whi[:, :, 0] * whi[:, :, 1]
|
|
|
|
return iou - (areai - union) / areai
|
|
|
|
|
|
def box_ciou(box1, box2, eps: float = 1e-7):
|
|
"""
|
|
Return complete intersection-over-union (Jaccard index) between two sets of boxes.
|
|
Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
|
|
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
|
|
Args:
|
|
boxes1 (Tensor[N, 4]): first set of boxes
|
|
boxes2 (Tensor[M, 4]): second set of boxes
|
|
eps (float, optional): small number to prevent division by zero. Default: 1e-7
|
|
Returns:
|
|
Tensor[N, M]: the NxM matrix containing the pairwise complete IoU values
|
|
for every element in boxes1 and boxes2
|
|
"""
|
|
|
|
def box_area(box):
|
|
# box = 4xn
|
|
return (box[2] - box[0]) * (box[3] - box[1])
|
|
|
|
area1 = box_area(box1.T)
|
|
area2 = box_area(box2.T)
|
|
|
|
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
|
union = (area1[:, None] + area2 - inter)
|
|
|
|
iou = inter / union
|
|
|
|
lti = torch.min(box1[:, None, :2], box2[:, :2])
|
|
rbi = torch.max(box1[:, None, 2:], box2[:, 2:])
|
|
|
|
whi = (rbi - lti).clamp(min=0) # [N,M,2]
|
|
diagonal_distance_squared = (whi[:, :, 0] ** 2) + (whi[:, :, 1] ** 2) + eps
|
|
|
|
# centers of boxes
|
|
x_p = (box1[:, None, 0] + box1[:, None, 2]) / 2
|
|
y_p = (box1[:, None, 1] + box1[:, None, 3]) / 2
|
|
x_g = (box2[:, 0] + box2[:, 2]) / 2
|
|
y_g = (box2[:, 1] + box2[:, 3]) / 2
|
|
# The distance between boxes' centers squared.
|
|
centers_distance_squared = (x_p - x_g) ** 2 + (y_p - y_g) ** 2
|
|
|
|
w_pred = box1[:, None, 2] - box1[:, None, 0]
|
|
h_pred = box1[:, None, 3] - box1[:, None, 1]
|
|
|
|
w_gt = box2[:, 2] - box2[:, 0]
|
|
h_gt = box2[:, 3] - box2[:, 1]
|
|
|
|
v = (4 / (torch.pi ** 2)) * torch.pow((torch.atan(w_gt / h_gt) - torch.atan(w_pred / h_pred)), 2)
|
|
with torch.no_grad():
|
|
alpha = v / (1 - iou + v + eps)
|
|
return iou - (centers_distance_squared / diagonal_distance_squared) - alpha * v
|
|
|
|
|
|
def box_diou(box1, box2, eps: float = 1e-7):
|
|
"""
|
|
Return distance intersection-over-union (Jaccard index) between two sets of boxes.
|
|
Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
|
|
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
|
|
Args:
|
|
boxes1 (Tensor[N, 4]): first set of boxes
|
|
boxes2 (Tensor[M, 4]): second set of boxes
|
|
eps (float, optional): small number to prevent division by zero. Default: 1e-7
|
|
Returns:
|
|
Tensor[N, M]: the NxM matrix containing the pairwise distance IoU values
|
|
for every element in boxes1 and boxes2
|
|
"""
|
|
|
|
def box_area(box):
|
|
# box = 4xn
|
|
return (box[2] - box[0]) * (box[3] - box[1])
|
|
|
|
area1 = box_area(box1.T)
|
|
area2 = box_area(box2.T)
|
|
|
|
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
|
union = (area1[:, None] + area2 - inter)
|
|
|
|
iou = inter / union
|
|
|
|
lti = torch.min(box1[:, None, :2], box2[:, :2])
|
|
rbi = torch.max(box1[:, None, 2:], box2[:, 2:])
|
|
|
|
whi = (rbi - lti).clamp(min=0) # [N,M,2]
|
|
diagonal_distance_squared = (whi[:, :, 0] ** 2) + (whi[:, :, 1] ** 2) + eps
|
|
|
|
# centers of boxes
|
|
x_p = (box1[:, None, 0] + box1[:, None, 2]) / 2
|
|
y_p = (box1[:, None, 1] + box1[:, None, 3]) / 2
|
|
x_g = (box2[:, 0] + box2[:, 2]) / 2
|
|
y_g = (box2[:, 1] + box2[:, 3]) / 2
|
|
# The distance between boxes' centers squared.
|
|
centers_distance_squared = (x_p - x_g) ** 2 + (y_p - y_g) ** 2
|
|
|
|
# The distance IoU is the IoU penalized by a normalized
|
|
# distance between boxes' centers squared.
|
|
return iou - (centers_distance_squared / diagonal_distance_squared)
|
|
|
|
|
|
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
|
|
labels=()):
|
|
"""Runs Non-Maximum Suppression (NMS) on inference results
|
|
|
|
Returns:
|
|
list of detections, on (n,6) tensor per image [xyxy, conf, cls]
|
|
"""
|
|
|
|
nc = prediction.shape[2] - 5 # number of classes
|
|
xc = prediction[..., 4] > conf_thres # candidates
|
|
|
|
# Settings
|
|
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
|
max_det = 300 # maximum number of detections per image
|
|
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
|
time_limit = 10.0 # seconds to quit after
|
|
redundant = True # require redundant detections
|
|
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
|
|
merge = False # use merge-NMS
|
|
|
|
t = time.time()
|
|
output = [torch.zeros((0, 6), 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
|
|
|
|
# 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
|
|
if nc == 1:
|
|
x[:, 5:] = x[:, 4:5] # for models with one class, cls_loss is 0 and cls_conf is always 0.5,
|
|
# so there is no need to multiplicate.
|
|
else:
|
|
x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf
|
|
|
|
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
|
|
box = xywh2xyxy(x[:, :4])
|
|
|
|
# Detections matrix nx6 (xyxy, conf, cls)
|
|
if multi_label:
|
|
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
|
|
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
|
|
else: # best class only
|
|
conf, j = x[:, 5:].max(1, keepdim=True)
|
|
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
|
|
|
# Filter by class
|
|
if classes is not None:
|
|
x = x[(x[:, 5:6] == 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[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence
|
|
|
|
# Batched NMS
|
|
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
|
|
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
|
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
|
|
if i.shape[0] > max_det: # limit detections
|
|
i = i[:max_det]
|
|
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
|
|
# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
|
|
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
|
|
weights = iou * scores[None] # box weights
|
|
x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
|
|
if redundant:
|
|
i = i[iou.sum(1) > 1] # require redundancy
|
|
|
|
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 non_max_suppression_kpt(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
|
|
labels=(), kpt_label=False, nc=None, nkpt=None):
|
|
"""Runs Non-Maximum Suppression (NMS) on inference results
|
|
|
|
Returns:
|
|
list of detections, on (n,6) tensor per image [xyxy, conf, cls]
|
|
"""
|
|
if nc is None:
|
|
nc = prediction.shape[2] - 5 if not kpt_label else prediction.shape[2] - 56 # number of classes
|
|
xc = prediction[..., 4] > conf_thres # candidates
|
|
|
|
# Settings
|
|
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
|
max_det = 300 # maximum number of detections per image
|
|
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
|
time_limit = 10.0 # seconds to quit after
|
|
redundant = True # require redundant detections
|
|
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
|
|
merge = False # use merge-NMS
|
|
|
|
t = time.time()
|
|
output = [torch.zeros((0,6), 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
|
|
|
|
# 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:5+nc] *= x[:, 4:5] # conf = obj_conf * cls_conf
|
|
|
|
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
|
|
box = xywh2xyxy(x[:, :4])
|
|
|
|
# Detections matrix nx6 (xyxy, conf, cls)
|
|
if multi_label:
|
|
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
|
|
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
|
|
else: # best class only
|
|
if not kpt_label:
|
|
conf, j = x[:, 5:].max(1, keepdim=True)
|
|
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
|
else:
|
|
kpts = x[:, 6:]
|
|
conf, j = x[:, 5:6].max(1, keepdim=True)
|
|
x = torch.cat((box, conf, j.float(), kpts), 1)[conf.view(-1) > conf_thres]
|
|
|
|
|
|
# Filter by class
|
|
if classes is not None:
|
|
x = x[(x[:, 5:6] == 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[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence
|
|
|
|
# Batched NMS
|
|
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
|
|
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
|
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
|
|
if i.shape[0] > max_det: # limit detections
|
|
i = i[:max_det]
|
|
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
|
|
# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
|
|
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
|
|
weights = iou * scores[None] # box weights
|
|
x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
|
|
if redundant:
|
|
i = i[iou.sum(1) > 1] # require redundancy
|
|
|
|
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 non_max_suppression_mask_conf(prediction, attn, bases, pooler, hyp, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False, mask_iou=None, vote=False):
|
|
|
|
if prediction.dtype is torch.float16:
|
|
prediction = prediction.float() # to FP32
|
|
nc = prediction[0].shape[1] - 5 # number of classes
|
|
xc = prediction[..., 4] > conf_thres # candidates
|
|
# Settings
|
|
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
|
max_det = 300 # maximum number of detections per image
|
|
time_limit = 10.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 = [None] * prediction.shape[0]
|
|
output_mask = [None] * prediction.shape[0]
|
|
output_mask_score = [None] * prediction.shape[0]
|
|
output_ac = [None] * prediction.shape[0]
|
|
output_ab = [None] * prediction.shape[0]
|
|
|
|
def RMS_contrast(masks):
|
|
mu = torch.mean(masks, dim=-1, keepdim=True)
|
|
return torch.sqrt(torch.mean((masks - mu)**2, dim=-1, keepdim=True))
|
|
|
|
|
|
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
|
|
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
|
|
box = xywh2xyxy(x[:, :4])
|
|
|
|
# If none remain process next image
|
|
if not x.shape[0]:
|
|
continue
|
|
|
|
a = attn[xi][xc[xi]]
|
|
base = bases[xi]
|
|
|
|
bboxes = Boxes(box)
|
|
pooled_bases = pooler([base[None]], [bboxes])
|
|
|
|
pred_masks = merge_bases(pooled_bases, a, hyp["attn_resolution"], hyp["num_base"]).view(a.shape[0], -1).sigmoid()
|
|
|
|
if mask_iou is not None:
|
|
mask_score = mask_iou[xi][xc[xi]][..., None]
|
|
else:
|
|
temp = pred_masks.clone()
|
|
temp[temp < 0.5] = 1 - temp[temp < 0.5]
|
|
mask_score = torch.exp(torch.log(temp).mean(dim=-1, keepdims=True))#torch.mean(temp, dim=-1, keepdims=True)
|
|
|
|
x[:, 5:] *= x[:, 4:5] * mask_score # x[:, 4:5] * * mask_conf * non_mask_conf # conf = obj_conf * cls_conf
|
|
|
|
if multi_label:
|
|
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
|
|
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
|
|
mask_score = mask_score[i]
|
|
if attn is not None:
|
|
pred_masks = pred_masks[i]
|
|
else: # best class only
|
|
conf, j = x[:, 5:].max(1, keepdim=True)
|
|
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
|
|
|
# Filter by class
|
|
if classes:
|
|
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
|
|
|
|
|
|
# If none remain process next image
|
|
n = x.shape[0] # number of boxes
|
|
if not n:
|
|
continue
|
|
|
|
# Batched NMS
|
|
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
|
|
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
|
# scores *= mask_score
|
|
i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
|
|
if i.shape[0] > max_det: # limit detections
|
|
i = i[:max_det]
|
|
|
|
|
|
all_candidates = []
|
|
all_boxes = []
|
|
if vote:
|
|
ious = box_iou(boxes[i], boxes) > iou_thres
|
|
for iou in ious:
|
|
selected_masks = pred_masks[iou]
|
|
k = min(10, selected_masks.shape[0])
|
|
_, tfive = torch.topk(scores[iou], k)
|
|
all_candidates.append(pred_masks[iou][tfive])
|
|
all_boxes.append(x[iou, :4][tfive])
|
|
#exit()
|
|
|
|
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
|
|
try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
|
|
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
|
|
weights = iou * scores[None] # box weights
|
|
x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
|
|
if redundant:
|
|
i = i[iou.sum(1) > 1] # require redundancy
|
|
except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139
|
|
print(x, i, x.shape, i.shape)
|
|
pass
|
|
|
|
output[xi] = x[i]
|
|
output_mask_score[xi] = mask_score[i]
|
|
output_ac[xi] = all_candidates
|
|
output_ab[xi] = all_boxes
|
|
if attn is not None:
|
|
output_mask[xi] = pred_masks[i]
|
|
if (time.time() - t) > time_limit:
|
|
break # time limit exceeded
|
|
|
|
return output, output_mask, output_mask_score, output_ac, output_ab
|
|
|
|
|
|
def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_optimizer()
|
|
# Strip optimizer from 'f' to finalize training, optionally save as 's'
|
|
x = torch.load(f, map_location=torch.device('cpu'))
|
|
if x.get('ema'):
|
|
x['model'] = x['ema'] # replace model with ema
|
|
for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates': # keys
|
|
x[k] = None
|
|
x['epoch'] = -1
|
|
x['model'].half() # to FP16
|
|
for p in x['model'].parameters():
|
|
p.requires_grad = False
|
|
torch.save(x, s or f)
|
|
mb = os.path.getsize(s or f) / 1E6 # filesize
|
|
print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")
|
|
|
|
|
|
def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''):
|
|
# Print mutation results to evolve.txt (for use with train.py --evolve)
|
|
a = '%10s' * len(hyp) % tuple(hyp.keys()) # hyperparam keys
|
|
b = '%10.3g' * len(hyp) % tuple(hyp.values()) # hyperparam values
|
|
c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
|
|
print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c))
|
|
|
|
if bucket:
|
|
url = 'gs://%s/evolve.txt' % bucket
|
|
if gsutil_getsize(url) > (os.path.getsize('evolve.txt') if os.path.exists('evolve.txt') else 0):
|
|
os.system('gsutil cp %s .' % url) # download evolve.txt if larger than local
|
|
|
|
with open('evolve.txt', 'a') as f: # append result
|
|
f.write(c + b + '\n')
|
|
x = np.unique(np.loadtxt('evolve.txt', ndmin=2), axis=0) # load unique rows
|
|
x = x[np.argsort(-fitness(x))] # sort
|
|
np.savetxt('evolve.txt', x, '%10.3g') # save sort by fitness
|
|
|
|
# Save yaml
|
|
for i, k in enumerate(hyp.keys()):
|
|
hyp[k] = float(x[0, i + 7])
|
|
with open(yaml_file, 'w') as f:
|
|
results = tuple(x[0, :7])
|
|
c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
|
|
f.write('# Hyperparameter Evolution Results\n# Generations: %g\n# Metrics: ' % len(x) + c + '\n\n')
|
|
yaml.dump(hyp, f, sort_keys=False)
|
|
|
|
if bucket:
|
|
os.system('gsutil cp evolve.txt %s gs://%s' % (yaml_file, bucket)) # upload
|
|
|
|
|
|
def apply_classifier(x, model, img, im0):
|
|
# applies a second stage classifier to yolo outputs
|
|
im0 = [im0] if isinstance(im0, np.ndarray) else im0
|
|
for i, d in enumerate(x): # per image
|
|
if d is not None and len(d):
|
|
d = d.clone()
|
|
|
|
# Reshape and pad cutouts
|
|
b = xyxy2xywh(d[:, :4]) # boxes
|
|
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # rectangle to square
|
|
b[:, 2:] = b[:, 2:] * 1.3 + 30 # pad
|
|
d[:, :4] = xywh2xyxy(b).long()
|
|
|
|
# Rescale boxes from img_size to im0 size
|
|
scale_coords(img.shape[2:], d[:, :4], im0[i].shape)
|
|
|
|
# Classes
|
|
pred_cls1 = d[:, 5].long()
|
|
ims = []
|
|
for j, a in enumerate(d): # per item
|
|
cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])]
|
|
im = cv2.resize(cutout, (224, 224)) # BGR
|
|
# cv2.imwrite('test%i.jpg' % j, cutout)
|
|
|
|
im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
|
im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32
|
|
im /= 255.0 # 0 - 255 to 0.0 - 1.0
|
|
ims.append(im)
|
|
|
|
pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction
|
|
x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections
|
|
|
|
return x
|
|
|
|
|
|
def increment_path(path, exist_ok=True, sep=''):
|
|
# Increment path, i.e. runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc.
|
|
path = Path(path) # os-agnostic
|
|
if (path.exists() and exist_ok) or (not path.exists()):
|
|
return str(path)
|
|
else:
|
|
dirs = glob.glob(f"{path}{sep}*") # similar paths
|
|
matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
|
|
i = [int(m.groups()[0]) for m in matches if m] # indices
|
|
n = max(i) + 1 if i else 2 # increment number
|
|
return f"{path}{sep}{n}" # update path
|