|
|
|
|
|
import logging |
|
import sys |
|
from copy import deepcopy |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
from .common import * |
|
from .experimental import * |
|
import os |
|
from yolo_utils import check_anchor_order, make_divisible, copy_attr, fuse_conv_and_bn, initialize_weights, scale_img |
|
|
|
|
|
class Detect(nn.Module): |
|
stride = None |
|
export = False |
|
|
|
def __init__(self, nc=80, anchors=(), ch=()): |
|
super(Detect, self).__init__() |
|
self.nc = nc |
|
self.no = nc + 5 |
|
self.nl = len(anchors) |
|
self.na = len(anchors[0]) // 2 |
|
self.grid = [torch.zeros(1)] * self.nl |
|
a = torch.tensor(anchors).float().view(self.nl, -1, 2) |
|
self.register_buffer("anchors", a) |
|
self.register_buffer("anchor_grid", a.clone().view(self.nl, 1, -1, 1, 1, 2)) |
|
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) |
|
|
|
def forward(self, x): |
|
|
|
z = [] |
|
self.training |= self.export |
|
for i in range(self.nl): |
|
x[i] = self.m[i](x[i]) |
|
bs, _, ny, nx = x[i].shape |
|
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() |
|
|
|
if not self.training: |
|
if self.grid[i].shape[2:4] != x[i].shape[2:4]: |
|
self.grid[i] = self._make_grid(nx, ny).to(x[i].device) |
|
|
|
y = x[i].sigmoid() |
|
y[..., 0:2] = (y[..., 0:2] * 2.0 - 0.5 + self.grid[i]) * self.stride[i] |
|
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] |
|
z.append(y.view(bs, -1, self.no)) |
|
|
|
return x if self.training else (torch.cat(z, 1), x) |
|
|
|
@staticmethod |
|
def _make_grid(nx=20, ny=20): |
|
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) |
|
return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() |
|
|
|
|
|
class Model(nn.Module): |
|
def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, anchors=None): |
|
super(Model, self).__init__() |
|
if isinstance(cfg, dict): |
|
self.yaml = cfg |
|
else: |
|
import yaml |
|
|
|
self.yaml_file = Path(cfg).name |
|
with open(cfg) as f: |
|
self.yaml = yaml.load(f, Loader=yaml.SafeLoader) |
|
|
|
|
|
ch = self.yaml["ch"] = self.yaml.get("ch", ch) |
|
if nc and nc != self.yaml["nc"]: |
|
logger.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}") |
|
self.yaml["nc"] = nc |
|
if anchors: |
|
logger.info(f"Overriding model.yaml anchors with anchors={anchors}") |
|
self.yaml["anchors"] = round(anchors) |
|
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) |
|
self.names = [str(i) for i in range(self.yaml["nc"])] |
|
|
|
|
|
|
|
m = self.model[-1] |
|
if isinstance(m, Detect): |
|
s = 256 |
|
m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) |
|
m.anchors /= m.stride.view(-1, 1, 1) |
|
check_anchor_order(m) |
|
self.stride = m.stride |
|
self._initialize_biases() |
|
|
|
|
|
|
|
initialize_weights(self) |
|
self.info() |
|
logger.info("") |
|
|
|
def forward(self, x, augment=False, profile=False): |
|
if augment: |
|
img_size = x.shape[-2:] |
|
s = [1, 0.83, 0.67] |
|
f = [None, 3, None] |
|
y = [] |
|
for si, fi in zip(s, f): |
|
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max())) |
|
yi = self.forward_once(xi)[0] |
|
|
|
yi[..., :4] /= si |
|
if fi == 2: |
|
yi[..., 1] = img_size[0] - yi[..., 1] |
|
elif fi == 3: |
|
yi[..., 0] = img_size[1] - yi[..., 0] |
|
y.append(yi) |
|
return torch.cat(y, 1), None |
|
else: |
|
return self.forward_once(x, profile) |
|
|
|
def forward_once(self, x, profile=False): |
|
y, dt = [], [] |
|
|
|
for m in self.model.modules(): |
|
if isinstance(m, nn.Upsample): |
|
m.recompute_scale_factor = None |
|
|
|
for m in self.model: |
|
if m.f != -1: |
|
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] |
|
|
|
x = m(x) |
|
|
|
y.append(x if m.i in self.save else None) |
|
|
|
if profile: |
|
print("%.1fms total" % sum(dt)) |
|
return x |
|
|
|
def _initialize_biases(self, cf=None): |
|
|
|
|
|
m = self.model[-1] |
|
for mi, s in zip(m.m, m.stride): |
|
b = mi.bias.view(m.na, -1) |
|
b.data[:, 4] += math.log(8 / (640 / s) ** 2) |
|
b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) |
|
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) |
|
|
|
def _print_biases(self): |
|
m = self.model[-1] |
|
for mi in m.m: |
|
b = mi.bias.detach().view(m.na, -1).T |
|
print(("%6g Conv2d.bias:" + "%10.3g" * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) |
|
|
|
|
|
|
|
|
|
|
|
|
|
def fuse(self): |
|
|
|
for m in self.model.modules(): |
|
if type(m) is Conv and hasattr(m, "bn"): |
|
m.conv = fuse_conv_and_bn(m.conv, m.bn) |
|
delattr(m, "bn") |
|
m.forward = m.fuseforward |
|
|
|
return self |
|
|
|
def nms(self, mode=True): |
|
present = type(self.model[-1]) is NMS |
|
if mode and not present: |
|
print("Adding NMS... ") |
|
m = NMS() |
|
m.f = -1 |
|
m.i = self.model[-1].i + 1 |
|
self.model.add_module(name="%s" % m.i, module=m) |
|
self.eval() |
|
elif not mode and present: |
|
print("Removing NMS... ") |
|
self.model = self.model[:-1] |
|
return self |
|
|
|
def autoshape(self): |
|
print("Adding autoShape... ") |
|
m = autoShape(self) |
|
copy_attr(m, self, include=("yaml", "nc", "hyp", "names", "stride"), exclude=()) |
|
return m |
|
|
|
|
|
|
|
|
|
|
|
def parse_model(d, ch): |
|
logger.info("\n%3s%18s%3s%10s %-40s%-30s" % ("", "from", "n", "params", "module", "arguments")) |
|
anchors, nc, gd, gw = ( |
|
d["anchors"], |
|
d["nc"], |
|
d["depth_multiple"], |
|
d["width_multiple"], |
|
) |
|
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors |
|
no = na * (nc + 5) |
|
|
|
layers, save, c2 = [], [], ch[-1] |
|
for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]): |
|
m = eval(m) if isinstance(m, str) else m |
|
for j, a in enumerate(args): |
|
try: |
|
args[j] = eval(a) if isinstance(a, str) else a |
|
except Exception as e: |
|
logger.error(e) |
|
|
|
n = max(round(n * gd), 1) if n > 1 else n |
|
if m in [ |
|
Conv, |
|
GhostConv, |
|
Bottleneck, |
|
GhostBottleneck, |
|
SPP, |
|
DWConv, |
|
MixConv2d, |
|
Focus, |
|
CrossConv, |
|
BottleneckCSP, |
|
C3, |
|
]: |
|
c1, c2 = ch[f], args[0] |
|
if c2 != no: |
|
c2 = make_divisible(c2 * gw, 8) |
|
|
|
args = [c1, c2, *args[1:]] |
|
if m in [BottleneckCSP, C3]: |
|
args.insert(2, n) |
|
n = 1 |
|
elif m is nn.BatchNorm2d: |
|
args = [ch[f]] |
|
elif m is Concat: |
|
c2 = sum([ch[x] for x in f]) |
|
elif m is Detect: |
|
args.append([ch[x] for x in f]) |
|
if isinstance(args[1], int): |
|
args[1] = [list(range(args[1] * 2))] * len(f) |
|
elif m is Contract: |
|
c2 = ch[f] * args[0] ** 2 |
|
elif m is Expand: |
|
c2 = ch[f] // args[0] ** 2 |
|
else: |
|
c2 = ch[f] |
|
|
|
m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) |
|
t = str(m)[8:-2].replace("__main__.", "") |
|
np = sum([x.numel() for x in m_.parameters()]) |
|
m_.i, m_.f, m_.type, m_.np = ( |
|
i, |
|
f, |
|
t, |
|
np, |
|
) |
|
logger.info("%3s%18s%3s%10.0f %-40s%-30s" % (i, f, n, np, t, args)) |
|
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) |
|
layers.append(m_) |
|
if i == 0: |
|
ch = [] |
|
ch.append(c2) |
|
return nn.Sequential(*layers), sorted(save) |
|
|