k20hcmus/FishEye8K
3
1import ast2import contextlib3import json4import math5import platform6import warnings7import zipfile8from collections import OrderedDict, namedtuple9from copy import copy10from pathlib import Path11from urllib.parse import urlparse12 13from typing import Optional14 15import cv216import numpy as np17import pandas as pd18import requests19import torch20import torch.nn as nn21from IPython.display import display22from PIL import Image23from torch.cuda import amp24 25from utils import TryExcept26from utils.dataloaders import exif_transpose, letterbox27from utils.general import (LOGGER, ROOT, Profile, check_requirements, check_suffix, check_version, colorstr,28 increment_path, is_notebook, make_divisible, non_max_suppression, scale_boxes,29 xywh2xyxy, xyxy2xywh, yaml_load)30from utils.plots import Annotator, colors, save_one_box31from utils.torch_utils import copy_attr, smart_inference_mode32 33 34def autopad(k, p=None, d=1): # kernel, padding, dilation35 # Pad to 'same' shape outputs36 if d > 1:37 k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size38 if p is None:39 p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad40 return p41 42 43class Conv(nn.Module):44 # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)45 default_act = nn.SiLU() # default activation46 47 def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):48 super().__init__()49 self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)50 self.bn = nn.BatchNorm2d(c2)51 self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()52 53 def forward(self, x):54 return self.act(self.bn(self.conv(x)))55 56 def forward_fuse(self, x):57 return self.act(self.conv(x))58 59class Convb(nn.Module):60 # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)61 default_act = nn.SiLU() # default activation62 63 def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):64 super().__init__()65 self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=True)66 self.bn = nn.BatchNorm2d(c2)67 self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()68 69 def forward(self, x):70 return self.act(self.bn(self.conv(x)))71 72 def forward_fuse(self, x):73 return self.act(self.conv(x))74 75 76class AConv(nn.Module):77 def __init__(self, c1, c2): # ch_in, ch_out, shortcut, kernels, groups, expand78 super().__init__()79 self.cv1 = Conv(c1, c2, 3, 2, 1)80 81 def forward(self, x):82 x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)83 return self.cv1(x)84 85 86class ADown(nn.Module):87 def __init__(self, c1, c2): # ch_in, ch_out, shortcut, kernels, groups, expand88 super().__init__()89 self.c = c2 // 290 self.cv1 = Conv(c1 // 2, self.c, 3, 2, 1)91 self.cv2 = Conv(c1 // 2, self.c, 1, 1, 0)92 93 def forward(self, x):94 x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)95 x1,x2 = x.chunk(2, 1)96 x1 = self.cv1(x1)97 x2 = torch.nn.functional.max_pool2d(x2, 3, 2, 1)98 x2 = self.cv2(x2)99 return torch.cat((x1, x2), 1)100 101 102class RepConvN(nn.Module):103 """RepConv is a basic rep-style block, including training and deploy status104 This code is based on https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py105 """106 default_act = nn.SiLU() # default activation107 108 def __init__(self, c1, c2, k=3, s=1, p=1, g=1, d=1, act=True, bn=False, deploy=False):109 super().__init__()110 assert k == 3 and p == 1111 self.g = g112 self.c1 = c1113 self.c2 = c2114 self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()115 116 self.bn = None117 self.conv1 = Conv(c1, c2, k, s, p=p, g=g, act=False)118 self.conv2 = Conv(c1, c2, 1, s, p=(p - k // 2), g=g, act=False)119 120 def forward_fuse(self, x):121 """Forward process"""122 return self.act(self.conv(x))123 124 def forward(self, x):125 """Forward process"""126 id_out = 0 if self.bn is None else self.bn(x)127 return self.act(self.conv1(x) + self.conv2(x) + id_out)128 129 def get_equivalent_kernel_bias(self):130 kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1)131 kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2)132 kernelid, biasid = self._fuse_bn_tensor(self.bn)133 return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid134 135 def _avg_to_3x3_tensor(self, avgp):136 channels = self.c1137 groups = self.g138 kernel_size = avgp.kernel_size139 input_dim = channels // groups140 k = torch.zeros((channels, input_dim, kernel_size, kernel_size))141 k[np.arange(channels), np.tile(np.arange(input_dim), groups), :, :] = 1.0 / kernel_size ** 2142 return k143 144 def _pad_1x1_to_3x3_tensor(self, kernel1x1):145 if kernel1x1 is None:146 return 0147 else:148 return torch.nn.functional.pad(kernel1x1, [1, 1, 1, 1])149 150 def _fuse_bn_tensor(self, branch):151 if branch is None:152 return 0, 0153 if isinstance(branch, Conv):154 kernel = branch.conv.weight155 running_mean = branch.bn.running_mean156 running_var = branch.bn.running_var157 gamma = branch.bn.weight158 beta = branch.bn.bias159 eps = branch.bn.eps160 elif isinstance(branch, nn.BatchNorm2d):161 if not hasattr(self, 'id_tensor'):162 input_dim = self.c1 // self.g163 kernel_value = np.zeros((self.c1, input_dim, 3, 3), dtype=np.float32)164 for i in range(self.c1):165 kernel_value[i, i % input_dim, 1, 1] = 1166 self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)167 kernel = self.id_tensor168 running_mean = branch.running_mean169 running_var = branch.running_var170 gamma = branch.weight171 beta = branch.bias172 eps = branch.eps173 std = (running_var + eps).sqrt()174 t = (gamma / std).reshape(-1, 1, 1, 1)175 return kernel * t, beta - running_mean * gamma / std176 177 def fuse_convs(self):178 if hasattr(self, 'conv'):179 return180 kernel, bias = self.get_equivalent_kernel_bias()181 self.conv = nn.Conv2d(in_channels=self.conv1.conv.in_channels,182 out_channels=self.conv1.conv.out_channels,183 kernel_size=self.conv1.conv.kernel_size,184 stride=self.conv1.conv.stride,185 padding=self.conv1.conv.padding,186 dilation=self.conv1.conv.dilation,187 groups=self.conv1.conv.groups,188 bias=True).requires_grad_(False)189 self.conv.weight.data = kernel190 self.conv.bias.data = bias191 for para in self.parameters():192 para.detach_()193 self.__delattr__('conv1')194 self.__delattr__('conv2')195 if hasattr(self, 'nm'):196 self.__delattr__('nm')197 if hasattr(self, 'bn'):198 self.__delattr__('bn')199 if hasattr(self, 'id_tensor'):200 self.__delattr__('id_tensor')201 202 203class SP(nn.Module):204 def __init__(self, k=3, s=1):205 super(SP, self).__init__()206 self.m = nn.MaxPool2d(kernel_size=k, stride=s, padding=k // 2)207 208 def forward(self, x):209 return self.m(x)210 211 212class MP(nn.Module):213 # Max pooling214 def __init__(self, k=2):215 super(MP, self).__init__()216 self.m = nn.MaxPool2d(kernel_size=k, stride=k)217 218 def forward(self, x):219 return self.m(x)220 221 222class ConvTranspose(nn.Module):223 # Convolution transpose 2d layer224 default_act = nn.SiLU() # default activation225 226 def __init__(self, c1, c2, k=2, s=2, p=0, bn=True, act=True):227 super().__init__()228 self.conv_transpose = nn.ConvTranspose2d(c1, c2, k, s, p, bias=not bn)229 self.bn = nn.BatchNorm2d(c2) if bn else nn.Identity()230 self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()231 232 def forward(self, x):233 return self.act(self.bn(self.conv_transpose(x)))234 235 236class DWConv(Conv):237 # Depth-wise convolution238 def __init__(self, c1, c2, k=1, s=1, d=1, act=True): # ch_in, ch_out, kernel, stride, dilation, activation239 super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act)240 241 242class DWConvTranspose2d(nn.ConvTranspose2d):243 # Depth-wise transpose convolution244 def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0): # ch_in, ch_out, kernel, stride, padding, padding_out245 super().__init__(c1, c2, k, s, p1, p2, groups=math.gcd(c1, c2))246 247class DWConvTranspose(nn.Module):248 # Convolution transpose 2d layer249 default_act = nn.SiLU() # default activation250 251 def __init__(self, c1, c2, k=2, s=2, p=0, bn=True, act=True):252 super().__init__()253 self.dwconv_transpose = DWConvTranspose2d(c1, c2, k, s, p)254 self.bn = nn.BatchNorm2d(c2) if bn else nn.Identity()255 self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()256 257 def forward(self, x):258 return self.act(self.bn(self.dwconv_transpose(x)))259 260 261class DFL(nn.Module):262 # DFL module263 def __init__(self, c1=17):264 super().__init__()265 self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)266 self.conv.weight.data[:] = nn.Parameter(torch.arange(c1, dtype=torch.float).view(1, c1, 1, 1)) # / 120.0267 self.c1 = c1268 # self.bn = nn.BatchNorm2d(4)269 270 def forward(self, x):271 b, c, a = x.shape # batch, channels, anchors272 return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)273 # return self.conv(x.view(b, self.c1, 4, a).softmax(1)).view(b, 4, a)274 275 276class BottleneckBase(nn.Module):277 # Standard bottleneck278 def __init__(self, c1, c2, shortcut=True, g=1, k=(1, 3), e=0.5): # ch_in, ch_out, shortcut, kernels, groups, expand279 super().__init__()280 c_ = int(c2 * e) # hidden channels281 self.cv1 = Conv(c1, c_, k[0], 1)282 self.cv2 = Conv(c_, c2, k[1], 1, g=g)283 self.add = shortcut and c1 == c2284 285 def forward(self, x):286 return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))287 288 289class RBottleneckBase(nn.Module):290 # Standard bottleneck291 def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 1), e=0.5): # ch_in, ch_out, shortcut, kernels, groups, expand292 super().__init__()293 c_ = int(c2 * e) # hidden channels294 self.cv1 = Conv(c1, c_, k[0], 1)295 self.cv2 = Conv(c_, c2, k[1], 1, g=g)296 self.add = shortcut and c1 == c2297 298 def forward(self, x):299 return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))300 301 302class RepNRBottleneckBase(nn.Module):303 # Standard bottleneck304 def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 1), e=0.5): # ch_in, ch_out, shortcut, kernels, groups, expand305 super().__init__()306 c_ = int(c2 * e) # hidden channels307 self.cv1 = RepConvN(c1, c_, k[0], 1)308 self.cv2 = Conv(c_, c2, k[1], 1, g=g)309 self.add = shortcut and c1 == c2310 311 def forward(self, x):312 return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))313 314 315class Bottleneck(nn.Module):316 # Standard bottleneck317 def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5): # ch_in, ch_out, shortcut, kernels, groups, expand318 super().__init__()319 c_ = int(c2 * e) # hidden channels320 self.cv1 = Conv(c1, c_, k[0], 1)321 self.cv2 = Conv(c_, c2, k[1], 1, g=g)322 self.add = shortcut and c1 == c2323 324 def forward(self, x):325 return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))326 327 328class RepNBottleneck(nn.Module):329 # Standard bottleneck330 def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5): # ch_in, ch_out, shortcut, kernels, groups, expand331 super().__init__()332 c_ = int(c2 * e) # hidden channels333 self.cv1 = RepConvN(c1, c_, k[0], 1)334 self.cv2 = Conv(c_, c2, k[1], 1, g=g)335 self.add = shortcut and c1 == c2336 337 def forward(self, x):338 return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))339 340 341class Res(nn.Module):342 # ResNet bottleneck343 def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion344 super(Res, self).__init__()345 c_ = int(c2 * e) # hidden channels346 self.cv1 = Conv(c1, c_, 1, 1)347 self.cv2 = Conv(c_, c_, 3, 1, g=g)348 self.cv3 = Conv(c_, c2, 1, 1)349 self.add = shortcut and c1 == c2350 351 def forward(self, x):352 return x + self.cv3(self.cv2(self.cv1(x))) if self.add else self.cv3(self.cv2(self.cv1(x)))353 354 355class RepNRes(nn.Module):356 # ResNet bottleneck357 def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion358 super(RepNRes, self).__init__()359 c_ = int(c2 * e) # hidden channels360 self.cv1 = Conv(c1, c_, 1, 1)361 self.cv2 = RepConvN(c_, c_, 3, 1, g=g)362 self.cv3 = Conv(c_, c2, 1, 1)363 self.add = shortcut and c1 == c2364 365 def forward(self, x):366 return x + self.cv3(self.cv2(self.cv1(x))) if self.add else self.cv3(self.cv2(self.cv1(x)))367 368 369class BottleneckCSP(nn.Module):370 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks371 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion372 super().__init__()373 c_ = int(c2 * e) # hidden channels374 self.cv1 = Conv(c1, c_, 1, 1)375 self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)376 self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)377 self.cv4 = Conv(2 * c_, c2, 1, 1)378 self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)379 self.act = nn.SiLU()380 self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))381 382 def forward(self, x):383 y1 = self.cv3(self.m(self.cv1(x)))384 y2 = self.cv2(x)385 return self.cv4(self.act(self.bn(torch.cat((y1, y2), 1))))386 387 388class CSP(nn.Module):389 # CSP Bottleneck with 3 convolutions390 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion391 super().__init__()392 c_ = int(c2 * e) # hidden channels393 self.cv1 = Conv(c1, c_, 1, 1)394 self.cv2 = Conv(c1, c_, 1, 1)395 self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2)396 self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))397 398 def forward(self, x):399 return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))400 401 402class RepNCSP(nn.Module):403 # CSP Bottleneck with 3 convolutions404 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion405 super().__init__()406 c_ = int(c2 * e) # hidden channels407 self.cv1 = Conv(c1, c_, 1, 1)408 self.cv2 = Conv(c1, c_, 1, 1)409 self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2)410 self.m = nn.Sequential(*(RepNBottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))411 412 def forward(self, x):413 return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))414 415 416class CSPBase(nn.Module):417 # CSP Bottleneck with 3 convolutions418 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion419 super().__init__()420 c_ = int(c2 * e) # hidden channels421 self.cv1 = Conv(c1, c_, 1, 1)422 self.cv2 = Conv(c1, c_, 1, 1)423 self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2)424 self.m = nn.Sequential(*(BottleneckBase(c_, c_, shortcut, g, e=1.0) for _ in range(n)))425 426 def forward(self, x):427 return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))428 429 430class SPP(nn.Module):431 # Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729432 def __init__(self, c1, c2, k=(5, 9, 13)):433 super().__init__()434 c_ = c1 // 2 # hidden channels435 self.cv1 = Conv(c1, c_, 1, 1)436 self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)437 self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])438 439 def forward(self, x):440 x = self.cv1(x)441 with warnings.catch_warnings():442 warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning443 return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))444 445 446class ASPP(torch.nn.Module):447 448 def __init__(self, in_channels, out_channels):449 super().__init__()450 kernel_sizes = [1, 3, 3, 1]451 dilations = [1, 3, 6, 1]452 paddings = [0, 3, 6, 0]453 self.aspp = torch.nn.ModuleList()454 for aspp_idx in range(len(kernel_sizes)):455 conv = torch.nn.Conv2d(456 in_channels,457 out_channels,458 kernel_size=kernel_sizes[aspp_idx],459 stride=1,460 dilation=dilations[aspp_idx],461 padding=paddings[aspp_idx],462 bias=True)463 self.aspp.append(conv)464 self.gap = torch.nn.AdaptiveAvgPool2d(1)465 self.aspp_num = len(kernel_sizes)466 for m in self.modules():467 if isinstance(m, torch.nn.Conv2d):468 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels469 m.weight.data.normal_(0, math.sqrt(2. / n))470 m.bias.data.fill_(0)471 472 def forward(self, x):473 avg_x = self.gap(x)474 out = []475 for aspp_idx in range(self.aspp_num):476 inp = avg_x if (aspp_idx == self.aspp_num - 1) else x477 out.append(F.relu_(self.aspp[aspp_idx](inp)))478 out[-1] = out[-1].expand_as(out[-2])479 out = torch.cat(out, dim=1)480 return out481 482 483class SPPCSPC(nn.Module):484 # CSP SPP https://github.com/WongKinYiu/CrossStagePartialNetworks485 def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, k=(5, 9, 13)):486 super(SPPCSPC, self).__init__()487 c_ = int(2 * c2 * e) # hidden channels488 self.cv1 = Conv(c1, c_, 1, 1)489 self.cv2 = Conv(c1, c_, 1, 1)490 self.cv3 = Conv(c_, c_, 3, 1)491 self.cv4 = Conv(c_, c_, 1, 1)492 self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])493 self.cv5 = Conv(4 * c_, c_, 1, 1)494 self.cv6 = Conv(c_, c_, 3, 1)495 self.cv7 = Conv(2 * c_, c2, 1, 1)496 497 def forward(self, x):498 x1 = self.cv4(self.cv3(self.cv1(x)))499 y1 = self.cv6(self.cv5(torch.cat([x1] + [m(x1) for m in self.m], 1)))500 y2 = self.cv2(x)501 return self.cv7(torch.cat((y1, y2), dim=1))502 503 504class SPPF(nn.Module):505 # Spatial Pyramid Pooling - Fast (SPPF) layer by Glenn Jocher506 def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))507 super().__init__()508 c_ = c1 // 2 # hidden channels509 self.cv1 = Conv(c1, c_, 1, 1)510 self.cv2 = Conv(c_ * 4, c2, 1, 1)511 self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)512 # self.m = SoftPool2d(kernel_size=k, stride=1, padding=k // 2)513 514 def forward(self, x):515 x = self.cv1(x)516 with warnings.catch_warnings():517 warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning518 y1 = self.m(x)519 y2 = self.m(y1)520 return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))521 522 523import torch.nn.functional as F524from torch.nn.modules.utils import _pair525 526 527class ReOrg(nn.Module):528 # yolo529 def __init__(self):530 super(ReOrg, self).__init__()531 532 def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)533 return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)534 535 536class Contract(nn.Module):537 # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)538 def __init__(self, gain=2):539 super().__init__()540 self.gain = gain541 542 def forward(self, x):543 b, c, h, w = x.size() # assert (h / s == 0) and (W / s == 0), 'Indivisible gain'544 s = self.gain545 x = x.view(b, c, h // s, s, w // s, s) # x(1,64,40,2,40,2)546 x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)547 return x.view(b, c * s * s, h // s, w // s) # x(1,256,40,40)548 549 550class Expand(nn.Module):551 # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)552 def __init__(self, gain=2):553 super().__init__()554 self.gain = gain555 556 def forward(self, x):557 b, c, h, w = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'558 s = self.gain559 x = x.view(b, s, s, c // s ** 2, h, w) # x(1,2,2,16,80,80)560 x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)561 return x.view(b, c // s ** 2, h * s, w * s) # x(1,16,160,160)562 563 564class Concat(nn.Module):565 # Concatenate a list of tensors along dimension566 def __init__(self, dimension=1):567 super().__init__()568 self.d = dimension569 570 def forward(self, x):571 return torch.cat(x, self.d)572 573 574class Shortcut(nn.Module):575 def __init__(self, dimension=0):576 super(Shortcut, self).__init__()577 self.d = dimension578 579 def forward(self, x):580 return x[0]+x[1]581 582 583class Silence(nn.Module):584 def __init__(self):585 super(Silence, self).__init__()586 def forward(self, x): 587 return x588 589 590##### GELAN ##### 591 592class SPPELAN(nn.Module):593 # spp-elan594 def __init__(self, c1, c2, c3): # ch_in, ch_out, number, shortcut, groups, expansion595 super().__init__()596 self.c = c3597 self.cv1 = Conv(c1, c3, 1, 1)598 self.cv2 = SP(5)599 self.cv3 = SP(5)600 self.cv4 = SP(5)601 self.cv5 = Conv(4*c3, c2, 1, 1)602 603 def forward(self, x):604 y = [self.cv1(x)]605 y.extend(m(y[-1]) for m in [self.cv2, self.cv3, self.cv4])606 return self.cv5(torch.cat(y, 1))607 608 609class RepNCSPELAN4(nn.Module):610 # csp-elan611 def __init__(self, c1, c2, c3, c4, c5=1): # ch_in, ch_out, number, shortcut, groups, expansion612 super().__init__()613 self.c = c3//2614 self.cv1 = Conv(c1, c3, 1, 1)615 self.cv2 = nn.Sequential(RepNCSP(c3//2, c4, c5), Conv(c4, c4, 3, 1))616 self.cv3 = nn.Sequential(RepNCSP(c4, c4, c5), Conv(c4, c4, 3, 1))617 self.cv4 = Conv(c3+(2*c4), c2, 1, 1)618 619 def forward(self, x):620 y = list(self.cv1(x).chunk(2, 1))621 y.extend((m(y[-1])) for m in [self.cv2, self.cv3])622 return self.cv4(torch.cat(y, 1))623 624 def forward_split(self, x):625 y = list(self.cv1(x).split((self.c, self.c), 1))626 y.extend(m(y[-1]) for m in [self.cv2, self.cv3])627 return self.cv4(torch.cat(y, 1))628 629#################630 631 632##### YOLOR #####633 634class ImplicitA(nn.Module):635 def __init__(self, channel):636 super(ImplicitA, self).__init__()637 self.channel = channel638 self.implicit = nn.Parameter(torch.zeros(1, channel, 1, 1))639 nn.init.normal_(self.implicit, std=.02) 640 641 def forward(self, x):642 return self.implicit + x643 644 645class ImplicitM(nn.Module):646 def __init__(self, channel):647 super(ImplicitM, self).__init__()648 self.channel = channel649 self.implicit = nn.Parameter(torch.ones(1, channel, 1, 1))650 nn.init.normal_(self.implicit, mean=1., std=.02) 651 652 def forward(self, x):653 return self.implicit * x654 655#################656 657 658##### CBNet #####659 660class CBLinear(nn.Module):661 def __init__(self, c1, c2s, k=1, s=1, p=None, g=1): # ch_in, ch_outs, kernel, stride, padding, groups662 super(CBLinear, self).__init__()663 self.c2s = c2s664 self.conv = nn.Conv2d(c1, sum(c2s), k, s, autopad(k, p), groups=g, bias=True)665 666 def forward(self, x):667 outs = self.conv(x).split(self.c2s, dim=1)668 return outs669 670class CBFuse(nn.Module):671 def __init__(self, idx):672 super(CBFuse, self).__init__()673 self.idx = idx674 675 def forward(self, xs):676 target_size = xs[-1].shape[2:]677 res = [F.interpolate(x[self.idx[i]], size=target_size, mode='nearest') for i, x in enumerate(xs[:-1])]678 out = torch.sum(torch.stack(res + xs[-1:]), dim=0)679 return out680 681#################682 683 684class DetectMultiBackend(nn.Module):685 # YOLO MultiBackend class for python inference on various backends686 def __init__(self, weights='yolo.pt', device=torch.device('cpu'), dnn=False, data=None, fp16=False, fuse=True):687 # Usage:688 # PyTorch: weights = *.pt689 # TorchScript: *.torchscript690 # ONNX Runtime: *.onnx691 # ONNX OpenCV DNN: *.onnx --dnn692 # OpenVINO: *_openvino_model693 # CoreML: *.mlmodel694 # TensorRT: *.engine695 # TensorFlow SavedModel: *_saved_model696 # TensorFlow GraphDef: *.pb697 # TensorFlow Lite: *.tflite698 # TensorFlow Edge TPU: *_edgetpu.tflite699 # PaddlePaddle: *_paddle_model700 from models.experimental import attempt_download, attempt_load # scoped to avoid circular import701 702 super().__init__()703 w = str(weights[0] if isinstance(weights, list) else weights)704 pt, jit, onnx, onnx_end2end, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, triton = self._model_type(w)705 fp16 &= pt or jit or onnx or engine # FP16706 nhwc = coreml or saved_model or pb or tflite or edgetpu # BHWC formats (vs torch BCWH)707 stride = 32 # default stride708 cuda = torch.cuda.is_available() and device.type != 'cpu' # use CUDA709 if not (pt or triton):710 w = attempt_download(w) # download if not local711 712 if pt: # PyTorch713 model = attempt_load(weights if isinstance(weights, list) else w, device=device, inplace=True, fuse=fuse)714 stride = max(int(model.stride.max()), 32) # model stride715 names = model.module.names if hasattr(model, 'module') else model.names # get class names716 model.half() if fp16 else model.float()717 self.model = model # explicitly assign for to(), cpu(), cuda(), half()718 elif jit: # TorchScript719 LOGGER.info(f'Loading {w} for TorchScript inference...')720 extra_files = {'config.txt': ''} # model metadata721 model = torch.jit.load(w, _extra_files=extra_files, map_location=device)722 model.half() if fp16 else model.float()723 if extra_files['config.txt']: # load metadata dict724 d = json.loads(extra_files['config.txt'],725 object_hook=lambda d: {int(k) if k.isdigit() else k: v726 for k, v in d.items()})727 stride, names = int(d['stride']), d['names']728 elif dnn: # ONNX OpenCV DNN729 LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')730 check_requirements('opencv-python>=4.5.4')731 net = cv2.dnn.readNetFromONNX(w)732 elif onnx: # ONNX Runtime733 LOGGER.info(f'Loading {w} for ONNX Runtime inference...')734 check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))735 import onnxruntime736 providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']737 session = onnxruntime.InferenceSession(w, providers=providers)738 output_names = [x.name for x in session.get_outputs()]739 meta = session.get_modelmeta().custom_metadata_map # metadata740 if 'stride' in meta:741 stride, names = int(meta['stride']), eval(meta['names'])742 elif xml: # OpenVINO743 LOGGER.info(f'Loading {w} for OpenVINO inference...')744 check_requirements('openvino') # requires openvino-dev: https://pypi.org/project/openvino-dev/745 from openvino.runtime import Core, Layout, get_batch746 ie = Core()747 if not Path(w).is_file(): # if not *.xml748 w = next(Path(w).glob('*.xml')) # get *.xml file from *_openvino_model dir749 network = ie.read_model(model=w, weights=Path(w).with_suffix('.bin'))750 if network.get_parameters()[0].get_layout().empty:751 network.get_parameters()[0].set_layout(Layout("NCHW"))752 batch_dim = get_batch(network)753 if batch_dim.is_static:754 batch_size = batch_dim.get_length()755 executable_network = ie.compile_model(network, device_name="CPU") # device_name="MYRIAD" for Intel NCS2756 stride, names = self._load_metadata(Path(w).with_suffix('.yaml')) # load metadata757 elif engine: # TensorRT758 LOGGER.info(f'Loading {w} for TensorRT inference...')759 import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download760 check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=7.0.0761 if device.type == 'cpu':762 device = torch.device('cuda:0')763 Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))764 logger = trt.Logger(trt.Logger.INFO)765 with open(w, 'rb') as f, trt.Runtime(logger) as runtime:766 model = runtime.deserialize_cuda_engine(f.read())767 context = model.create_execution_context()768 bindings = OrderedDict()769 output_names = []770 fp16 = False # default updated below771 dynamic = False772 for i in range(model.num_bindings):773 name = model.get_binding_name(i)774 dtype = trt.nptype(model.get_binding_dtype(i))775 if model.binding_is_input(i):776 if -1 in tuple(model.get_binding_shape(i)): # dynamic777 dynamic = True778 context.set_binding_shape(i, tuple(model.get_profile_shape(0, i)[2]))779 if dtype == np.float16:780 fp16 = True781 else: # output782 output_names.append(name)783 shape = tuple(context.get_binding_shape(i))784 im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)785 bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))786 binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())787 batch_size = bindings['images'].shape[0] # if dynamic, this is instead max batch size788 elif coreml: # CoreML789 LOGGER.info(f'Loading {w} for CoreML inference...')790 import coremltools as ct791 model = ct.models.MLModel(w)792 elif saved_model: # TF SavedModel793 LOGGER.info(f'Loading {w} for TensorFlow SavedModel inference...')794 import tensorflow as tf795 keras = False # assume TF1 saved_model796 model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)797 elif pb: # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt798 LOGGER.info(f'Loading {w} for TensorFlow GraphDef inference...')799 import tensorflow as tf800 801 def wrap_frozen_graph(gd, inputs, outputs):802 x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) # wrapped803 ge = x.graph.as_graph_element804 return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))805 806 def gd_outputs(gd):807 name_list, input_list = [], []808 for node in gd.node: # tensorflow.core.framework.node_def_pb2.NodeDef809 name_list.append(node.name)810 input_list.extend(node.input)811 return sorted(f'{x}:0' for x in list(set(name_list) - set(input_list)) if not x.startswith('NoOp'))812 813 gd = tf.Graph().as_graph_def() # TF GraphDef814 with open(w, 'rb') as f:815 gd.ParseFromString(f.read())816 frozen_func = wrap_frozen_graph(gd, inputs="x:0", outputs=gd_outputs(gd))817 elif tflite or edgetpu: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python818 try: # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu819 from tflite_runtime.interpreter import Interpreter, load_delegate820 except ImportError:821 import tensorflow as tf822 Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate,823 if edgetpu: # TF Edge TPU https://coral.ai/software/#edgetpu-runtime824 LOGGER.info(f'Loading {w} for TensorFlow Lite Edge TPU inference...')825 delegate = {826 'Linux': 'libedgetpu.so.1',827 'Darwin': 'libedgetpu.1.dylib',828 'Windows': 'edgetpu.dll'}[platform.system()]829 interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)])830 else: # TFLite831 LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')832 interpreter = Interpreter(model_path=w) # load TFLite model833 interpreter.allocate_tensors() # allocate834 input_details = interpreter.get_input_details() # inputs835 output_details = interpreter.get_output_details() # outputs836 # load metadata837 with contextlib.suppress(zipfile.BadZipFile):838 with zipfile.ZipFile(w, "r") as model:839 meta_file = model.namelist()[0]840 meta = ast.literal_eval(model.read(meta_file).decode("utf-8"))841 stride, names = int(meta['stride']), meta['names']842 elif tfjs: # TF.js843 raise NotImplementedError('ERROR: YOLO TF.js inference is not supported')844 elif paddle: # PaddlePaddle845 LOGGER.info(f'Loading {w} for PaddlePaddle inference...')846 check_requirements('paddlepaddle-gpu' if cuda else 'paddlepaddle')847 import paddle.inference as pdi848 if not Path(w).is_file(): # if not *.pdmodel849 w = next(Path(w).rglob('*.pdmodel')) # get *.pdmodel file from *_paddle_model dir850 weights = Path(w).with_suffix('.pdiparams')851 config = pdi.Config(str(w), str(weights))852 if cuda:853 config.enable_use_gpu(memory_pool_init_size_mb=2048, device_id=0)854 predictor = pdi.create_predictor(config)855 input_handle = predictor.get_input_handle(predictor.get_input_names()[0])856 output_names = predictor.get_output_names()857 elif triton: # NVIDIA Triton Inference Server858 LOGGER.info(f'Using {w} as Triton Inference Server...')859 check_requirements('tritonclient[all]')860 from utils.triton import TritonRemoteModel861 model = TritonRemoteModel(url=w)862 nhwc = model.runtime.startswith("tensorflow")863 else:864 raise NotImplementedError(f'ERROR: {w} is not a supported format')865 866 # class names867 if 'names' not in locals():868 names = yaml_load(data)['names'] if data else {i: f'class{i}' for i in range(999)}869 if names[0] == 'n01440764' and len(names) == 1000: # ImageNet870 names = yaml_load(ROOT / 'data/ImageNet.yaml')['names'] # human-readable names871 872 self.__dict__.update(locals()) # assign all variables to self873 874 def forward(self, im, augment=False, visualize=False):875 # YOLO MultiBackend inference876 b, ch, h, w = im.shape # batch, channel, height, width877 if self.fp16 and im.dtype != torch.float16:878 im = im.half() # to FP16879 if self.nhwc:880 im = im.permute(0, 2, 3, 1) # torch BCHW to numpy BHWC shape(1,320,192,3)881 882 if self.pt: # PyTorch883 y = self.model(im, augment=augment, visualize=visualize) if augment or visualize else self.model(im)884 elif self.jit: # TorchScript885 y = self.model(im)886 elif self.dnn: # ONNX OpenCV DNN887 im = im.cpu().numpy() # torch to numpy888 self.net.setInput(im)889 y = self.net.forward()890 elif self.onnx: # ONNX Runtime891 im = im.cpu().numpy() # torch to numpy892 y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})893 elif self.xml: # OpenVINO894 im = im.cpu().numpy() # FP32895 y = list(self.executable_network([im]).values())896 elif self.engine: # TensorRT897 if self.dynamic and im.shape != self.bindings['images'].shape:898 i = self.model.get_binding_index('images')899 self.context.set_binding_shape(i, im.shape) # reshape if dynamic900 self.bindings['images'] = self.bindings['images']._replace(shape=im.shape)901 for name in self.output_names:902 i = self.model.get_binding_index(name)903 self.bindings[name].data.resize_(tuple(self.context.get_binding_shape(i)))904 s = self.bindings['images'].shape905 assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"906 self.binding_addrs['images'] = int(im.data_ptr())907 self.context.execute_v2(list(self.binding_addrs.values()))908 y = [self.bindings[x].data for x in sorted(self.output_names)]909 elif self.coreml: # CoreML910 im = im.cpu().numpy()911 im = Image.fromarray((im[0] * 255).astype('uint8'))912 # im = im.resize((192, 320), Image.ANTIALIAS)913 y = self.model.predict({'image': im}) # coordinates are xywh normalized914 if 'confidence' in y:915 box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels916 conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)917 y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)918 else:919 y = list(reversed(y.values())) # reversed for segmentation models (pred, proto)920 elif self.paddle: # PaddlePaddle921 im = im.cpu().numpy().astype(np.float32)922 self.input_handle.copy_from_cpu(im)923 self.predictor.run()924 y = [self.predictor.get_output_handle(x).copy_to_cpu() for x in self.output_names]925 elif self.triton: # NVIDIA Triton Inference Server926 y = self.model(im)927 else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)928 im = im.cpu().numpy()929 if self.saved_model: # SavedModel930 y = self.model(im, training=False) if self.keras else self.model(im)931 elif self.pb: # GraphDef932 y = self.frozen_func(x=self.tf.constant(im))933 else: # Lite or Edge TPU934 input = self.input_details[0]935 int8 = input['dtype'] == np.uint8 # is TFLite quantized uint8 model936 if int8:937 scale, zero_point = input['quantization']938 im = (im / scale + zero_point).astype(np.uint8) # de-scale939 self.interpreter.set_tensor(input['index'], im)940 self.interpreter.invoke()941 y = []942 for output in self.output_details:943 x = self.interpreter.get_tensor(output['index'])944 if int8:945 scale, zero_point = output['quantization']946 x = (x.astype(np.float32) - zero_point) * scale # re-scale947 y.append(x)948 y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]949 y[0][..., :4] *= [w, h, w, h] # xywh normalized to pixels950 951 if isinstance(y, (list, tuple)):952 return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y]953 else:954 return self.from_numpy(y)955 956 def from_numpy(self, x):957 return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x958 959 def warmup(self, imgsz=(1, 3, 640, 640)):960 # Warmup model by running inference once961 warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton962 if any(warmup_types) and (self.device.type != 'cpu' or self.triton):963 im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input964 for _ in range(2 if self.jit else 1): #965 self.forward(im) # warmup966 967 @staticmethod968 def _model_type(p='path/to/model.pt'):969 # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx970 # types = [pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle]971 from export import export_formats972 from utils.downloads import is_url973 sf = list(export_formats().Suffix) # export suffixes974 if not is_url(p, check=False):975 check_suffix(p, sf) # checks976 url = urlparse(p) # if url may be Triton inference server977 types = [s in Path(p).name for s in sf]978 types[8] &= not types[9] # tflite &= not edgetpu979 triton = not any(types) and all([any(s in url.scheme for s in ["http", "grpc"]), url.netloc])980 return types + [triton]981 982 @staticmethod983 def _load_metadata(f=Path('path/to/meta.yaml')):984 # Load metadata from meta.yaml if it exists985 if f.exists():986 d = yaml_load(f)987 return d['stride'], d['names'] # assign stride, names988 return None, None989 990 991class AutoShape(nn.Module):992 # YOLO input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS993 conf = 0.25 # NMS confidence threshold994 iou = 0.45 # NMS IoU threshold995 agnostic = False # NMS class-agnostic996 multi_label = False # NMS multiple labels per box997 classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs998 max_det = 1000 # maximum number of detections per image999 amp = False # Automatic Mixed Precision (AMP) inference1000 1001 def __init__(self, model, verbose=True):1002 super().__init__()1003 if verbose:1004 LOGGER.info('Adding AutoShape... ')1005 copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=()) # copy attributes1006 self.dmb = isinstance(model, DetectMultiBackend) # DetectMultiBackend() instance1007 self.pt = not self.dmb or model.pt # PyTorch model1008 self.model = model.eval()1009 if self.pt:1010 m = self.model.model.model[-1] if self.dmb else self.model.model[-1] # Detect()1011 m.inplace = False # Detect.inplace=False for safe multithread inference1012 m.export = True # do not output loss values1013 1014 def _apply(self, fn):1015 # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers1016 self = super()._apply(fn)1017 from models.yolo import Detect, Segment1018 if self.pt:1019 m = self.model.model.model[-1] if self.dmb else self.model.model[-1] # Detect()1020 if isinstance(m, (Detect, Segment)):1021 for k in 'stride', 'anchor_grid', 'stride_grid', 'grid':1022 x = getattr(m, k)1023 setattr(m, k, list(map(fn, x))) if isinstance(x, (list, tuple)) else setattr(m, k, fn(x))1024 return self1025 1026 @smart_inference_mode()1027 def forward(self, ims, size=640, augment=False, profile=False):1028 # Inference from various sources. For size(height=640, width=1280), RGB images example inputs are:1029 # file: ims = 'data/images/zidane.jpg' # str or PosixPath1030 # URI: = 'https://ultralytics.com/images/zidane.jpg'1031 # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)1032 # PIL: = Image.open('image.jpg') or ImageGrab.grab() # HWC x(640,1280,3)1033 # numpy: = np.zeros((640,1280,3)) # HWC1034 # torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)1035 # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images1036 1037 dt = (Profile(), Profile(), Profile())1038 with dt[0]:1039 if isinstance(size, int): # expand1040 size = (size, size)1041 p = next(self.model.parameters()) if self.pt else torch.empty(1, device=self.model.device) # param1042 autocast = self.amp and (p.device.type != 'cpu') # Automatic Mixed Precision (AMP) inference1043 if isinstance(ims, torch.Tensor): # torch1044 with amp.autocast(autocast):1045 return self.model(ims.to(p.device).type_as(p), augment=augment) # inference1046 1047 # Pre-process1048 n, ims = (len(ims), list(ims)) if isinstance(ims, (list, tuple)) else (1, [ims]) # number, list of images1049 shape0, shape1, files = [], [], [] # image and inference shapes, filenames1050 for i, im in enumerate(ims):1051 f = f'image{i}' # filename1052 if isinstance(im, (str, Path)): # filename or uri1053 im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im1054 im = np.asarray(exif_transpose(im))1055 elif isinstance(im, Image.Image): # PIL Image1056 im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f1057 files.append(Path(f).with_suffix('.jpg').name)1058 if im.shape[0] < 5: # image in CHW1059 im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)1060 im = im[..., :3] if im.ndim == 3 else cv2.cvtColor(im, cv2.COLOR_GRAY2BGR) # enforce 3ch input1061 s = im.shape[:2] # HWC1062 shape0.append(s) # image shape1063 g = max(size) / max(s) # gain1064 shape1.append([int(y * g) for y in s])1065 ims[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update1066 shape1 = [make_divisible(x, self.stride) for x in np.array(shape1).max(0)] # inf shape1067 x = [letterbox(im, shape1, auto=False)[0] for im in ims] # pad1068 x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2))) # stack and BHWC to BCHW1069 x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/321070 1071 with amp.autocast(autocast):1072 # Inference1073 with dt[1]:1074 y = self.model(x, augment=augment) # forward1075 1076 # Post-process1077 with dt[2]:1078 y = non_max_suppression(y if self.dmb else y[0],1079 self.conf,1080 self.iou,1081 self.classes,1082 self.agnostic,1083 self.multi_label,1084 max_det=self.max_det) # NMS1085 for i in range(n):1086 scale_boxes(shape1, y[i][:, :4], shape0[i])1087 1088 return Detections(ims, y, files, dt, self.names, x.shape)1089 1090 1091class Detections:1092 # YOLO detections class for inference results1093 def __init__(self, ims, pred, files, times=(0, 0, 0), names=None, shape=None):1094 super().__init__()1095 d = pred[0].device # device1096 gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in ims] # normalizations1097 self.ims = ims # list of images as numpy arrays1098 self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)1099 self.names = names # class names1100 self.files = files # image filenames1101 self.times = times # profiling times1102 self.xyxy = pred # xyxy pixels1103 self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels1104 self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized1105 self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized1106 self.n = len(self.pred) # number of images (batch size)1107 self.t = tuple(x.t / self.n * 1E3 for x in times) # timestamps (ms)1108 self.s = tuple(shape) # inference BCHW shape1109 1110 def _run(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path('')):1111 s, crops = '', []1112 for i, (im, pred) in enumerate(zip(self.ims, self.pred)):1113 s += f'\nimage {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string1114 if pred.shape[0]:1115 for c in pred[:, -1].unique():1116 n = (pred[:, -1] == c).sum() # detections per class1117 s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string1118 s = s.rstrip(', ')1119 if show or save or render or crop:1120 annotator = Annotator(im, example=str(self.names))1121 for *box, conf, cls in reversed(pred): # xyxy, confidence, class1122 label = f'{self.names[int(cls)]} {conf:.2f}'1123 if crop:1124 file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None1125 crops.append({1126 'box': box,1127 'conf': conf,1128 'cls': cls,1129 'label': label,1130 'im': save_one_box(box, im, file=file, save=save)})1131 else: # all others1132 annotator.box_label(box, label if labels else '', color=colors(cls))1133 im = annotator.im1134 else:1135 s += '(no detections)'1136 1137 im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np1138 if show:1139 display(im) if is_notebook() else im.show(self.files[i])1140 if save:1141 f = self.files[i]1142 im.save(save_dir / f) # save1143 if i == self.n - 1:1144 LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")1145 if render:1146 self.ims[i] = np.asarray(im)1147 if pprint:1148 s = s.lstrip('\n')1149 return f'{s}\nSpeed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {self.s}' % self.t1150 if crop:1151 if save:1152 LOGGER.info(f'Saved results to {save_dir}\n')1153 return crops1154 1155 @TryExcept('Showing images is not supported in this environment')1156 def show(self, labels=True):1157 self._run(show=True, labels=labels) # show results1158 1159 def save(self, labels=True, save_dir='runs/detect/exp', exist_ok=False):1160 save_dir = increment_path(save_dir, exist_ok, mkdir=True) # increment save_dir1161 self._run(save=True, labels=labels, save_dir=save_dir) # save results1162 1163 def crop(self, save=True, save_dir='runs/detect/exp', exist_ok=False):1164 save_dir = increment_path(save_dir, exist_ok, mkdir=True) if save else None1165 return self._run(crop=True, save=save, save_dir=save_dir) # crop results1166 1167 def render(self, labels=True):1168 self._run(render=True, labels=labels) # render results1169 return self.ims1170 1171 def pandas(self):1172 # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])1173 new = copy(self) # return copy1174 ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns1175 cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns1176 for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):1177 a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update1178 setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])1179 return new1180 1181 def tolist(self):1182 # return a list of Detections objects, i.e. 'for result in results.tolist():'1183 r = range(self.n) # iterable1184 x = [Detections([self.ims[i]], [self.pred[i]], [self.files[i]], self.times, self.names, self.s) for i in r]1185 # for d in x:1186 # for k in ['ims', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:1187 # setattr(d, k, getattr(d, k)[0]) # pop out of list1188 return x1189 1190 def print(self):1191 LOGGER.info(self.__str__())1192 1193 def __len__(self): # override len(results)1194 return self.n1195 1196 def __str__(self): # override print(results)1197 return self._run(pprint=True) # print results1198 1199 def __repr__(self):1200 return f'YOLO {self.__class__} instance\n' + self.__str__()