|
|
import math
|
|
|
import threading
|
|
|
import time
|
|
|
import contextlib
|
|
|
import numpy as np
|
|
|
import torch
|
|
|
import torchvision
|
|
|
from models.common import DetectMultiBackend
|
|
|
import cv2
|
|
|
from utils.plots import Annotator
|
|
|
import pkg_resources as pkg
|
|
|
import platform
|
|
|
import logging
|
|
|
import logging.config
|
|
|
import os
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
from pathlib import Path
|
|
|
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)
|
|
|
|
|
|
|
|
|
LOGGING_NAME = "yolov5"
|
|
|
def set_logging(name=LOGGING_NAME, verbose=True):
|
|
|
# sets up logging for the given name
|
|
|
rank = int(os.getenv('RANK', -1)) # rank in world for Multi-GPU trainings
|
|
|
level = logging.INFO if verbose and rank in {-1, 0} else logging.ERROR
|
|
|
logging.config.dictConfig({
|
|
|
"version": 1,
|
|
|
"disable_existing_loggers": False,
|
|
|
"formatters": {
|
|
|
name: {
|
|
|
"format": "%(message)s"}},
|
|
|
"handlers": {
|
|
|
name: {
|
|
|
"class": "logging.StreamHandler",
|
|
|
"formatter": name,
|
|
|
"level": level,}},
|
|
|
"loggers": {
|
|
|
name: {
|
|
|
"level": level,
|
|
|
"handlers": [name],
|
|
|
"propagate": False,}}})
|
|
|
set_logging(LOGGING_NAME) # run before defining LOGGER
|
|
|
LOGGER = logging.getLogger(LOGGING_NAME) # define globally (used in train.py, val.py, detect.py, etc.)
|
|
|
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_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False, hard=False, verbose=False):
|
|
|
# Check version vs. required version
|
|
|
current, minimum = (pkg.parse_version(x) for x in (current, minimum))
|
|
|
result = (current == minimum) if pinned else (current >= minimum) # bool
|
|
|
s = f'WARNING ⚠️ {name}{minimum} is required by YOLOv5, but {name}{current} is currently installed' # string
|
|
|
if hard:
|
|
|
assert result, emojis(s) # assert min requirements met
|
|
|
if verbose and not result:
|
|
|
LOGGER.warning(s)
|
|
|
return result
|
|
|
|
|
|
|
|
|
def smart_inference_mode(torch_1_9=check_version(torch.__version__, '1.9.0')):
|
|
|
# Applies torch.inference_mode() decorator if torch>=1.9.0 else torch.no_grad() decorator
|
|
|
def decorate(fn):
|
|
|
return (torch.inference_mode if torch_1_9 else torch.no_grad)()(fn)
|
|
|
|
|
|
return decorate
|
|
|
|
|
|
|
|
|
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 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 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 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 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 letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
|
|
|
# Resize and pad image while meeting stride-multiple constraints
|
|
|
shape = im.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 val 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, stride), np.mod(dh, stride) # 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
|
|
|
im = cv2.resize(im, 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))
|
|
|
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
|
|
|
return im, ratio, (dw, dh)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def non_max_suppression(
|
|
|
prediction,
|
|
|
conf_thres=0.25,
|
|
|
iou_thres=0.45,
|
|
|
classes=None,
|
|
|
agnostic=False,
|
|
|
multi_label=False,
|
|
|
labels=(),
|
|
|
max_det=300,
|
|
|
nm=0, # number of masks
|
|
|
):
|
|
|
"""Non-Maximum Suppression (NMS) on inference results to reject overlapping detections
|
|
|
|
|
|
Returns:
|
|
|
list of detections, on (n,6) tensor per image [xyxy, conf, cls]
|
|
|
"""
|
|
|
# YOLOv5 model in validation model, output = (inference_out, loss_out)
|
|
|
prediction = prediction[0] # select only inference output
|
|
|
nc = prediction.shape[2] - nm - 5 # number of classes
|
|
|
xc = prediction[..., 4] > conf_thres # candidates
|
|
|
# Checks
|
|
|
mi = 5 + nc # mask start index
|
|
|
output = [torch.zeros((0, 6 + nm), device=prediction.device)]
|
|
|
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]):
|
|
|
lb = labels[xi]
|
|
|
v = torch.zeros((len(lb), nc + nm + 5), device=x.device)
|
|
|
v[:, :4] = lb[:, 1:5] # box
|
|
|
v[:, 4] = 1.0 # conf
|
|
|
v[range(len(lb)), lb[:, 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:] *= x[:, 4:5] # conf = obj_conf * cls_conf
|
|
|
|
|
|
# Box/Mask
|
|
|
box = xywh2xyxy(x[:, :4]) # center_x, center_y, width, height) to (x1, y1, x2, y2)
|
|
|
mask = x[:, mi:] # zero columns if no masks
|
|
|
|
|
|
# Detections matrix nx6 (xyxy, conf, cls)
|
|
|
|
|
|
conf, j = x[:, 5:mi].max(1, keepdim=True)
|
|
|
x = torch.cat((box, conf, j.float(), mask), 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
|
|
|
|
|
|
x = x[x[:, 4].argsort(descending=True)] # sort by confidence
|
|
|
|
|
|
# Batched NMS
|
|
|
c = x[:, 5:6] # classes
|
|
|
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
|
|
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
|
|
|
# limit detections
|
|
|
i = i[:max_det]
|
|
|
# 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
|
|
|
i = i[iou.sum(1) > 1] # require redundancy
|
|
|
|
|
|
output[xi] = x[i]
|
|
|
return output
|
|
|
|
|
|
def is_writeable(dir, test=False):
|
|
|
# Return True if directory has write permissions, test opening a file with write permissions if test=True
|
|
|
if not test:
|
|
|
return os.access(dir, os.W_OK) # possible issues on Windows
|
|
|
file = Path(dir) / 'tmp.txt'
|
|
|
try:
|
|
|
with open(file, 'w'): # open file with write permissions
|
|
|
pass
|
|
|
file.unlink() # remove file
|
|
|
return True
|
|
|
except OSError:
|
|
|
return False
|
|
|
def user_config_dir(dir='Ultralytics', env_var='YOLOV5_CONFIG_DIR'):
|
|
|
# Return path of user configuration directory. Prefer environment variable if exists. Make dir if required.
|
|
|
env = os.getenv(env_var)
|
|
|
if env:
|
|
|
path = Path(env) # use environment variable
|
|
|
else:
|
|
|
cfg = {'Windows': 'AppData/Roaming', 'Linux': '.config', 'Darwin': 'Library/Application Support'} # 3 OS dirs
|
|
|
path = Path.home() / cfg.get(platform.system(), '') # OS-specific config dir
|
|
|
path = (path if is_writeable(path) else Path('/tmp')) / dir # GCP and AWS lambda fix, only /tmp is writeable
|
|
|
path.mkdir(exist_ok=True) # make if required
|
|
|
return path
|
|
|
|
|
|
CONFIG_DIR = user_config_dir() # Ultralytics settings dir
|
|
|
FONT = 'Arial.ttf' # https://ultralytics.com/assets/Arial.ttf
|
|
|
|
|
|
def check_font(font=FONT, progress=False):
|
|
|
# Download font to CONFIG_DIR if necessary
|
|
|
font = Path(font)
|
|
|
file = CONFIG_DIR / font.name
|
|
|
if not font.exists() and not file.exists():
|
|
|
url = f'https://ultralytics.com/assets/{font.name}'
|
|
|
LOGGER.info(f'Downloading {url} to {file}...')
|
|
|
torch.hub.download_url_to_file(url, str(file), progress=progress)
|
|
|
def check_python(minimum='3.7.0'):
|
|
|
# Check current python version vs. required python version
|
|
|
check_version(platform.python_version(), minimum, name='Python ', hard=True)
|
|
|
|
|
|
class TryExcept(contextlib.ContextDecorator):
|
|
|
# YOLOv5 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager
|
|
|
def __init__(self, msg=''):
|
|
|
self.msg = msg
|
|
|
|
|
|
def __enter__(self):
|
|
|
pass
|
|
|
|
|
|
def __exit__(self, exc_type, value, traceback):
|
|
|
if value:
|
|
|
print(emojis(f"{self.msg}{': ' if self.msg else ''}{value}"))
|
|
|
return True
|
|
|
class TryExcept(contextlib.ContextDecorator):
|
|
|
# YOLOv5 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager
|
|
|
def __init__(self, msg=''):
|
|
|
self.msg = msg
|
|
|
|
|
|
def __enter__(self):
|
|
|
pass
|
|
|
|
|
|
def __exit__(self, exc_type, value, traceback):
|
|
|
if value:
|
|
|
print(emojis(f"{self.msg}{': ' if self.msg else ''}{value}"))
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
|
|
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']
|
|
|
|
|
|
|
|
|
@TryExcept()
|
|
|
# def check_requirements(requirements='requirements.txt', exclude=(), install=True, cmds=''):
|
|
|
# # Check installed dependencies meet YOLOv5 requirements (pass *.txt file or list of packages or single package str)
|
|
|
# prefix = colorstr('red', 'bold', 'requirements:')
|
|
|
# check_python() # check python version
|
|
|
# if isinstance(requirements, Path): # requirements.txt file
|
|
|
# file = requirements.resolve()
|
|
|
# assert file.exists(), f"{prefix} {file} not found, check failed."
|
|
|
# with file.open() as f:
|
|
|
# requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(f) if x.name not in exclude]
|
|
|
# elif isinstance(requirements, str):
|
|
|
# requirements = [requirements]
|
|
|
#
|
|
|
# s = ''
|
|
|
# n = 0
|
|
|
# for r in requirements:
|
|
|
# try:
|
|
|
# pkg.require(r)
|
|
|
# except (pkg.VersionConflict, pkg.DistributionNotFound): # exception if requirements not met
|
|
|
# s += f'"{r}" '
|
|
|
# n += 1
|
|
|
#
|
|
|
# if s and install and AUTOINSTALL: # check environment variable
|
|
|
# LOGGER.info(f"{prefix} YOLOv5 requirement{'s' * (n > 1)} {s}not found, attempting AutoUpdate...")
|
|
|
# try:
|
|
|
# # assert check_online(), "AutoUpdate skipped (offline)"
|
|
|
# LOGGER.info(check_output(f'pip install {s} {cmds}', shell=True).decode())
|
|
|
# source = file 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"
|
|
|
# LOGGER.info(s)
|
|
|
# except Exception as e:
|
|
|
# LOGGER.warning(f'{prefix} ❌ {e}')
|
|
|
|
|
|
# def check_pil_font(font=FONT, size=10):
|
|
|
# # Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary
|
|
|
# font = Path(font)
|
|
|
# font = font if font.exists() else (CONFIG_DIR / font.name)
|
|
|
# try:
|
|
|
# return ImageFont.truetype(str(font) if font.exists() else font.name, size)
|
|
|
# except Exception: # download if missing
|
|
|
# try:
|
|
|
# check_font(font)
|
|
|
# return ImageFont.truetype(str(font), size)
|
|
|
# except TypeError:
|
|
|
# check_requirements('Pillow>=8.4.0') # known issue https://github.com/ultralytics/yolov5/issues/5374
|
|
|
# except URLError: # not online
|
|
|
# return ImageFont.load_default()
|
|
|
|
|
|
|
|
|
|
|
|
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
|
|
|
# class Annotator:
|
|
|
# # YOLOv5 Annotator for train/val mosaics and jpgs and detect/hub inference annotations
|
|
|
# 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.'
|
|
|
# non_ascii = not is_ascii(example) # non-latin labels, i.e. asian, arabic, cyrillic
|
|
|
# self.pil = pil or non_ascii
|
|
|
# if self.pil: # use PIL
|
|
|
# self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
|
|
# self.draw = ImageDraw.Draw(self.im)
|
|
|
# self.font = check_pil_font(font='Arial.Unicode.ttf' if non_ascii else font,
|
|
|
# size=font_size or max(round(sum(self.im.size) / 2 * 0.035), 12))
|
|
|
# else: # use cv2
|
|
|
# self.im = 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
|
|
|
# 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 masks(self, masks, colors, im_gpu=None, alpha=0.5):
|
|
|
# """Plot masks at once.
|
|
|
# Args:
|
|
|
# masks (tensor): predicted masks on cuda, shape: [n, h, w]
|
|
|
# colors (List[List[Int]]): colors for predicted masks, [[r, g, b] * n]
|
|
|
# im_gpu (tensor): img is in cuda, shape: [3, h, w], range: [0, 1]
|
|
|
# alpha (float): mask transparency: 0.0 fully transparent, 1.0 opaque
|
|
|
# """
|
|
|
# if self.pil:
|
|
|
# # convert to numpy first
|
|
|
# self.im = np.asarray(self.im).copy()
|
|
|
# if im_gpu is None:
|
|
|
# # Add multiple masks of shape(h,w,n) with colors list([r,g,b], [r,g,b], ...)
|
|
|
# if len(masks) == 0:
|
|
|
# return
|
|
|
# if isinstance(masks, torch.Tensor):
|
|
|
# masks = torch.as_tensor(masks, dtype=torch.uint8)
|
|
|
# masks = masks.permute(1, 2, 0).contiguous()
|
|
|
# masks = masks.cpu().numpy()
|
|
|
# # masks = np.ascontiguousarray(masks.transpose(1, 2, 0))
|
|
|
# masks = scale_image(masks.shape[:2], masks, self.im.shape)
|
|
|
# masks = np.asarray(masks, dtype=np.float32)
|
|
|
# colors = np.asarray(colors, dtype=np.float32) # shape(n,3)
|
|
|
# s = masks.sum(2, keepdims=True).clip(0, 1) # add all masks together
|
|
|
# masks = (masks @ colors).clip(0, 255) # (h,w,n) @ (n,3) = (h,w,3)
|
|
|
# self.im[:] = masks * alpha + self.im * (1 - s * alpha)
|
|
|
# else:
|
|
|
# if len(masks) == 0:
|
|
|
# self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
|
|
|
# colors = torch.tensor(colors, device=im_gpu.device, dtype=torch.float32) / 255.0
|
|
|
# colors = colors[:, None, None] # shape(n,1,1,3)
|
|
|
# masks = masks.unsqueeze(3) # shape(n,h,w,1)
|
|
|
# masks_color = masks * (colors * alpha) # shape(n,h,w,3)
|
|
|
#
|
|
|
# inv_alph_masks = (1 - masks * alpha).cumprod(0) # shape(n,h,w,1)
|
|
|
# mcs = (masks_color * inv_alph_masks).sum(0) * 2 # mask color summand shape(n,h,w,3)
|
|
|
#
|
|
|
# im_gpu = im_gpu.flip(dims=[0]) # flip channel
|
|
|
# im_gpu = im_gpu.permute(1, 2, 0).contiguous() # shape(h,w,3)
|
|
|
# im_gpu = im_gpu * inv_alph_masks[-1] + mcs
|
|
|
# im_mask = (im_gpu * 255).byte().cpu().numpy()
|
|
|
# self.im[:] = scale_image(im_gpu.shape, im_mask, self.im.shape)
|
|
|
# if self.pil:
|
|
|
# # convert im back to PIL and update draw
|
|
|
# self.fromarray(self.im)
|
|
|
#
|
|
|
# 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), anchor='top'):
|
|
|
# # Add text to image (PIL-only)
|
|
|
# if anchor == 'bottom': # start y from font bottom
|
|
|
# w, h = self.font.getsize(text) # text width, height
|
|
|
# xy[1] += 1 - h
|
|
|
# self.draw.text(xy, text, fill=txt_color, font=self.font)
|
|
|
#
|
|
|
# def fromarray(self, im):
|
|
|
# # Update self.im from a numpy array
|
|
|
# self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
|
|
# self.draw = ImageDraw.Draw(self.im)
|
|
|
#
|
|
|
# def result(self):
|
|
|
# # Return annotated image as array
|
|
|
# return np.asarray(self.im)
|
|
|
|
|
|
@smart_inference_mode()
|
|
|
def run():
|
|
|
# Load model
|
|
|
device = torch.device('cuda:0')
|
|
|
model = DetectMultiBackend(weights='fire.pt', device=device, dnn=False, data=False, fp16=True)
|
|
|
stride, names, pt = model.stride, model.names, model.pt
|
|
|
#IP摄像头
|
|
|
cap = cv2.VideoCapture("http://admin:admin@10.129.50.72:8081") # URL of the IP camera
|
|
|
# 读取图片
|
|
|
while True:
|
|
|
ret, img0 = cap.read()
|
|
|
#图片旋转
|
|
|
img0 = np.rot90(img0,3)
|
|
|
|
|
|
img0 = np.array(img0)
|
|
|
img0 = cv2.cvtColor(img0, cv2.COLOR_BGRA2BGR)
|
|
|
|
|
|
|
|
|
# 处理图片
|
|
|
im = letterbox(img0, (640, 640), stride=32, auto=True)[0] # padded resize
|
|
|
im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
|
|
|
im = np.ascontiguousarray(im) # contiguous
|
|
|
im = torch.from_numpy(im).to(model.device)
|
|
|
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
|
|
|
im /= 255 # 0 - 255 to 0.0 - 1.0
|
|
|
if len(im.shape) == 3:
|
|
|
im = im[None] # expand for batch dim
|
|
|
|
|
|
# 推理
|
|
|
pred = model(im, augment=False, visualize=False)
|
|
|
# 非极大值抑制
|
|
|
pred = non_max_suppression(pred, conf_thres=0.4, iou_thres=0.05, classes=None, max_det=1000)
|
|
|
|
|
|
# 处理推理内容
|
|
|
|
|
|
for i, det in enumerate(pred):
|
|
|
# 画框
|
|
|
annotator = Annotator(img0, line_width=2)
|
|
|
if len(det):
|
|
|
target_list = []
|
|
|
# 将转换后的图片画框结果转换成原图上的结果
|
|
|
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], img0.shape).round()
|
|
|
for *xyxy, conf, cls in reversed(det): # 处理推理出来每个目标的信息
|
|
|
# 将xyxy(左上角+右下角)格式转为xywh(中心点+宽长)格式，并除上w，h做归一化，转化为列表再保存
|
|
|
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4))).view(-1).tolist() # normalized xywh
|
|
|
|
|
|
annotator.box_label(xyxy, label=f'[{names[int(cls)]} {conf:.2f}]',
|
|
|
color=(34, 139, 34),
|
|
|
txt_color=(0, 191, 255))
|
|
|
|
|
|
|
|
|
target_list.append(xywh)
|
|
|
|
|
|
|
|
|
|
|
|
im0 = annotator.result()
|
|
|
cv2.imshow('window', im0)
|
|
|
cv2.waitKey(1)
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
run()
|
