First model version

67bb36a over 2 years ago

No virus

10.5 kB

	import math
	import numpy as np
	import torch
	import torch.nn as nn
	from PIL import Image, ImageDraw


	def autopad(k, p=None): # kernel, padding
	# Pad to 'same'
	if p is None:
	p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
	return p


	class DepthSeperabelConv2d(nn.Module):
	"""
	DepthSeperable Convolution 2d with residual connection
	"""

	def __init__(self, inplanes, planes, kernel_size=3, stride=1, downsample=None, act=True):
	super(DepthSeperabelConv2d, self).__init__()
	self.depthwise = nn.Sequential(
	nn.Conv2d(inplanes, inplanes, kernel_size, stride=stride, groups=inplanes, padding=kernel_size//2, bias=False),
	nn.BatchNorm2d(inplanes, momentum=BN_MOMENTUM)
	)
	# self.depthwise = nn.Conv2d(inplanes, inplanes, kernel_size, stride=stride, groups=inplanes, padding=1, bias=False)
	# self.pointwise = nn.Conv2d(inplanes, planes, 1, bias=False)

	self.pointwise = nn.Sequential(
	nn.Conv2d(inplanes, planes, 1, bias=False),
	nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
	)
	self.downsample = downsample
	self.stride = stride
	try:
	self.act = nn.Hardswish() if act else nn.Identity()
	except:
	self.act = nn.Identity()

	def forward(self, x):
	#residual = x

	out = self.depthwise(x)
	out = self.act(out)
	out = self.pointwise(out)

	if self.downsample is not None:
	residual = self.downsample(x)
	out = self.act(out)

	return out



	class SharpenConv(nn.Module):
	# SharpenConv convolution
	def __init__(self, c1, c2, k=3, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
	super(SharpenConv, self).__init__()
	sobel_kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype='float32')
	kenel_weight = np.vstack([sobel_kernel]c2c1).reshape(c2,c1,3,3)
	self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
	self.conv.weight.data = torch.from_numpy(kenel_weight)
	self.conv.weight.requires_grad = False
	self.bn = nn.BatchNorm2d(c2)
	try:
	self.act = nn.Hardswish() if act else nn.Identity()
	except:
	self.act = nn.Identity()

	def forward(self, x):
	return self.act(self.bn(self.conv(x)))

	def fuseforward(self, x):
	return self.act(self.conv(x))


	class Conv(nn.Module):
	# Standard convolution
	def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
	super(Conv, self).__init__()
	self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
	self.bn = nn.BatchNorm2d(c2)
	try:
	self.act = nn.Hardswish() if act else nn.Identity()
	except:
	self.act = nn.Identity()

	def forward(self, x):
	return self.act(self.bn(self.conv(x)))

	def fuseforward(self, x):
	return self.act(self.conv(x))


	class Bottleneck(nn.Module):
	# Standard bottleneck
	def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
	super(Bottleneck, self).__init__()
	c_ = int(c2 * e) # hidden channels
	self.cv1 = Conv(c1, c_, 1, 1)
	self.cv2 = Conv(c_, c2, 3, 1, g=g)
	self.add = shortcut and c1 == c2

	def forward(self, x):
	return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))


	class BottleneckCSP(nn.Module):
	# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
	def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
	super(BottleneckCSP, self).__init__()
	c_ = int(c2 * e) # hidden channels
	self.cv1 = Conv(c1, c_, 1, 1)
	self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
	self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
	self.cv4 = Conv(2 * c_, c2, 1, 1)
	self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
	self.act = nn.LeakyReLU(0.1, inplace=True)
	self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])

	def forward(self, x):
	y1 = self.cv3(self.m(self.cv1(x)))
	y2 = self.cv2(x)
	return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))


	class SPP(nn.Module):
	# Spatial pyramid pooling layer used in YOLOv3-SPP
	def __init__(self, c1, c2, k=(5, 9, 13)):
	super(SPP, self).__init__()
	c_ = c1 // 2 # hidden channels
	self.cv1 = Conv(c1, c_, 1, 1)
	self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
	self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])

	def forward(self, x):
	x = self.cv1(x)
	return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))


	class Focus(nn.Module):
	# Focus wh information into c-space
	# slice concat conv
	def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
	super(Focus, self).__init__()
	self.conv = Conv(c1 * 4, c2, k, s, p, g, act)

	def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
	return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))


	class Concat(nn.Module):
	# Concatenate a list of tensors along dimension
	def __init__(self, dimension=1):
	super(Concat, self).__init__()
	self.d = dimension

	def forward(self, x):
	""" print("***********************")
	for f in x:
	print(f.shape) """
	return torch.cat(x, self.d)


	class Detect(nn.Module):
	stride = None # strides computed during build

	def __init__(self, nc=13, anchors=(), ch=()): # detection layer
	super(Detect, self).__init__()
	self.nc = nc # number of classes
	self.no = nc + 5 # number of outputs per anchor 85
	self.nl = len(anchors) # number of detection layers 3
	self.na = len(anchors[0]) // 2 # number of anchors 3
	self.grid = [torch.zeros(1)] * self.nl # init grid
	a = torch.tensor(anchors).float().view(self.nl, -1, 2)
	self.register_buffer('anchors', a) # shape(nl,na,2)
	self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2)
	self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv

	def forward(self, x):
	z = [] # inference output
	for i in range(self.nl):
	x[i] = self.m[i](x[i]) # conv
	# print(str(i)+str(x[i].shape))
	bs, _, ny, nx = x[i].shape # x(bs,255,w,w) to x(bs,3,w,w,85)
	x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
	# print(str(i)+str(x[i].shape))

	if not self.training: # inference
	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()
	#print("**")
	#print(y.shape) #[1, 3, w, h, 85]
	#print(self.grid[i].shape) #[1, 3, w, h, 2]
	y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy
	y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
	"""print("**")
	print(y.shape) #[1, 3, w, h, 85]
	print(y.view(bs, -1, self.no).shape) #[1, 3wh, 85]"""
	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 Detections:
	# detections class for YOLOv5 inference results
	def __init__(self, imgs, pred, names=None):
	super(Detections, self).__init__()
	d = pred[0].device # device
	gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations
	self.imgs = imgs # list of images as numpy arrays
	self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
	self.names = names # class names
	self.xyxy = pred # xyxy pixels
	self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
	self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
	self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
	self.n = len(self.pred)

	def display(self, pprint=False, show=False, save=False):
	colors = color_list()
	for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
	str = f'Image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
	if pred is not None:
	for c in pred[:, -1].unique():
	n = (pred[:, -1] == c).sum() # detections per class
	str += f'{n} {self.names[int(c)]}s, ' # add to string
	if show or save:
	img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
	for *box, conf, cls in pred: # xyxy, confidence, class
	# str += '%s %.2f, ' % (names[int(cls)], conf) # label
	ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot
	if save:
	f = f'results{i}.jpg'
	str += f"saved to '{f}'"
	img.save(f) # save
	if show:
	img.show(f'Image {i}') # show
	if pprint:
	print(str)

	def print(self):
	self.display(pprint=True) # print results

	def show(self):
	self.display(show=True) # show results

	def save(self):
	self.display(save=True) # save results

	def __len__(self):
	return self.n

	def tolist(self):
	# return a list of Detections objects, i.e. 'for result in results.tolist():'
	x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)]
	for d in x:
	for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
	setattr(d, k, getattr(d, k)[0]) # pop out of list"""