parent
a2edd5fdca
commit
f73dc04ff6
@ -0,0 +1,10 @@
|
||||
DEEPSORT:
|
||||
REID_CKPT: "deep_sort/deep/checkpoint/ckpt.t7"
|
||||
MAX_DIST: 0.2
|
||||
MIN_CONFIDENCE: 0.3
|
||||
NMS_MAX_OVERLAP: 0.5
|
||||
MAX_IOU_DISTANCE: 0.7
|
||||
MAX_AGE: 70
|
||||
N_INIT: 3
|
||||
NN_BUDGET: 100
|
||||
|
@ -0,0 +1,3 @@
|
||||
# Deep Sort
|
||||
|
||||
This is the implemention of deep sort with pytorch.
|
@ -0,0 +1,21 @@
|
||||
from .deep_sort import DeepSort
|
||||
|
||||
|
||||
__all__ = ['DeepSort', 'build_tracker']
|
||||
|
||||
|
||||
def build_tracker(cfg, use_cuda):
|
||||
return DeepSort(cfg.DEEPSORT.REID_CKPT,
|
||||
max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
|
||||
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
|
||||
max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET, use_cuda=use_cuda)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
Binary file not shown.
@ -0,0 +1,15 @@
|
||||
import torch
|
||||
|
||||
features = torch.load("features.pth")
|
||||
qf = features["qf"]
|
||||
ql = features["ql"]
|
||||
gf = features["gf"]
|
||||
gl = features["gl"]
|
||||
|
||||
scores = qf.mm(gf.t())
|
||||
res = scores.topk(5, dim=1)[1][:,0]
|
||||
top1correct = gl[res].eq(ql).sum().item()
|
||||
|
||||
print("Acc top1:{:.3f}".format(top1correct/ql.size(0)))
|
||||
|
||||
|
@ -0,0 +1,55 @@
|
||||
import torch
|
||||
import torchvision.transforms as transforms
|
||||
import numpy as np
|
||||
import cv2
|
||||
import logging
|
||||
|
||||
from .model import Net
|
||||
|
||||
class Extractor(object):
|
||||
def __init__(self, model_path, use_cuda=True):
|
||||
self.net = Net(reid=True)
|
||||
self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
|
||||
state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)['net_dict']
|
||||
self.net.load_state_dict(state_dict)
|
||||
logger = logging.getLogger("root.tracker")
|
||||
logger.info("Loading weights from {}... Done!".format(model_path))
|
||||
self.net.to(self.device)
|
||||
self.size = (64, 128)
|
||||
self.norm = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
|
||||
|
||||
def _preprocess(self, im_crops):
|
||||
"""
|
||||
TODO:
|
||||
1. to float with scale from 0 to 1
|
||||
2. resize to (64, 128) as Market1501 dataset did
|
||||
3. concatenate to a numpy array
|
||||
3. to torch Tensor
|
||||
4. normalize
|
||||
"""
|
||||
def _resize(im, size):
|
||||
return cv2.resize(im.astype(np.float32)/255., size)
|
||||
|
||||
im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze(0) for im in im_crops], dim=0).float()
|
||||
return im_batch
|
||||
|
||||
|
||||
def __call__(self, im_crops):
|
||||
im_batch = self._preprocess(im_crops)
|
||||
with torch.no_grad():
|
||||
im_batch = im_batch.to(self.device)
|
||||
features = self.net(im_batch)
|
||||
return features.cpu().numpy()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
img = cv2.imread("demo.jpg")[:,:,(2,1,0)]
|
||||
extr = Extractor("checkpoint/ckpt.t7")
|
||||
feature = extr(img)
|
||||
print(feature.shape)
|
||||
|
@ -0,0 +1,104 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
def __init__(self, c_in, c_out,is_downsample=False):
|
||||
super(BasicBlock,self).__init__()
|
||||
self.is_downsample = is_downsample
|
||||
if is_downsample:
|
||||
self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=2, padding=1, bias=False)
|
||||
else:
|
||||
self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=1, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(c_out)
|
||||
self.relu = nn.ReLU(True)
|
||||
self.conv2 = nn.Conv2d(c_out,c_out,3,stride=1,padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(c_out)
|
||||
if is_downsample:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, 1, stride=2, bias=False),
|
||||
nn.BatchNorm2d(c_out)
|
||||
)
|
||||
elif c_in != c_out:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, 1, stride=1, bias=False),
|
||||
nn.BatchNorm2d(c_out)
|
||||
)
|
||||
self.is_downsample = True
|
||||
|
||||
def forward(self,x):
|
||||
y = self.conv1(x)
|
||||
y = self.bn1(y)
|
||||
y = self.relu(y)
|
||||
y = self.conv2(y)
|
||||
y = self.bn2(y)
|
||||
if self.is_downsample:
|
||||
x = self.downsample(x)
|
||||
return F.relu(x.add(y),True)
|
||||
|
||||
def make_layers(c_in,c_out,repeat_times, is_downsample=False):
|
||||
blocks = []
|
||||
for i in range(repeat_times):
|
||||
if i ==0:
|
||||
blocks += [BasicBlock(c_in,c_out, is_downsample=is_downsample),]
|
||||
else:
|
||||
blocks += [BasicBlock(c_out,c_out),]
|
||||
return nn.Sequential(*blocks)
|
||||
|
||||
class Net(nn.Module):
|
||||
def __init__(self, num_classes=751 ,reid=False):
|
||||
super(Net,self).__init__()
|
||||
# 3 128 64
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(3,64,3,stride=1,padding=1),
|
||||
nn.BatchNorm2d(64),
|
||||
nn.ReLU(inplace=True),
|
||||
# nn.Conv2d(32,32,3,stride=1,padding=1),
|
||||
# nn.BatchNorm2d(32),
|
||||
# nn.ReLU(inplace=True),
|
||||
nn.MaxPool2d(3,2,padding=1),
|
||||
)
|
||||
# 32 64 32
|
||||
self.layer1 = make_layers(64,64,2,False)
|
||||
# 32 64 32
|
||||
self.layer2 = make_layers(64,128,2,True)
|
||||
# 64 32 16
|
||||
self.layer3 = make_layers(128,256,2,True)
|
||||
# 128 16 8
|
||||
self.layer4 = make_layers(256,512,2,True)
|
||||
# 256 8 4
|
||||
self.avgpool = nn.AvgPool2d((8,4),1)
|
||||
# 256 1 1
|
||||
self.reid = reid
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Linear(512, 256),
|
||||
nn.BatchNorm1d(256),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout(),
|
||||
nn.Linear(256, num_classes),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
x = self.avgpool(x)
|
||||
x = x.view(x.size(0),-1)
|
||||
# B x 128
|
||||
if self.reid:
|
||||
x = x.div(x.norm(p=2,dim=1,keepdim=True))
|
||||
return x
|
||||
# classifier
|
||||
x = self.classifier(x)
|
||||
return x
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
net = Net()
|
||||
x = torch.randn(4,3,128,64)
|
||||
y = net(x)
|
||||
import ipdb; ipdb.set_trace()
|
||||
|
||||
|
@ -0,0 +1,106 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
def __init__(self, c_in, c_out,is_downsample=False):
|
||||
super(BasicBlock,self).__init__()
|
||||
self.is_downsample = is_downsample
|
||||
if is_downsample:
|
||||
self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=2, padding=1, bias=False)
|
||||
else:
|
||||
self.conv1 = nn.Conv2d(c_in, c_out, 3, stride=1, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(c_out)
|
||||
self.relu = nn.ReLU(True)
|
||||
self.conv2 = nn.Conv2d(c_out,c_out,3,stride=1,padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(c_out)
|
||||
if is_downsample:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, 1, stride=2, bias=False),
|
||||
nn.BatchNorm2d(c_out)
|
||||
)
|
||||
elif c_in != c_out:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, 1, stride=1, bias=False),
|
||||
nn.BatchNorm2d(c_out)
|
||||
)
|
||||
self.is_downsample = True
|
||||
|
||||
def forward(self,x):
|
||||
y = self.conv1(x)
|
||||
y = self.bn1(y)
|
||||
y = self.relu(y)
|
||||
y = self.conv2(y)
|
||||
y = self.bn2(y)
|
||||
if self.is_downsample:
|
||||
x = self.downsample(x)
|
||||
return F.relu(x.add(y),True)
|
||||
|
||||
def make_layers(c_in,c_out,repeat_times, is_downsample=False):
|
||||
blocks = []
|
||||
for i in range(repeat_times):
|
||||
if i ==0:
|
||||
blocks += [BasicBlock(c_in,c_out, is_downsample=is_downsample),]
|
||||
else:
|
||||
blocks += [BasicBlock(c_out,c_out),]
|
||||
return nn.Sequential(*blocks)
|
||||
|
||||
class Net(nn.Module):
|
||||
def __init__(self, num_classes=625 ,reid=False):
|
||||
super(Net,self).__init__()
|
||||
# 3 128 64
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(3,32,3,stride=1,padding=1),
|
||||
nn.BatchNorm2d(32),
|
||||
nn.ELU(inplace=True),
|
||||
nn.Conv2d(32,32,3,stride=1,padding=1),
|
||||
nn.BatchNorm2d(32),
|
||||
nn.ELU(inplace=True),
|
||||
nn.MaxPool2d(3,2,padding=1),
|
||||
)
|
||||
# 32 64 32
|
||||
self.layer1 = make_layers(32,32,2,False)
|
||||
# 32 64 32
|
||||
self.layer2 = make_layers(32,64,2,True)
|
||||
# 64 32 16
|
||||
self.layer3 = make_layers(64,128,2,True)
|
||||
# 128 16 8
|
||||
self.dense = nn.Sequential(
|
||||
nn.Dropout(p=0.6),
|
||||
nn.Linear(128*16*8, 128),
|
||||
nn.BatchNorm1d(128),
|
||||
nn.ELU(inplace=True)
|
||||
)
|
||||
# 256 1 1
|
||||
self.reid = reid
|
||||
self.batch_norm = nn.BatchNorm1d(128)
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Linear(128, num_classes),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
|
||||
x = x.view(x.size(0),-1)
|
||||
if self.reid:
|
||||
x = self.dense[0](x)
|
||||
x = self.dense[1](x)
|
||||
x = x.div(x.norm(p=2,dim=1,keepdim=True))
|
||||
return x
|
||||
x = self.dense(x)
|
||||
# B x 128
|
||||
# classifier
|
||||
x = self.classifier(x)
|
||||
return x
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
net = Net(reid=True)
|
||||
x = torch.randn(4,3,128,64)
|
||||
y = net(x)
|
||||
import ipdb; ipdb.set_trace()
|
||||
|
||||
|
@ -0,0 +1,77 @@
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
import torchvision
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from model import Net
|
||||
|
||||
parser = argparse.ArgumentParser(description="Train on market1501")
|
||||
parser.add_argument("--data-dir",default='data',type=str)
|
||||
parser.add_argument("--no-cuda",action="store_true")
|
||||
parser.add_argument("--gpu-id",default=0,type=int)
|
||||
args = parser.parse_args()
|
||||
|
||||
# device
|
||||
device = "cuda:{}".format(args.gpu_id) if torch.cuda.is_available() and not args.no_cuda else "cpu"
|
||||
if torch.cuda.is_available() and not args.no_cuda:
|
||||
cudnn.benchmark = True
|
||||
|
||||
# data loader
|
||||
root = args.data_dir
|
||||
query_dir = os.path.join(root,"query")
|
||||
gallery_dir = os.path.join(root,"gallery")
|
||||
transform = torchvision.transforms.Compose([
|
||||
torchvision.transforms.Resize((128,64)),
|
||||
torchvision.transforms.ToTensor(),
|
||||
torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
||||
])
|
||||
queryloader = torch.utils.data.DataLoader(
|
||||
torchvision.datasets.ImageFolder(query_dir, transform=transform),
|
||||
batch_size=64, shuffle=False
|
||||
)
|
||||
galleryloader = torch.utils.data.DataLoader(
|
||||
torchvision.datasets.ImageFolder(gallery_dir, transform=transform),
|
||||
batch_size=64, shuffle=False
|
||||
)
|
||||
|
||||
# net definition
|
||||
net = Net(reid=True)
|
||||
assert os.path.isfile("./checkpoint/ckpt.t7"), "Error: no checkpoint file found!"
|
||||
print('Loading from checkpoint/ckpt.t7')
|
||||
checkpoint = torch.load("./checkpoint/ckpt.t7")
|
||||
net_dict = checkpoint['net_dict']
|
||||
net.load_state_dict(net_dict, strict=False)
|
||||
net.eval()
|
||||
net.to(device)
|
||||
|
||||
# compute features
|
||||
query_features = torch.tensor([]).float()
|
||||
query_labels = torch.tensor([]).long()
|
||||
gallery_features = torch.tensor([]).float()
|
||||
gallery_labels = torch.tensor([]).long()
|
||||
|
||||
with torch.no_grad():
|
||||
for idx,(inputs,labels) in enumerate(queryloader):
|
||||
inputs = inputs.to(device)
|
||||
features = net(inputs).cpu()
|
||||
query_features = torch.cat((query_features, features), dim=0)
|
||||
query_labels = torch.cat((query_labels, labels))
|
||||
|
||||
for idx,(inputs,labels) in enumerate(galleryloader):
|
||||
inputs = inputs.to(device)
|
||||
features = net(inputs).cpu()
|
||||
gallery_features = torch.cat((gallery_features, features), dim=0)
|
||||
gallery_labels = torch.cat((gallery_labels, labels))
|
||||
|
||||
gallery_labels -= 2
|
||||
|
||||
# save features
|
||||
features = {
|
||||
"qf": query_features,
|
||||
"ql": query_labels,
|
||||
"gf": gallery_features,
|
||||
"gl": gallery_labels
|
||||
}
|
||||
torch.save(features,"features.pth")
|
After Width: | Height: | Size: 59 KiB |
@ -0,0 +1,189 @@
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
import torchvision
|
||||
|
||||
from model import Net
|
||||
|
||||
parser = argparse.ArgumentParser(description="Train on market1501")
|
||||
parser.add_argument("--data-dir",default='data',type=str)
|
||||
parser.add_argument("--no-cuda",action="store_true")
|
||||
parser.add_argument("--gpu-id",default=0,type=int)
|
||||
parser.add_argument("--lr",default=0.1, type=float)
|
||||
parser.add_argument("--interval",'-i',default=20,type=int)
|
||||
parser.add_argument('--resume', '-r',action='store_true')
|
||||
args = parser.parse_args()
|
||||
|
||||
# device
|
||||
device = "cuda:{}".format(args.gpu_id) if torch.cuda.is_available() and not args.no_cuda else "cpu"
|
||||
if torch.cuda.is_available() and not args.no_cuda:
|
||||
cudnn.benchmark = True
|
||||
|
||||
# data loading
|
||||
root = args.data_dir
|
||||
train_dir = os.path.join(root,"train")
|
||||
test_dir = os.path.join(root,"test")
|
||||
transform_train = torchvision.transforms.Compose([
|
||||
torchvision.transforms.RandomCrop((128,64),padding=4),
|
||||
torchvision.transforms.RandomHorizontalFlip(),
|
||||
torchvision.transforms.ToTensor(),
|
||||
torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
||||
])
|
||||
transform_test = torchvision.transforms.Compose([
|
||||
torchvision.transforms.Resize((128,64)),
|
||||
torchvision.transforms.ToTensor(),
|
||||
torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
||||
])
|
||||
trainloader = torch.utils.data.DataLoader(
|
||||
torchvision.datasets.ImageFolder(train_dir, transform=transform_train),
|
||||
batch_size=64,shuffle=True
|
||||
)
|
||||
testloader = torch.utils.data.DataLoader(
|
||||
torchvision.datasets.ImageFolder(test_dir, transform=transform_test),
|
||||
batch_size=64,shuffle=True
|
||||
)
|
||||
num_classes = max(len(trainloader.dataset.classes), len(testloader.dataset.classes))
|
||||
|
||||
# net definition
|
||||
start_epoch = 0
|
||||
net = Net(num_classes=num_classes)
|
||||
if args.resume:
|
||||
assert os.path.isfile("./checkpoint/ckpt.t7"), "Error: no checkpoint file found!"
|
||||
print('Loading from checkpoint/ckpt.t7')
|
||||
checkpoint = torch.load("./checkpoint/ckpt.t7")
|
||||
# import ipdb; ipdb.set_trace()
|
||||
net_dict = checkpoint['net_dict']
|
||||
net.load_state_dict(net_dict)
|
||||
best_acc = checkpoint['acc']
|
||||
start_epoch = checkpoint['epoch']
|
||||
net.to(device)
|
||||
|
||||
# loss and optimizer
|
||||
criterion = torch.nn.CrossEntropyLoss()
|
||||
optimizer = torch.optim.SGD(net.parameters(), args.lr, momentum=0.9, weight_decay=5e-4)
|
||||
best_acc = 0.
|
||||
|
||||
# train function for each epoch
|
||||
def train(epoch):
|
||||
print("\nEpoch : %d"%(epoch+1))
|
||||
net.train()
|
||||
training_loss = 0.
|
||||
train_loss = 0.
|
||||
correct = 0
|
||||
total = 0
|
||||
interval = args.interval
|
||||
start = time.time()
|
||||
for idx, (inputs, labels) in enumerate(trainloader):
|
||||
# forward
|
||||
inputs,labels = inputs.to(device),labels.to(device)
|
||||
outputs = net(inputs)
|
||||
loss = criterion(outputs, labels)
|
||||
|
||||
# backward
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
# accumurating
|
||||
training_loss += loss.item()
|
||||
train_loss += loss.item()
|
||||
correct += outputs.max(dim=1)[1].eq(labels).sum().item()
|
||||
total += labels.size(0)
|
||||
|
||||
# print
|
||||
if (idx+1)%interval == 0:
|
||||
end = time.time()
|
||||
print("[progress:{:.1f}%]time:{:.2f}s Loss:{:.5f} Correct:{}/{} Acc:{:.3f}%".format(
|
||||
100.*(idx+1)/len(trainloader), end-start, training_loss/interval, correct, total, 100.*correct/total
|
||||
))
|
||||
training_loss = 0.
|
||||
start = time.time()
|
||||
|
||||
return train_loss/len(trainloader), 1.- correct/total
|
||||
|
||||
def test(epoch):
|
||||
global best_acc
|
||||
net.eval()
|
||||
test_loss = 0.
|
||||
correct = 0
|
||||
total = 0
|
||||
start = time.time()
|
||||
with torch.no_grad():
|
||||
for idx, (inputs, labels) in enumerate(testloader):
|
||||
inputs, labels = inputs.to(device), labels.to(device)
|
||||
outputs = net(inputs)
|
||||
loss = criterion(outputs, labels)
|
||||
|
||||
test_loss += loss.item()
|
||||
correct += outputs.max(dim=1)[1].eq(labels).sum().item()
|
||||
total += labels.size(0)
|
||||
|
||||
print("Testing ...")
|
||||
end = time.time()
|
||||
print("[progress:{:.1f}%]time:{:.2f}s Loss:{:.5f} Correct:{}/{} Acc:{:.3f}%".format(
|
||||
100.*(idx+1)/len(testloader), end-start, test_loss/len(testloader), correct, total, 100.*correct/total
|
||||
))
|
||||
|
||||
# saving checkpoint
|
||||
acc = 100.*correct/total
|
||||
if acc > best_acc:
|
||||
best_acc = acc
|
||||
print("Saving parameters to checkpoint/ckpt.t7")
|
||||
checkpoint = {
|
||||
'net_dict':net.state_dict(),
|
||||
'acc':acc,
|
||||
'epoch':epoch,
|
||||
}
|
||||
if not os.path.isdir('checkpoint'):
|
||||
os.mkdir('checkpoint')
|
||||
torch.save(checkpoint, './checkpoint/ckpt.t7')
|
||||
|
||||
return test_loss/len(testloader), 1.- correct/total
|
||||
|
||||
# plot figure
|
||||
x_epoch = []
|
||||
record = {'train_loss':[], 'train_err':[], 'test_loss':[], 'test_err':[]}
|
||||
fig = plt.figure()
|
||||
ax0 = fig.add_subplot(121, title="loss")
|
||||
ax1 = fig.add_subplot(122, title="top1err")
|
||||
def draw_curve(epoch, train_loss, train_err, test_loss, test_err):
|
||||
global record
|
||||
record['train_loss'].append(train_loss)
|
||||
record['train_err'].append(train_err)
|
||||
record['test_loss'].append(test_loss)
|
||||
record['test_err'].append(test_err)
|
||||
|
||||
x_epoch.append(epoch)
|
||||
ax0.plot(x_epoch, record['train_loss'], 'bo-', label='train')
|
||||
ax0.plot(x_epoch, record['test_loss'], 'ro-', label='val')
|
||||
ax1.plot(x_epoch, record['train_err'], 'bo-', label='train')
|
||||
ax1.plot(x_epoch, record['test_err'], 'ro-', label='val')
|
||||
if epoch == 0:
|
||||
ax0.legend()
|
||||
ax1.legend()
|
||||
fig.savefig("train.jpg")
|
||||
|
||||
# lr decay
|
||||
def lr_decay():
|
||||
global optimizer
|
||||
for params in optimizer.param_groups:
|
||||
params['lr'] *= 0.1
|
||||
lr = params['lr']
|
||||
print("Learning rate adjusted to {}".format(lr))
|
||||
|
||||
def main():
|
||||
for epoch in range(start_epoch, start_epoch+40):
|
||||
train_loss, train_err = train(epoch)
|
||||
test_loss, test_err = test(epoch)
|
||||
draw_curve(epoch, train_loss, train_err, test_loss, test_err)
|
||||
if (epoch+1)%20==0:
|
||||
lr_decay()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
@ -0,0 +1,131 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .deep.feature_extractor import Extractor
|
||||
from .sort.nn_matching import NearestNeighborDistanceMetric
|
||||
from .sort.preprocessing import non_max_suppression
|
||||
from .sort.detection import Detection
|
||||
from .sort.tracker import Tracker
|
||||
|
||||
|
||||
__all__ = ['DeepSort']
|
||||
|
||||
|
||||
class DeepSort(object):
|
||||
def __init__(self, model_path, max_dist=0.2, min_confidence=0.3, nms_max_overlap=1.0, max_iou_distance=0.7, max_age=70, n_init=3, nn_budget=100, use_cuda=True):
|
||||
self.min_confidence = min_confidence
|
||||
self.nms_max_overlap = nms_max_overlap
|
||||
|
||||
self.extractor = Extractor(model_path, use_cuda=use_cuda)
|
||||
|
||||
max_cosine_distance = max_dist
|
||||
nn_budget = 100
|
||||
metric = NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
|
||||
|
||||
# tracker maintain a list contains(self.tracks) for each Track object
|
||||
self.tracker = Tracker(metric, max_iou_distance=max_iou_distance, max_age=max_age, n_init=n_init)
|
||||
|
||||
def update(self, bbox_xywh, confidences, ori_img):
|
||||
# bbox_xywh (#obj,4), [xc,yc, w, h] bounding box for each person
|
||||
# conf (#obj,1)
|
||||
|
||||
self.height, self.width = ori_img.shape[:2]
|
||||
|
||||
# get appearance feature with neural network (Deep) *********************************************************
|
||||
features = self._get_features(bbox_xywh, ori_img)
|
||||
|
||||
bbox_tlwh = self._xywh_to_tlwh(bbox_xywh) # # [cx,cy,w,h] -> [x1,y1,w,h] top left
|
||||
|
||||
# generate detections class object for each person *********************************************************
|
||||
# filter object with less confidence
|
||||
# each Detection obj maintain the location(bbox_tlwh), confidence(conf), and appearance feature
|
||||
detections = [Detection(bbox_tlwh[i], conf, features[i]) for i,conf in enumerate(confidences) if conf>self.min_confidence]
|
||||
|
||||
# run on non-maximum supression (useless) *******************************************************************
|
||||
boxes = np.array([d.tlwh for d in detections])
|
||||
scores = np.array([d.confidence for d in detections])
|
||||
indices = non_max_suppression(boxes, self.nms_max_overlap, scores) # Here, nms_max_overlap is 1
|
||||
detections = [detections[i] for i in indices]
|
||||
|
||||
# update tracker ********************************************************************************************
|
||||
self.tracker.predict() # predict based on t-1 info
|
||||
# for first frame, this function do nothing
|
||||
|
||||
# detections is the measurement results as time T
|
||||
self.tracker.update(detections)
|
||||
|
||||
# output bbox identities ************************************************************************************
|
||||
outputs = []
|
||||
for track in self.tracker.tracks:
|
||||
|
||||
if not track.is_confirmed() or track.time_since_update > 1:
|
||||
continue
|
||||
|
||||
box = track.to_tlwh() # (xc,yc,a,h) to (x1,y1,w,h)
|
||||
x1,y1,x2,y2 = self._tlwh_to_xyxy(box)
|
||||
track_id = track.track_id
|
||||
outputs.append(np.array([x1,y1,x2,y2,track_id], dtype=np.int))
|
||||
if len(outputs) > 0:
|
||||
outputs = np.stack(outputs,axis=0) # (#obj, 5) (x1,y1,x2,y2,ID)
|
||||
return outputs
|
||||
|
||||
|
||||
"""
|
||||
TODO:
|
||||
Convert bbox from xc_yc_w_h to xtl_ytl_w_h
|
||||
Thanks JieChen91@github.com for reporting this bug!
|
||||
"""
|
||||
@staticmethod
|
||||
def _xywh_to_tlwh(bbox_xywh):
|
||||
if isinstance(bbox_xywh, np.ndarray):
|
||||
bbox_tlwh = bbox_xywh.copy()
|
||||
elif isinstance(bbox_xywh, torch.Tensor):
|
||||
bbox_tlwh = bbox_xywh.clone()
|
||||
bbox_tlwh[:,0] = bbox_xywh[:,0] - bbox_xywh[:,2]/2.
|
||||
bbox_tlwh[:,1] = bbox_xywh[:,1] - bbox_xywh[:,3]/2.
|
||||
return bbox_tlwh
|
||||
|
||||
|
||||
def _xywh_to_xyxy(self, bbox_xywh):
|
||||
x,y,w,h = bbox_xywh
|
||||
x1 = max(int(x-w/2),0)
|
||||
x2 = min(int(x+w/2),self.width-1)
|
||||
y1 = max(int(y-h/2),0)
|
||||
y2 = min(int(y+h/2),self.height-1)
|
||||
return x1,y1,x2,y2
|
||||
|
||||
def _tlwh_to_xyxy(self, bbox_tlwh):
|
||||
"""
|
||||
TODO:
|
||||
Convert bbox from xtl_ytl_w_h to xc_yc_w_h
|
||||
Thanks JieChen91@github.com for reporting this bug!
|
||||
"""
|
||||
x,y,w,h = bbox_tlwh
|
||||
x1 = max(int(x),0)
|
||||
x2 = min(int(x+w),self.width-1)
|
||||
y1 = max(int(y),0)
|
||||
y2 = min(int(y+h),self.height-1)
|
||||
return x1,y1,x2,y2
|
||||
|
||||
def _xyxy_to_tlwh(self, bbox_xyxy):
|
||||
x1,y1,x2,y2 = bbox_xyxy
|
||||
|
||||
t = x1
|
||||
l = y1
|
||||
w = int(x2-x1)
|
||||
h = int(y2-y1)
|
||||
return t,l,w,h
|
||||
|
||||
def _get_features(self, bbox_xywh, ori_img):
|
||||
im_crops = []
|
||||
for box in bbox_xywh:
|
||||
x1,y1,x2,y2 = self._xywh_to_xyxy(box)
|
||||
im = ori_img[y1:y2,x1:x2]
|
||||
im_crops.append(im)
|
||||
if im_crops:
|
||||
features = self.extractor(im_crops)
|
||||
else:
|
||||
features = np.array([])
|
||||
return features
|
||||
|
||||
|
@ -0,0 +1,49 @@
|
||||
# vim: expandtab:ts=4:sw=4
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Detection(object):
|
||||
"""
|
||||
This class represents a bounding box detection in a single image.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tlwh : array_like
|
||||
Bounding box in format `(x, y, w, h)`.
|
||||
confidence : float
|
||||
Detector confidence score.
|
||||
feature : array_like
|
||||
A feature vector that describes the object contained in this image.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
tlwh : ndarray
|
||||
Bounding box in format `(top left x, top left y, width, height)`.
|
||||
confidence : ndarray
|
||||
Detector confidence score.
|
||||
feature : ndarray | NoneType
|
||||
A feature vector that describes the object contained in this image.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, tlwh, confidence, feature):
|
||||
self.tlwh = np.asarray(tlwh, dtype=np.float) # x1, y1, w, h
|
||||
self.confidence = float(confidence)
|
||||
self.feature = np.asarray(feature, dtype=np.float32)
|
||||
|
||||
def to_tlbr(self):
|
||||
"""Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
|
||||
`(top left, bottom right)`.
|
||||
"""
|
||||
ret = self.tlwh.copy()
|
||||
ret[2:] += ret[:2]
|
||||
return ret
|
||||
|
||||
def to_xyah(self):
|
||||
"""Convert bounding box to format `(center x, center y, aspect ratio,
|
||||
height)`, where the aspect ratio is `width / height`.
|
||||
"""
|
||||
ret = self.tlwh.copy()
|
||||
ret[:2] += ret[2:] / 2
|
||||
ret[2] /= ret[3]
|
||||
return ret
|
@ -0,0 +1,81 @@
|
||||
# vim: expandtab:ts=4:sw=4
|
||||
from __future__ import absolute_import
|
||||
import numpy as np
|
||||
from . import linear_assignment
|
||||
|
||||
|
||||
def iou(bbox, candidates):
|
||||
"""Computer intersection over union.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
bbox : ndarray
|
||||
A bounding box in format `(top left x, top left y, width, height)`.
|
||||
candidates : ndarray
|
||||
A matrix of candidate bounding boxes (one per row) in the same format
|
||||
as `bbox`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
The intersection over union in [0, 1] between the `bbox` and each
|
||||
candidate. A higher score means a larger fraction of the `bbox` is
|
||||
occluded by the candidate.
|
||||
|
||||
"""
|
||||
bbox_tl, bbox_br = bbox[:2], bbox[:2] + bbox[2:]
|
||||
candidates_tl = candidates[:, :2]
|
||||
candidates_br = candidates[:, :2] + candidates[:, 2:]
|
||||
|
||||
tl = np.c_[np.maximum(bbox_tl[0], candidates_tl[:, 0])[:, np.newaxis],
|
||||
np.maximum(bbox_tl[1], candidates_tl[:, 1])[:, np.newaxis]]
|
||||
br = np.c_[np.minimum(bbox_br[0], candidates_br[:, 0])[:, np.newaxis],
|
||||
np.minimum(bbox_br[1], candidates_br[:, 1])[:, np.newaxis]]
|
||||
wh = np.maximum(0., br - tl)
|
||||
|
||||
area_intersection = wh.prod(axis=1)
|
||||
area_bbox = bbox[2:].prod()
|
||||
area_candidates = candidates[:, 2:].prod(axis=1)
|
||||
return area_intersection / (area_bbox + area_candidates - area_intersection)
|
||||
|
||||
|
||||
def iou_cost(tracks, detections, track_indices=None,
|
||||
detection_indices=None):
|
||||
"""An intersection over union distance metric.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tracks : List[deep_sort.track.Track]
|
||||
A list of tracks.
|
||||
detections : List[deep_sort.detection.Detection]
|
||||
A list of detections.
|
||||
track_indices : Optional[List[int]]
|
||||
A list of indices to tracks that should be matched. Defaults to
|
||||
all `tracks`.
|
||||
detection_indices : Optional[List[int]]
|
||||
A list of indices to detections that should be matched. Defaults
|
||||
to all `detections`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
Returns a cost matrix of shape
|
||||
len(track_indices), len(detection_indices) where entry (i, j) is
|
||||
`1 - iou(tracks[track_indices[i]], detections[detection_indices[j]])`.
|
||||
|
||||
"""
|
||||
if track_indices is None:
|
||||
track_indices = np.arange(len(tracks))
|
||||
if detection_indices is None:
|
||||
detection_indices = np.arange(len(detections))
|
||||
|
||||
cost_matrix = np.zeros((len(track_indices), len(detection_indices)))
|
||||
for row, track_idx in enumerate(track_indices):
|
||||
if tracks[track_idx].time_since_update > 1:
|
||||
cost_matrix[row, :] = linear_assignment.INFTY_COST
|
||||
continue
|
||||
|
||||
bbox = tracks[track_idx].to_tlwh()
|
||||
candidates = np.asarray([detections[i].tlwh for i in detection_indices])
|
||||
cost_matrix[row, :] = 1. - iou(bbox, candidates)
|
||||
return cost_matrix
|
@ -0,0 +1,192 @@
|
||||
# vim: expandtab:ts=4:sw=4
|
||||
from __future__ import absolute_import
|
||||
import numpy as np
|
||||
# from sklearn.utils.linear_assignment_ import linear_assignment
|
||||
from scipy.optimize import linear_sum_assignment as linear_assignment
|
||||
from . import kalman_filter
|
||||
|
||||
|
||||
INFTY_COST = 1e+5
|
||||
|
||||
|
||||
def min_cost_matching(
|
||||
distance_metric, max_distance, tracks, detections, track_indices=None,
|
||||
detection_indices=None):
|
||||
"""Solve linear assignment problem.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray
|
||||
The distance metric is given a list of tracks and detections as well as
|
||||
a list of N track indices and M detection indices. The metric should
|
||||
return the NxM dimensional cost matrix, where element (i, j) is the
|
||||
association cost between the i-th track in the given track indices and
|
||||
the j-th detection in the given detection_indices.
|
||||
max_distance : float
|
||||
Gating threshold. Associations with cost larger than this value are
|
||||
disregarded.
|
||||
tracks : List[track.Track]
|
||||
A list of predicted tracks at the current time step.
|
||||
detections : List[detection.Detection]
|
||||
A list of detections at the current time step.
|
||||
track_indices : List[int]
|
||||
List of track indices that maps rows in `cost_matrix` to tracks in
|
||||
`tracks` (see description above).
|
||||
detection_indices : List[int]
|
||||
List of detection indices that maps columns in `cost_matrix` to
|
||||
detections in `detections` (see description above).
|
||||
|
||||
Returns
|
||||
-------
|
||||
(List[(int, int)], List[int], List[int])
|
||||
Returns a tuple with the following three entries:
|
||||
* A list of matched track and detection indices.
|
||||
* A list of unmatched track indices.
|
||||
* A list of unmatched detection indices.
|
||||
|
||||
"""
|
||||
if track_indices is None:
|
||||
track_indices = np.arange(len(tracks))
|
||||
if detection_indices is None:
|
||||
detection_indices = np.arange(len(detections))
|
||||
|
||||
if len(detection_indices) == 0 or len(track_indices) == 0:
|
||||
return [], track_indices, detection_indices # Nothing to match.
|
||||
|
||||
cost_matrix = distance_metric(
|
||||
tracks, detections, track_indices, detection_indices)
|
||||
cost_matrix[cost_matrix > max_distance] = max_distance + 1e-5
|
||||
|
||||
row_indices, col_indices = linear_assignment(cost_matrix)
|
||||
|
||||
matches, unmatched_tracks, unmatched_detections = [], [], []
|
||||
for col, detection_idx in enumerate(detection_indices):
|
||||
if col not in col_indices:
|
||||
unmatched_detections.append(detection_idx)
|
||||
for row, track_idx in enumerate(track_indices):
|
||||
if row not in row_indices:
|
||||
unmatched_tracks.append(track_idx)
|
||||
for row, col in zip(row_indices, col_indices):
|
||||
track_idx = track_indices[row]
|
||||
detection_idx = detection_indices[col]
|
||||
if cost_matrix[row, col] > max_distance:
|
||||
unmatched_tracks.append(track_idx)
|
||||
unmatched_detections.append(detection_idx)
|
||||
else:
|
||||
matches.append((track_idx, detection_idx))
|
||||
return matches, unmatched_tracks, unmatched_detections
|
||||
|
||||
|
||||
def matching_cascade(
|
||||
distance_metric, max_distance, cascade_depth, tracks, detections,
|
||||
track_indices=None, detection_indices=None):
|
||||
"""Run matching cascade.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray
|
||||
The distance metric is given a list of tracks and detections as well as
|
||||
a list of N track indices and M detection indices. The metric should
|
||||
return the NxM dimensional cost matrix, where element (i, j) is the
|
||||
association cost between the i-th track in the given track indices and
|
||||
the j-th detection in the given detection indices.
|
||||
max_distance : float
|
||||
Gating threshold. Associations with cost larger than this value are
|
||||
disregarded.
|
||||
cascade_depth: int
|
||||
The cascade depth, should be se to the maximum track age.
|
||||
tracks : List[track.Track]
|
||||
A list of predicted tracks at the current time step.
|
||||
detections : List[detection.Detection]
|
||||
A list of detections at the current time step.
|
||||
track_indices : Optional[List[int]]
|
||||
List of track indices that maps rows in `cost_matrix` to tracks in
|
||||
`tracks` (see description above). Defaults to all tracks.
|
||||
detection_indices : Optional[List[int]]
|
||||
List of detection indices that maps columns in `cost_matrix` to
|
||||
detections in `detections` (see description above). Defaults to all
|
||||
detections.
|
||||
|
||||
Returns
|
||||
-------
|
||||
(List[(int, int)], List[int], List[int])
|
||||
Returns a tuple with the following three entries:
|
||||
* A list of matched track and detection indices.
|
||||
* A list of unmatched track indices.
|
||||
* A list of unmatched detection indices.
|
||||
|
||||
"""
|
||||
if track_indices is None:
|
||||
track_indices = list(range(len(tracks)))
|
||||
if detection_indices is None:
|
||||
detection_indices = list(range(len(detections)))
|
||||
|
||||
unmatched_detections = detection_indices
|
||||
matches = []
|
||||
for level in range(cascade_depth):
|
||||
if len(unmatched_detections) == 0: # No detections left
|
||||
break
|
||||
|
||||
track_indices_l = [
|
||||
k for k in track_indices
|
||||
if tracks[k].time_since_update == 1 + level
|
||||
]
|
||||
if len(track_indices_l) == 0: # Nothing to match at this level
|
||||
continue
|
||||
|
||||
matches_l, _, unmatched_detections = \
|
||||
min_cost_matching(
|
||||
distance_metric, max_distance, tracks, detections,
|
||||
track_indices_l, unmatched_detections)
|
||||
matches += matches_l
|
||||
unmatched_tracks = list(set(track_indices) - set(k for k, _ in matches))
|
||||
return matches, unmatched_tracks, unmatched_detections
|
||||
|
||||
|
||||
def gate_cost_matrix(
|
||||
kf, cost_matrix, tracks, detections, track_indices, detection_indices,
|
||||
gated_cost=INFTY_COST, only_position=False):
|
||||
"""Invalidate infeasible entries in cost matrix based on the state
|
||||
distributions obtained by Kalman filtering.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
kf : The Kalman filter.
|
||||
cost_matrix : ndarray
|
||||
The NxM dimensional cost matrix, where N is the number of track indices
|
||||
and M is the number of detection indices, such that entry (i, j) is the
|
||||
association cost between `tracks[track_indices[i]]` and
|
||||
`detections[detection_indices[j]]`.
|
||||
tracks : List[track.Track]
|
||||
A list of predicted tracks at the current time step.
|
||||
detections : List[detection.Detection]
|
||||
A list of detections at the current time step.
|
||||
track_indices : List[int]
|
||||
List of track indices that maps rows in `cost_matrix` to tracks in
|
||||
`tracks` (see description above).
|
||||
detection_indices : List[int]
|
||||
List of detection indices that maps columns in `cost_matrix` to
|
||||
detections in `detections` (see description above).
|
||||
gated_cost : Optional[float]
|
||||
Entries in the cost matrix corresponding to infeasible associations are
|
||||
set this value. Defaults to a very large value.
|
||||
only_position : Optional[bool]
|
||||
If True, only the x, y position of the state distribution is considered
|
||||
during gating. Defaults to False.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
Returns the modified cost matrix.
|
||||
|
||||
"""
|
||||
gating_dim = 2 if only_position else 4
|
||||
gating_threshold = kalman_filter.chi2inv95[gating_dim]
|
||||
measurements = np.asarray(
|
||||
[detections[i].to_xyah() for i in detection_indices])
|
||||
for row, track_idx in enumerate(track_indices):
|
||||
track = tracks[track_idx]
|
||||
gating_distance = kf.gating_distance(
|
||||
track.mean, track.covariance, measurements, only_position)
|
||||
cost_matrix[row, gating_distance > gating_threshold] = gated_cost
|
||||
return cost_matrix
|
@ -0,0 +1,177 @@
|
||||
# vim: expandtab:ts=4:sw=4
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _pdist(a, b):
|
||||
"""Compute pair-wise squared distance between points in `a` and `b`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a : array_like
|
||||
An NxM matrix of N samples of dimensionality M.
|
||||
b : array_like
|
||||
An LxM matrix of L samples of dimensionality M.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
Returns a matrix of size len(a), len(b) such that eleement (i, j)
|
||||
contains the squared distance between `a[i]` and `b[j]`.
|
||||
|
||||
"""
|
||||
a, b = np.asarray(a), np.asarray(b)
|
||||
if len(a) == 0 or len(b) == 0:
|
||||
return np.zeros((len(a), len(b)))
|
||||
a2, b2 = np.square(a).sum(axis=1), np.square(b).sum(axis=1)
|
||||
r2 = -2. * np.dot(a, b.T) + a2[:, None] + b2[None, :]
|
||||
r2 = np.clip(r2, 0., float(np.inf))
|
||||
return r2
|
||||
|
||||
|
||||
def _cosine_distance(a, b, data_is_normalized=False):
|
||||
"""Compute pair-wise cosine distance between points in `a` and `b`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a : array_like
|
||||
An NxM matrix of N samples of dimensionality M.
|
||||
b : array_like
|
||||
An LxM matrix of L samples of dimensionality M.
|
||||
data_is_normalized : Optional[bool]
|
||||
If True, assumes rows in a and b are unit length vectors.
|
||||
Otherwise, a and b are explicitly normalized to lenght 1.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
Returns a matrix of size len(a), len(b) such that eleement (i, j)
|
||||
contains the squared distance between `a[i]` and `b[j]`.
|
||||
|
||||
"""
|
||||
if not data_is_normalized:
|
||||
a = np.asarray(a) / np.linalg.norm(a, axis=1, keepdims=True)
|
||||
b = np.asarray(b) / np.linalg.norm(b, axis=1, keepdims=True)
|
||||
return 1. - np.dot(a, b.T)
|
||||
|
||||
|
||||
def _nn_euclidean_distance(x, y):
|
||||
""" Helper function for nearest neighbor distance metric (Euclidean).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : ndarray
|
||||
A matrix of N row-vectors (sample points).
|
||||
y : ndarray
|
||||
A matrix of M row-vectors (query points).
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
A vector of length M that contains for each entry in `y` the
|
||||
smallest Euclidean distance to a sample in `x`.
|
||||
|
||||
"""
|
||||
distances = _pdist(x, y)
|
||||
return np.maximum(0.0, distances.min(axis=0))
|
||||
|
||||
|
||||
def _nn_cosine_distance(x, y):
|
||||
""" Helper function for nearest neighbor distance metric (cosine).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : ndarray
|
||||
A matrix of N row-vectors (sample points).
|
||||
y : ndarray
|
||||
A matrix of M row-vectors (query points).
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
A vector of length M that contains for each entry in `y` the
|
||||
smallest cosine distance to a sample in `x`.
|
||||
|
||||
"""
|
||||
distances = _cosine_distance(x, y)
|
||||
return distances.min(axis=0)
|
||||
|
||||
|
||||
class NearestNeighborDistanceMetric(object):
|
||||
"""
|
||||
A nearest neighbor distance metric that, for each target, returns
|
||||
the closest distance to any sample that has been observed so far.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
metric : str
|
||||
Either "euclidean" or "cosine".
|
||||
matching_threshold: float
|
||||
The matching threshold. Samples with larger distance are considered an
|
||||
invalid match.
|
||||
budget : Optional[int]
|
||||
If not None, fix samples per class to at most this number. Removes
|
||||
the oldest samples when the budget is reached.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
samples : Dict[int -> List[ndarray]]
|
||||
A dictionary that maps from target identities to the list of samples
|
||||
that have been observed so far.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, metric, matching_threshold, budget=None):
|
||||
|
||||
|
||||
if metric == "euclidean":
|
||||
self._metric = _nn_euclidean_distance
|
||||
elif metric == "cosine":
|
||||
self._metric = _nn_cosine_distance
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid metric; must be either 'euclidean' or 'cosine'")
|
||||
self.matching_threshold = matching_threshold
|
||||
self.budget = budget
|
||||
self.samples = {}
|
||||
|
||||
def partial_fit(self, features, targets, active_targets):
|
||||
"""Update the distance metric with new data.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
features : ndarray
|
||||
An NxM matrix of N features of dimensionality M.
|
||||
targets : ndarray
|
||||
An integer array of associated target identities.
|
||||
active_targets : List[int]
|
||||
A list of targets that are currently present in the scene.
|
||||
|
||||
"""
|
||||
for feature, target in zip(features, targets):
|
||||
self.samples.setdefault(target, []).append(feature)
|
||||
if self.budget is not None:
|
||||
self.samples[target] = self.samples[target][-self.budget:]
|
||||
self.samples = {k: self.samples[k] for k in active_targets}
|
||||
|
||||
def distance(self, features, targets):
|
||||
"""Compute distance between features and targets.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
features : ndarray
|
||||
An NxM matrix of N features of dimensionality M.
|
||||
targets : List[int]
|
||||
A list of targets to match the given `features` against.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
Returns a cost matrix of shape len(targets), len(features), where
|
||||
element (i, j) contains the closest squared distance between
|
||||
`targets[i]` and `features[j]`.
|
||||
|
||||
"""
|
||||
cost_matrix = np.zeros((len(targets), len(features)))
|
||||
for i, target in enumerate(targets):
|
||||
cost_matrix[i, :] = self._metric(self.samples[target], features)
|
||||
return cost_matrix
|
@ -0,0 +1,73 @@
|
||||
# vim: expandtab:ts=4:sw=4
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
||||
def non_max_suppression(boxes, max_bbox_overlap, scores=None):
|
||||
"""Suppress overlapping detections.
|
||||
|
||||
Original code from [1]_ has been adapted to include confidence score.
|
||||
|
||||
.. [1] http://www.pyimagesearch.com/2015/02/16/
|
||||
faster-non-maximum-suppression-python/
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
>>> boxes = [d.roi for d in detections]
|
||||
>>> scores = [d.confidence for d in detections]
|
||||
>>> indices = non_max_suppression(boxes, max_bbox_overlap, scores)
|
||||
>>> detections = [detections[i] for i in indices]
|
||||
|
||||
Parameters
|
||||
----------
|
||||
boxes : ndarray
|
||||
Array of ROIs (x, y, width, height).
|
||||
max_bbox_overlap : float
|
||||
ROIs that overlap more than this values are suppressed.
|
||||
scores : Optional[array_like]
|
||||
Detector confidence score.
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[int]
|
||||
Returns indices of detections that have survived non-maxima suppression.
|
||||
|
||||
"""
|
||||
if len(boxes) == 0:
|
||||
return []
|
||||
|
||||
boxes = boxes.astype(np.float)
|
||||
pick = []
|
||||
|
||||
x1 = boxes[:, 0]
|
||||
y1 = boxes[:, 1]
|
||||
x2 = boxes[:, 2] + boxes[:, 0]
|
||||
y2 = boxes[:, 3] + boxes[:, 1]
|
||||
|
||||
area = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
if scores is not None:
|
||||
idxs = np.argsort(scores)
|
||||
else:
|
||||
idxs = np.argsort(y2)
|
||||
|
||||
while len(idxs) > 0:
|
||||
last = len(idxs) - 1
|
||||
i = idxs[last]
|
||||
pick.append(i)
|
||||
|
||||
xx1 = np.maximum(x1[i], x1[idxs[:last]])
|
||||
yy1 = np.maximum(y1[i], y1[idxs[:last]])
|
||||
xx2 = np.minimum(x2[i], x2[idxs[:last]])
|
||||
yy2 = np.minimum(y2[i], y2[idxs[:last]])
|
||||
|
||||
w = np.maximum(0, xx2 - xx1 + 1)
|
||||
h = np.maximum(0, yy2 - yy1 + 1)
|
||||
|
||||
overlap = (w * h) / area[idxs[:last]]
|
||||
|
||||
idxs = np.delete(
|
||||
idxs, np.concatenate(
|
||||
([last], np.where(overlap > max_bbox_overlap)[0])))
|
||||
|
||||
return pick
|
@ -0,0 +1,167 @@
|
||||
# vim: expandtab:ts=4:sw=4
|
||||
|
||||
|
||||
class TrackState:
|
||||
"""
|
||||
Enumeration type for the single target track state. Newly created tracks are
|
||||
classified as `tentative` until enough evidence has been collected. Then,
|
||||
the track state is changed to `confirmed`. Tracks that are no longer alive
|
||||
are classified as `deleted` to mark them for removal from the set of active
|
||||
tracks.
|
||||
|
||||
"""
|
||||
|
||||
Tentative = 1
|
||||
Confirmed = 2
|
||||
Deleted = 3
|
||||
|
||||
|
||||
class Track:
|
||||
"""
|
||||
A single target track with state space `(x, y, a, h)` and associated
|
||||
velocities, where `(x, y)` is the center of the bounding box, `a` is the
|
||||
aspect ratio and `h` is the height.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mean : ndarray
|
||||
Mean vector of the initial state distribution.
|
||||
covariance : ndarray
|
||||
Covariance matrix of the initial state distribution.
|
||||
track_id : int
|
||||
A unique track identifier.
|
||||
n_init : int
|
||||
Number of consecutive detections before the track is confirmed. The
|
||||
track state is set to `Deleted` if a miss occurs within the first
|
||||
`n_init` frames.
|
||||
max_age : int
|
||||
The maximum number of consecutive misses before the track state is
|
||||
set to `Deleted`.
|
||||
feature : Optional[ndarray]
|
||||
Feature vector of the detection this track originates from. If not None,
|
||||
this feature is added to the `features` cache.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
mean : ndarray
|
||||
Mean vector of the initial state distribution.
|
||||
covariance : ndarray
|
||||
Covariance matrix of the initial state distribution.
|
||||
track_id : int
|
||||
A unique track identifier.
|
||||
hits : int
|
||||
Total number of measurement updates.
|
||||
age : int
|
||||
Total number of frames since first occurance.
|
||||
time_since_update : int
|
||||
Total number of frames since last measurement update.
|
||||
state : TrackState
|
||||
The current track state.
|
||||
features : List[ndarray]
|
||||
A cache of features. On each measurement update, the associated feature
|
||||
vector is added to this list.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, mean, covariance, track_id, n_init, max_age,
|
||||
feature=None):
|
||||
#
|
||||
self.mean = mean
|
||||
self.covariance = covariance
|
||||
self.track_id = track_id
|
||||
self.hits = 1
|
||||
self.age = 1
|
||||
self.time_since_update = 0
|
||||
|
||||
self.state = TrackState.Tentative
|
||||
self.features = []
|
||||
if feature is not None:
|
||||
self.features.append(feature)
|
||||
|
||||
self._n_init = n_init
|
||||
self._max_age = max_age
|
||||
|
||||
def to_tlwh(self):
|
||||
"""Get current position in bounding box format `(top left x, top left y,
|
||||
width, height)`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
The bounding box.
|
||||
|
||||
"""
|
||||
ret = self.mean[:4].copy() # xc,yc, a, h
|
||||
ret[2] *= ret[3]
|
||||
ret[:2] -= ret[2:] / 2
|
||||
return ret
|
||||
|
||||
def to_tlbr(self):
|
||||
"""Get current position in bounding box format `(min x, miny, max x,
|
||||
max y)`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ndarray
|
||||
The bounding box.
|
||||
|
||||
"""
|
||||
ret = self.to_tlwh()
|
||||
ret[2:] = ret[:2] + ret[2:]
|
||||
return ret
|
||||
|
||||
def predict(self, kf):
|
||||
"""Propagate the state distribution to the current time step using a
|
||||
Kalman filter prediction step.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
kf : kalman_filter.KalmanFilter
|
||||
The Kalman filter.
|
||||
|
||||
"""
|
||||
self.mean, self.covariance = kf.predict(self.mean, self.covariance)
|
||||
self.age += 1
|
||||
self.time_since_update += 1
|
||||
|
||||
def update(self, kf, detection):
|
||||
"""Perform Kalman filter measurement update step and update the feature
|
||||
cache.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
kf : kalman_filter.KalmanFilter
|
||||
The Kalman filter.
|
||||
detection : Detection
|
||||
The associated detection.
|
||||
|
||||
"""
|
||||
self.mean, self.covariance = kf.update(
|
||||
self.mean, self.covariance, detection.to_xyah())
|
||||
self.features.append(detection.feature)
|
||||
|
||||
self.hits += 1
|
||||
self.time_since_update = 0
|
||||
if self.state == TrackState.Tentative and self.hits >= self._n_init:
|
||||
self.state = TrackState.Confirmed
|
||||
|
||||
def mark_missed(self):
|
||||
"""Mark this track as missed (no association at the current time step).
|
||||
"""
|
||||
if self.state == TrackState.Tentative:
|
||||
self.state = TrackState.Deleted
|
||||
elif self.time_since_update > self._max_age:
|
||||
self.state = TrackState.Deleted
|
||||
|
||||
def is_tentative(self):
|
||||
"""Returns True if this track is tentative (unconfirmed).
|
||||
"""
|
||||
return self.state == TrackState.Tentative
|
||||
|
||||
def is_confirmed(self):
|
||||
"""Returns True if this track is confirmed."""
|
||||
return self.state == TrackState.Confirmed
|
||||
|
||||
def is_deleted(self):
|
||||
"""Returns True if this track is dead and should be deleted."""
|
||||
return self.state == TrackState.Deleted
|
@ -0,0 +1,263 @@
|
||||
from yolov5.utils.general import (
|
||||
check_img_size, non_max_suppression, scale_coords, xyxy2xywh)
|
||||
from yolov5.utils.torch_utils import select_device, time_synchronized
|
||||
from yolov5.utils.datasets import letterbox
|
||||
|
||||
from utils_ds.parser import get_config
|
||||
from utils_ds.draw import draw_boxes
|
||||
from deep_sort import build_tracker
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
import numpy as np
|
||||
import warnings
|
||||
import cv2
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
|
||||
import sys
|
||||
|
||||
currentUrl = os.path.dirname(__file__)
|
||||
sys.path.append(os.path.abspath(os.path.join(currentUrl, 'yolov5')))
|
||||
|
||||
|
||||
cudnn.benchmark = True
|
||||
|
||||
|
||||
class VideoTracker(object):
|
||||
def __init__(self, args):
|
||||
print('Initialize DeepSORT & YOLO-V5')
|
||||
# ***************** Initialize ******************************************************
|
||||
self.args = args
|
||||
|
||||
self.img_size = args.img_size # image size in detector, default is 640
|
||||
self.frame_interval = args.frame_interval # frequency
|
||||
|
||||
self.device = select_device(args.device)
|
||||
self.half = self.device.type != 'cpu' # half precision only supported on CUDA
|
||||
|
||||
# create video capture ****************
|
||||
if args.display:
|
||||
cv2.namedWindow("test", cv2.WINDOW_NORMAL)
|
||||
cv2.resizeWindow("test", args.display_width, args.display_height)
|
||||
|
||||
if args.cam != -1:
|
||||
print("Using webcam " + str(args.cam))
|
||||
self.vdo = cv2.VideoCapture(args.cam)
|
||||
if not self.vdo.isOpened():
|
||||
raise ValueError(f"Error opening camera {args.cam}")
|
||||
else:
|
||||
self.vdo = cv2.VideoCapture()
|
||||
|
||||
# ***************************** initialize DeepSORT **********************************
|
||||
cfg = get_config()
|
||||
cfg.merge_from_file(args.config_deepsort)
|
||||
|
||||
use_cuda = self.device.type != 'cpu' and torch.cuda.is_available()
|
||||
self.deepsort = build_tracker(cfg, use_cuda=use_cuda)
|
||||
|
||||
# ***************************** initialize YOLO-V5 **********************************
|
||||
self.detector = torch.load(args.weights, map_location=self.device)['model'].float() # load to FP32
|
||||
self.detector.to(self.device).eval()
|
||||
if self.half:
|
||||
self.detector.half() # to FP16
|
||||
|
||||
self.names = self.detector.module.names if hasattr(self.detector, 'module') else self.detector.names
|
||||
|
||||
print('Done..')
|
||||
if self.device == 'cpu':
|
||||
warnings.warn("Running in cpu mode which maybe very slow!", UserWarning)
|
||||
|
||||
def __enter__(self):
|
||||
# ************************* Load video from camera *************************
|
||||
if self.args.cam != -1:
|
||||
print('Camera ...')
|
||||
ret, frame = self.vdo.read()
|
||||
assert ret, "Error: Camera error"
|
||||
self.im_width = int(self.vdo.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
self.im_height = int(self.vdo.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
|
||||
# ************************* Load video from file *************************
|
||||
else:
|
||||
assert os.path.isfile(self.args.input_path), "Path error"
|
||||
self.vdo.open(self.args.input_path)
|
||||
self.im_width = int(self.vdo.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
self.im_height = int(self.vdo.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
assert self.vdo.isOpened()
|
||||
print('Done. Load video file ', self.args.input_path)
|
||||
|
||||
# ************************* create output *************************
|
||||
if self.args.save_path:
|
||||
os.makedirs(self.args.save_path, exist_ok=True)
|
||||
# path of saved video and results
|
||||
self.save_video_path = os.path.join(self.args.save_path, "results.mp4")
|
||||
|
||||
# create video writer
|
||||
fourcc = cv2.VideoWriter_fourcc(*self.args.fourcc)
|
||||
self.writer = cv2.VideoWriter(self.save_video_path, fourcc,
|
||||
self.vdo.get(cv2.CAP_PROP_FPS), (self.im_width, self.im_height))
|
||||
print('Done. Create output file ', self.save_video_path)
|
||||
|
||||
if self.args.save_txt:
|
||||
os.makedirs(self.args.save_txt, exist_ok=True)
|
||||
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, exc_traceback):
|
||||
self.vdo.release()
|
||||
self.writer.release()
|
||||
if exc_type:
|
||||
print(exc_type, exc_value, exc_traceback)
|
||||
|
||||
def run(self):
|
||||
yolo_time, sort_time, avg_fps = [], [], []
|
||||
t_start = time.time()
|
||||
|
||||
idx_frame = 0
|
||||
last_out = None
|
||||
while self.vdo.grab():
|
||||
# Inference *********************************************************************
|
||||
t0 = time.time()
|
||||
_, img0 = self.vdo.retrieve()
|
||||
|
||||
if idx_frame % self.args.frame_interval == 0:
|
||||
outputs, yt, st = self.image_track(img0) # (#ID, 5) x1,y1,x2,y2,id
|
||||
last_out = outputs
|
||||
yolo_time.append(yt)
|
||||
sort_time.append(st)
|
||||
print('Frame %d Done. YOLO-time:(%.3fs) SORT-time:(%.3fs)' % (idx_frame, yt, st))
|
||||
else:
|
||||
outputs = last_out # directly use prediction in last frames
|
||||
t1 = time.time()
|
||||
avg_fps.append(t1 - t0)
|
||||
|
||||
# post-processing ***************************************************************
|
||||
# visualize bbox ********************************
|
||||
if len(outputs) > 0:
|
||||
bbox_xyxy = outputs[:, :4]
|
||||
identities = outputs[:, -1]
|
||||
img0 = draw_boxes(img0, bbox_xyxy, identities) # BGR
|
||||
|
||||
# add FPS information on output video
|
||||
text_scale = max(1, img0.shape[1] // 1600)
|
||||
cv2.putText(img0, 'frame: %d fps: %.2f ' % (idx_frame, len(avg_fps) / sum(avg_fps)),
|
||||
(20, 20 + text_scale), cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 255), thickness=2)
|
||||
|
||||
# display on window ******************************
|
||||
if self.args.display:
|
||||
cv2.imshow("test", img0)
|
||||
if cv2.waitKey(1) == ord('q'): # q to quit
|
||||
cv2.destroyAllWindows()
|
||||
break
|
||||
|
||||
# save to video file *****************************
|
||||
if self.args.save_path:
|
||||
self.writer.write(img0)
|
||||
|
||||
if self.args.save_txt:
|
||||
with open(self.args.save_txt + str(idx_frame).zfill(4) + '.txt', 'a') as f:
|
||||
for i in range(len(outputs)):
|
||||
x1, y1, x2, y2, idx = outputs[i]
|
||||
f.write('{}\t{}\t{}\t{}\t{}\n'.format(x1, y1, x2, y2, idx))
|
||||
|
||||
|
||||
|
||||
idx_frame += 1
|
||||
|
||||
print('Avg YOLO time (%.3fs), Sort time (%.3fs) per frame' % (sum(yolo_time) / len(yolo_time),
|
||||
sum(sort_time)/len(sort_time)))
|
||||
t_end = time.time()
|
||||
print('Total time (%.3fs), Total Frame: %d' % (t_end - t_start, idx_frame))
|
||||
|
||||
def image_track(self, im0):
|
||||
"""
|
||||
:param im0: original image, BGR format
|
||||
:return:
|
||||
"""
|
||||
# preprocess ************************************************************
|
||||
# Padded resize
|
||||
img = letterbox(im0, new_shape=self.img_size)[0]
|
||||
# Convert
|
||||
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
||||
img = np.ascontiguousarray(img)
|
||||
|
||||
# numpy to tensor
|
||||
img = torch.from_numpy(img).to(self.device)
|
||||
img = img.half() if self.half else img.float() # uint8 to fp16/32
|
||||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||||
if img.ndimension() == 3:
|
||||
img = img.unsqueeze(0)
|
||||
s = '%gx%g ' % img.shape[2:] # print string
|
||||
|
||||
# Detection time *********************************************************
|
||||
# Inference
|
||||
t1 = time_synchronized()
|
||||
with torch.no_grad():
|
||||
pred = self.detector(img, augment=self.args.augment)[0] # list: bz * [ (#obj, 6)]
|
||||
|
||||
# Apply NMS and filter object other than person (cls:0)
|
||||
pred = non_max_suppression(pred, self.args.conf_thres, self.args.iou_thres,
|
||||
classes=self.args.classes, agnostic=self.args.agnostic_nms)
|
||||
t2 = time_synchronized()
|
||||
|
||||
# get all obj ************************************************************
|
||||
det = pred[0] # for video, bz is 1
|
||||
if det is not None and len(det): # det: (#obj, 6) x1 y1 x2 y2 conf cls
|
||||
|
||||
# Rescale boxes from img_size to original im0 size
|
||||
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
|
||||
|
||||
# Print results. statistics of number of each obj
|
||||
for c in det[:, -1].unique():
|
||||
n = (det[:, -1] == c).sum() # detections per class
|
||||
s += '%g %ss, ' % (n, self.names[int(c)]) # add to string
|
||||
|
||||
bbox_xywh = xyxy2xywh(det[:, :4]).cpu()
|
||||
confs = det[:, 4:5].cpu()
|
||||
|
||||
# ****************************** deepsort ****************************
|
||||
outputs = self.deepsort.update(bbox_xywh, confs, im0)
|
||||
# (#ID, 5) x1,y1,x2,y2,track_ID
|
||||
else:
|
||||
outputs = torch.zeros((0, 5))
|
||||
|
||||
t3 = time.time()
|
||||
return outputs, t2-t1, t3-t2
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
# input and output
|
||||
parser.add_argument('--input_path', type=str, default='input_480.mp4', help='source') # file/folder, 0 for webcam
|
||||
parser.add_argument('--save_path', type=str, default='output/', help='output folder') # output folder
|
||||
parser.add_argument("--frame_interval", type=int, default=2)
|
||||
parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--save_txt', default='output/predict/', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
|
||||
# camera only
|
||||
parser.add_argument("--display", action="store_true")
|
||||
parser.add_argument("--display_width", type=int, default=800)
|
||||
parser.add_argument("--display_height", type=int, default=600)
|
||||
parser.add_argument("--camera", action="store", dest="cam", type=int, default="-1")
|
||||
|
||||
# YOLO-V5 parameters
|
||||
parser.add_argument('--weights', type=str, default='yolov5/weights/yolov5s.pt', help='model.pt path')
|
||||
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
||||
parser.add_argument('--conf-thres', type=float, default=0.5, help='object confidence threshold')
|
||||
parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
|
||||
parser.add_argument('--classes', nargs='+', type=int, default=[0], help='filter by class')
|
||||
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
|
||||
parser.add_argument('--augment', action='store_true', help='augmented inference')
|
||||
|
||||
# deepsort parameters
|
||||
parser.add_argument("--config_deepsort", type=str, default="./configs/deep_sort.yaml")
|
||||
|
||||
args = parser.parse_args()
|
||||
args.img_size = check_img_size(args.img_size)
|
||||
print(args)
|
||||
|
||||
with VideoTracker(args) as vdo_trk:
|
||||
vdo_trk.run()
|
||||
|
@ -0,0 +1,14 @@
|
||||
# pip install -U -r requirements.txt
|
||||
Cython
|
||||
matplotlib>=3.2.2
|
||||
numpy>=1.18.5
|
||||
opencv-python>=4.1.2
|
||||
pillow
|
||||
easydict
|
||||
# pycocotools>=2.0
|
||||
PyYAML>=5.3
|
||||
scipy>=1.4.1
|
||||
tensorboard>=2.2
|
||||
#torch>=1.6.0
|
||||
#torchvision>=0.7.0
|
||||
tqdm>=4.41.0
|
@ -0,0 +1,13 @@
|
||||
from os import environ
|
||||
|
||||
|
||||
def assert_in(file, files_to_check):
|
||||
if file not in files_to_check:
|
||||
raise AssertionError("{} does not exist in the list".format(str(file)))
|
||||
return True
|
||||
|
||||
|
||||
def assert_in_env(check_list: list):
|
||||
for item in check_list:
|
||||
assert_in(item, environ.keys())
|
||||
return True
|
@ -0,0 +1,36 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
|
||||
|
||||
|
||||
def compute_color_for_labels(label):
|
||||
"""
|
||||
Simple function that adds fixed color depending on the class
|
||||
"""
|
||||
color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette]
|
||||
return tuple(color)
|
||||
|
||||
|
||||
def draw_boxes(img, bbox, identities=None, offset=(0,0)):
|
||||
for i,box in enumerate(bbox):
|
||||
x1,y1,x2,y2 = [int(i) for i in box]
|
||||
x1 += offset[0]
|
||||
x2 += offset[0]
|
||||
y1 += offset[1]
|
||||
y2 += offset[1]
|
||||
# box text and bar
|
||||
id = int(identities[i]) if identities is not None else 0
|
||||
color = compute_color_for_labels(id)
|
||||
label = '{}{:d}'.format("", id)
|
||||
t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2 , 2)[0]
|
||||
cv2.rectangle(img,(x1, y1),(x2,y2),color,3)
|
||||
cv2.rectangle(img,(x1, y1),(x1+t_size[0]+3,y1+t_size[1]+4), color,-1)
|
||||
cv2.putText(img,label,(x1,y1+t_size[1]+4), cv2.FONT_HERSHEY_PLAIN, 2, [255,255,255], 2)
|
||||
return img
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
for i in range(82):
|
||||
print(compute_color_for_labels(i))
|
@ -0,0 +1,103 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import copy
|
||||
import motmetrics as mm
|
||||
mm.lap.default_solver = 'lap'
|
||||
from utils.io import read_results, unzip_objs
|
||||
|
||||
|
||||
class Evaluator(object):
|
||||
|
||||
def __init__(self, data_root, seq_name, data_type):
|
||||
self.data_root = data_root
|
||||
self.seq_name = seq_name
|
||||
self.data_type = data_type
|
||||
|
||||
self.load_annotations()
|
||||
self.reset_accumulator()
|
||||
|
||||
def load_annotations(self):
|
||||
assert self.data_type == 'mot'
|
||||
|
||||
gt_filename = os.path.join(self.data_root, self.seq_name, 'gt', 'gt.txt')
|
||||
self.gt_frame_dict = read_results(gt_filename, self.data_type, is_gt=True)
|
||||
self.gt_ignore_frame_dict = read_results(gt_filename, self.data_type, is_ignore=True)
|
||||
|
||||
def reset_accumulator(self):
|
||||
self.acc = mm.MOTAccumulator(auto_id=True)
|
||||
|
||||
def eval_frame(self, frame_id, trk_tlwhs, trk_ids, rtn_events=False):
|
||||
# results
|
||||
trk_tlwhs = np.copy(trk_tlwhs)
|
||||
trk_ids = np.copy(trk_ids)
|
||||
|
||||
# gts
|
||||
gt_objs = self.gt_frame_dict.get(frame_id, [])
|
||||
gt_tlwhs, gt_ids = unzip_objs(gt_objs)[:2]
|
||||
|
||||
# ignore boxes
|
||||
ignore_objs = self.gt_ignore_frame_dict.get(frame_id, [])
|
||||
ignore_tlwhs = unzip_objs(ignore_objs)[0]
|
||||
|
||||
|
||||
# remove ignored results
|
||||
keep = np.ones(len(trk_tlwhs), dtype=bool)
|
||||
iou_distance = mm.distances.iou_matrix(ignore_tlwhs, trk_tlwhs, max_iou=0.5)
|
||||
if len(iou_distance) > 0:
|
||||
match_is, match_js = mm.lap.linear_sum_assignment(iou_distance)
|
||||
match_is, match_js = map(lambda a: np.asarray(a, dtype=int), [match_is, match_js])
|
||||
match_ious = iou_distance[match_is, match_js]
|
||||
|
||||
match_js = np.asarray(match_js, dtype=int)
|
||||
match_js = match_js[np.logical_not(np.isnan(match_ious))]
|
||||
keep[match_js] = False
|
||||
trk_tlwhs = trk_tlwhs[keep]
|
||||
trk_ids = trk_ids[keep]
|
||||
|
||||
# get distance matrix
|
||||
iou_distance = mm.distances.iou_matrix(gt_tlwhs, trk_tlwhs, max_iou=0.5)
|
||||
|
||||
# acc
|
||||
self.acc.update(gt_ids, trk_ids, iou_distance)
|
||||
|
||||
if rtn_events and iou_distance.size > 0 and hasattr(self.acc, 'last_mot_events'):
|
||||
events = self.acc.last_mot_events # only supported by https://github.com/longcw/py-motmetrics
|
||||
else:
|
||||
events = None
|
||||
return events
|
||||
|
||||
def eval_file(self, filename):
|
||||
self.reset_accumulator()
|
||||
|
||||
result_frame_dict = read_results(filename, self.data_type, is_gt=False)
|
||||
frames = sorted(list(set(self.gt_frame_dict.keys()) | set(result_frame_dict.keys())))
|
||||
for frame_id in frames:
|
||||
trk_objs = result_frame_dict.get(frame_id, [])
|
||||
trk_tlwhs, trk_ids = unzip_objs(trk_objs)[:2]
|
||||
self.eval_frame(frame_id, trk_tlwhs, trk_ids, rtn_events=False)
|
||||
|
||||
return self.acc
|
||||
|
||||
@staticmethod
|
||||
def get_summary(accs, names, metrics=('mota', 'num_switches', 'idp', 'idr', 'idf1', 'precision', 'recall')):
|
||||
names = copy.deepcopy(names)
|
||||
if metrics is None:
|
||||
metrics = mm.metrics.motchallenge_metrics
|
||||
metrics = copy.deepcopy(metrics)
|
||||
|
||||
mh = mm.metrics.create()
|
||||
summary = mh.compute_many(
|
||||
accs,
|
||||
metrics=metrics,
|
||||
names=names,
|
||||
generate_overall=True
|
||||
)
|
||||
|
||||
return summary
|
||||
|
||||
@staticmethod
|
||||
def save_summary(summary, filename):
|
||||
import pandas as pd
|
||||
writer = pd.ExcelWriter(filename)
|
||||
summary.to_excel(writer)
|
||||
writer.save()
|
@ -0,0 +1,133 @@
|
||||
import os
|
||||
from typing import Dict
|
||||
import numpy as np
|
||||
|
||||
# from utils.log import get_logger
|
||||
|
||||
|
||||
def write_results(filename, results, data_type):
|
||||
if data_type == 'mot':
|
||||
save_format = '{frame},{id},{x1},{y1},{w},{h},-1,-1,-1,-1\n'
|
||||
elif data_type == 'kitti':
|
||||
save_format = '{frame} {id} pedestrian 0 0 -10 {x1} {y1} {x2} {y2} -10 -10 -10 -1000 -1000 -1000 -10\n'
|
||||
else:
|
||||
raise ValueError(data_type)
|
||||
|
||||
with open(filename, 'w') as f:
|
||||
for frame_id, tlwhs, track_ids in results:
|
||||
if data_type == 'kitti':
|
||||
frame_id -= 1
|
||||
for tlwh, track_id in zip(tlwhs, track_ids):
|
||||
if track_id < 0:
|
||||
continue
|
||||
x1, y1, w, h = tlwh
|
||||
x2, y2 = x1 + w, y1 + h
|
||||
line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h)
|
||||
f.write(line)
|
||||
|
||||
|
||||
# def write_results(filename, results_dict: Dict, data_type: str):
|
||||
# if not filename:
|
||||
# return
|
||||
# path = os.path.dirname(filename)
|
||||
# if not os.path.exists(path):
|
||||
# os.makedirs(path)
|
||||
|
||||
# if data_type in ('mot', 'mcmot', 'lab'):
|
||||
# save_format = '{frame},{id},{x1},{y1},{w},{h},1,-1,-1,-1\n'
|
||||
# elif data_type == 'kitti':
|
||||
# save_format = '{frame} {id} pedestrian -1 -1 -10 {x1} {y1} {x2} {y2} -1 -1 -1 -1000 -1000 -1000 -10 {score}\n'
|
||||
# else:
|
||||
# raise ValueError(data_type)
|
||||
|
||||
# with open(filename, 'w') as f:
|
||||
# for frame_id, frame_data in results_dict.items():
|
||||
# if data_type == 'kitti':
|
||||
# frame_id -= 1
|
||||
# for tlwh, track_id in frame_data:
|
||||
# if track_id < 0:
|
||||
# continue
|
||||
# x1, y1, w, h = tlwh
|
||||
# x2, y2 = x1 + w, y1 + h
|
||||
# line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h, score=1.0)
|
||||
# f.write(line)
|
||||
# logger.info('Save results to {}'.format(filename))
|
||||
|
||||
|
||||
def read_results(filename, data_type: str, is_gt=False, is_ignore=False):
|
||||
if data_type in ('mot', 'lab'):
|
||||
read_fun = read_mot_results
|
||||
else:
|
||||
raise ValueError('Unknown data type: {}'.format(data_type))
|
||||
|
||||
return read_fun(filename, is_gt, is_ignore)
|
||||
|
||||
|
||||
"""
|
||||
labels={'ped', ... % 1
|
||||
'person_on_vhcl', ... % 2
|
||||
'car', ... % 3
|
||||
'bicycle', ... % 4
|
||||
'mbike', ... % 5
|
||||
'non_mot_vhcl', ... % 6
|
||||
'static_person', ... % 7
|
||||
'distractor', ... % 8
|
||||
'occluder', ... % 9
|
||||
'occluder_on_grnd', ... %10
|
||||
'occluder_full', ... % 11
|
||||
'reflection', ... % 12
|
||||
'crowd' ... % 13
|
||||
};
|
||||
"""
|
||||
|
||||
|
||||
def read_mot_results(filename, is_gt, is_ignore):
|
||||
valid_labels = {1}
|
||||
ignore_labels = {2, 7, 8, 12}
|
||||
results_dict = dict()
|
||||
if os.path.isfile(filename):
|
||||
with open(filename, 'r') as f:
|
||||
for line in f.readlines():
|
||||
linelist = line.split(',')
|
||||
if len(linelist) < 7:
|
||||
continue
|
||||
fid = int(linelist[0])
|
||||
if fid < 1:
|
||||
continue
|
||||
results_dict.setdefault(fid, list())
|
||||
|
||||
if is_gt:
|
||||
if 'MOT16-' in filename or 'MOT17-' in filename:
|
||||
label = int(float(linelist[7]))
|
||||
mark = int(float(linelist[6]))
|
||||
if mark == 0 or label not in valid_labels:
|
||||
continue
|
||||
score = 1
|
||||
elif is_ignore:
|
||||
if 'MOT16-' in filename or 'MOT17-' in filename:
|
||||
label = int(float(linelist[7]))
|
||||
vis_ratio = float(linelist[8])
|
||||
if label not in ignore_labels and vis_ratio >= 0:
|
||||
continue
|
||||
else:
|
||||
continue
|
||||
score = 1
|
||||
else:
|
||||
score = float(linelist[6])
|
||||
|
||||
tlwh = tuple(map(float, linelist[2:6]))
|
||||
target_id = int(linelist[1])
|
||||
|
||||
results_dict[fid].append((tlwh, target_id, score))
|
||||
|
||||
return results_dict
|
||||
|
||||
|
||||
def unzip_objs(objs):
|
||||
if len(objs) > 0:
|
||||
tlwhs, ids, scores = zip(*objs)
|
||||
else:
|
||||
tlwhs, ids, scores = [], [], []
|
||||
tlwhs = np.asarray(tlwhs, dtype=float).reshape(-1, 4)
|
||||
|
||||
return tlwhs, ids, scores
|
@ -0,0 +1,383 @@
|
||||
"""
|
||||
References:
|
||||
https://medium.com/analytics-vidhya/creating-a-custom-logging-mechanism-for-real-time-object-detection-using-tdd-4ca2cfcd0a2f
|
||||
"""
|
||||
import json
|
||||
from os import makedirs
|
||||
from os.path import exists, join
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class JsonMeta(object):
|
||||
HOURS = 3
|
||||
MINUTES = 59
|
||||
SECONDS = 59
|
||||
PATH_TO_SAVE = 'LOGS'
|
||||
DEFAULT_FILE_NAME = 'remaining'
|
||||
|
||||
|
||||
class BaseJsonLogger(object):
|
||||
"""
|
||||
This is the base class that returns __dict__ of its own
|
||||
it also returns the dicts of objects in the attributes that are list instances
|
||||
|
||||
"""
|
||||
|
||||
def dic(self):
|
||||
# returns dicts of objects
|
||||
out = {}
|
||||
for k, v in self.__dict__.items():
|
||||
if hasattr(v, 'dic'):
|
||||
out[k] = v.dic()
|
||||
elif isinstance(v, list):
|
||||
out[k] = self.list(v)
|
||||
else:
|
||||
out[k] = v
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def list(values):
|
||||
# applies the dic method on items in the list
|
||||
return [v.dic() if hasattr(v, 'dic') else v for v in values]
|
||||
|
||||
|
||||
class Label(BaseJsonLogger):
|
||||
"""
|
||||
For each bounding box there are various categories with confidences. Label class keeps track of that information.
|
||||
"""
|
||||
|
||||
def __init__(self, category: str, confidence: float):
|
||||
self.category = category
|
||||
self.confidence = confidence
|
||||
|
||||
|
||||
class Bbox(BaseJsonLogger):
|
||||
"""
|
||||
This module stores the information for each frame and use them in JsonParser
|
||||
Attributes:
|
||||
labels (list): List of label module.
|
||||
top (int):
|
||||
left (int):
|
||||
width (int):
|
||||
height (int):
|
||||
|
||||
Args:
|
||||
bbox_id (float):
|
||||
top (int):
|
||||
left (int):
|
||||
width (int):
|
||||
height (int):
|
||||
|
||||
References:
|
||||
Check Label module for better understanding.
|
||||
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, bbox_id, top, left, width, height):
|
||||
self.labels = []
|
||||
self.bbox_id = bbox_id
|
||||
self.top = top
|
||||
self.left = left
|
||||
self.width = width
|
||||
self.height = height
|
||||
|
||||
def add_label(self, category, confidence):
|
||||
# adds category and confidence only if top_k is not exceeded.
|
||||
self.labels.append(Label(category, confidence))
|
||||
|
||||
def labels_full(self, value):
|
||||
return len(self.labels) == value
|
||||
|
||||
|
||||
class Frame(BaseJsonLogger):
|
||||
"""
|
||||
This module stores the information for each frame and use them in JsonParser
|
||||
Attributes:
|
||||
timestamp (float): The elapsed time of captured frame
|
||||
frame_id (int): The frame number of the captured video
|
||||
bboxes (list of Bbox objects): Stores the list of bbox objects.
|
||||
|
||||
References:
|
||||
Check Bbox class for better information
|
||||
|
||||
Args:
|
||||
timestamp (float):
|
||||
frame_id (int):
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, frame_id: int, timestamp: float = None):
|
||||
self.frame_id = frame_id
|
||||
self.timestamp = timestamp
|
||||
self.bboxes = []
|
||||
|
||||
def add_bbox(self, bbox_id: int, top: int, left: int, width: int, height: int):
|
||||
bboxes_ids = [bbox.bbox_id for bbox in self.bboxes]
|
||||
if bbox_id not in bboxes_ids:
|
||||
self.bboxes.append(Bbox(bbox_id, top, left, width, height))
|
||||
else:
|
||||
raise ValueError("Frame with id: {} already has a Bbox with id: {}".format(self.frame_id, bbox_id))
|
||||
|
||||
def add_label_to_bbox(self, bbox_id: int, category: str, confidence: float):
|
||||
bboxes = {bbox.id: bbox for bbox in self.bboxes}
|
||||
if bbox_id in bboxes.keys():
|
||||
res = bboxes.get(bbox_id)
|
||||
res.add_label(category, confidence)
|
||||
else:
|
||||
raise ValueError('the bbox with id: {} does not exists!'.format(bbox_id))
|
||||
|
||||
|
||||
class BboxToJsonLogger(BaseJsonLogger):
|
||||
"""
|
||||
ُ This module is designed to automate the task of logging jsons. An example json is used
|
||||
to show the contents of json file shortly
|
||||
Example:
|
||||
{
|
||||
"video_details": {
|
||||
"frame_width": 1920,
|
||||
"frame_height": 1080,
|
||||
"frame_rate": 20,
|
||||
"video_name": "/home/gpu/codes/MSD/pedestrian_2/project/public/camera1.avi"
|
||||
},
|
||||
"frames": [
|
||||
{
|
||||
"frame_id": 329,
|
||||
"timestamp": 3365.1254
|
||||
"bboxes": [
|
||||
{
|
||||
"labels": [
|
||||
{
|
||||
"category": "pedestrian",
|
||||
"confidence": 0.9
|
||||
}
|
||||
],
|
||||
"bbox_id": 0,
|
||||
"top": 1257,
|
||||
"left": 138,
|
||||
"width": 68,
|
||||
"height": 109
|
||||
}
|
||||
]
|
||||
}],
|
||||
|
||||
Attributes:
|
||||
frames (dict): It's a dictionary that maps each frame_id to json attributes.
|
||||
video_details (dict): information about video file.
|
||||
top_k_labels (int): shows the allowed number of labels
|
||||
start_time (datetime object): we use it to automate the json output by time.
|
||||
|
||||
Args:
|
||||
top_k_labels (int): shows the allowed number of labels
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, top_k_labels: int = 1):
|
||||
self.frames = {}
|
||||
self.video_details = self.video_details = dict(frame_width=None, frame_height=None, frame_rate=None,
|
||||
video_name=None)
|
||||
self.top_k_labels = top_k_labels
|
||||
self.start_time = datetime.now()
|
||||
|
||||
def set_top_k(self, value):
|
||||
self.top_k_labels = value
|
||||
|
||||
def frame_exists(self, frame_id: int) -> bool:
|
||||
"""
|
||||
Args:
|
||||
frame_id (int):
|
||||
|
||||
Returns:
|
||||
bool: true if frame_id is recognized
|
||||
"""
|
||||
return frame_id in self.frames.keys()
|
||||
|
||||
def add_frame(self, frame_id: int, timestamp: float = None) -> None:
|
||||
"""
|
||||
Args:
|
||||
frame_id (int):
|
||||
timestamp (float): opencv captured frame time property
|
||||
|
||||
Raises:
|
||||
ValueError: if frame_id would not exist in class frames attribute
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
"""
|
||||
if not self.frame_exists(frame_id):
|
||||
self.frames[frame_id] = Frame(frame_id, timestamp)
|
||||
else:
|
||||
raise ValueError("Frame id: {} already exists".format(frame_id))
|
||||
|
||||
def bbox_exists(self, frame_id: int, bbox_id: int) -> bool:
|
||||
"""
|
||||
Args:
|
||||
frame_id:
|
||||
bbox_id:
|
||||
|
||||
Returns:
|
||||
bool: if bbox exists in frame bboxes list
|
||||
"""
|
||||
bboxes = []
|
||||
if self.frame_exists(frame_id=frame_id):
|
||||
bboxes = [bbox.bbox_id for bbox in self.frames[frame_id].bboxes]
|
||||
return bbox_id in bboxes
|
||||
|
||||
def find_bbox(self, frame_id: int, bbox_id: int):
|
||||
"""
|
||||
|
||||
Args:
|
||||
frame_id:
|
||||
bbox_id:
|
||||
|
||||
Returns:
|
||||
bbox_id (int):
|
||||
|
||||
Raises:
|
||||
ValueError: if bbox_id does not exist in the bbox list of specific frame.
|
||||
"""
|
||||
if not self.bbox_exists(frame_id, bbox_id):
|
||||
raise ValueError("frame with id: {} does not contain bbox with id: {}".format(frame_id, bbox_id))
|
||||
bboxes = {bbox.bbox_id: bbox for bbox in self.frames[frame_id].bboxes}
|
||||
return bboxes.get(bbox_id)
|
||||
|
||||
def add_bbox_to_frame(self, frame_id: int, bbox_id: int, top: int, left: int, width: int, height: int) -> None:
|
||||
"""
|
||||
|
||||
Args:
|
||||
frame_id (int):
|
||||
bbox_id (int):
|
||||
top (int):
|
||||
left (int):
|
||||
width (int):
|
||||
height (int):
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
ValueError: if bbox_id already exist in frame information with frame_id
|
||||
ValueError: if frame_id does not exist in frames attribute
|
||||
"""
|
||||
if self.frame_exists(frame_id):
|
||||
frame = self.frames[frame_id]
|
||||
if not self.bbox_exists(frame_id, bbox_id):
|
||||
frame.add_bbox(bbox_id, top, left, width, height)
|
||||
else:
|
||||
raise ValueError(
|
||||
"frame with frame_id: {} already contains the bbox with id: {} ".format(frame_id, bbox_id))
|
||||
else:
|
||||
raise ValueError("frame with frame_id: {} does not exist".format(frame_id))
|
||||
|
||||
def add_label_to_bbox(self, frame_id: int, bbox_id: int, category: str, confidence: float):
|
||||
"""
|
||||
Args:
|
||||
frame_id:
|
||||
bbox_id:
|
||||
category:
|
||||
confidence: the confidence value returned from yolo detection
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
ValueError: if labels quota (top_k_labels) exceeds.
|
||||
"""
|
||||
bbox = self.find_bbox(frame_id, bbox_id)
|
||||
if not bbox.labels_full(self.top_k_labels):
|
||||
bbox.add_label(category, confidence)
|
||||
else:
|
||||
raise ValueError("labels in frame_id: {}, bbox_id: {} is fulled".format(frame_id, bbox_id))
|
||||
|
||||
def add_video_details(self, frame_width: int = None, frame_height: int = None, frame_rate: int = None,
|
||||
video_name: str = None):
|
||||
self.video_details['frame_width'] = frame_width
|
||||
self.video_details['frame_height'] = frame_height
|
||||
self.video_details['frame_rate'] = frame_rate
|
||||
self.video_details['video_name'] = video_name
|
||||
|
||||
def output(self):
|
||||
output = {'video_details': self.video_details}
|
||||
result = list(self.frames.values())
|
||||
output['frames'] = [item.dic() for item in result]
|
||||
return output
|
||||
|
||||
def json_output(self, output_name):
|
||||
"""
|
||||
Args:
|
||||
output_name:
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Notes:
|
||||
It creates the json output with `output_name` name.
|
||||
"""
|
||||
if not output_name.endswith('.json'):
|
||||
output_name += '.json'
|
||||
with open(output_name, 'w') as file:
|
||||
json.dump(self.output(), file)
|
||||
file.close()
|
||||
|
||||
def set_start(self):
|
||||
self.start_time = datetime.now()
|
||||
|
||||
def schedule_output_by_time(self, output_dir=JsonMeta.PATH_TO_SAVE, hours: int = 0, minutes: int = 0,
|
||||
seconds: int = 60) -> None:
|
||||
"""
|
||||
Notes:
|
||||
Creates folder and then periodically stores the jsons on that address.
|
||||
|
||||
Args:
|
||||
output_dir (str): the directory where output files will be stored
|
||||
hours (int):
|
||||
minutes (int):
|
||||
seconds (int):
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
"""
|
||||
end = datetime.now()
|
||||
interval = 0
|
||||
interval += abs(min([hours, JsonMeta.HOURS]) * 3600)
|
||||
interval += abs(min([minutes, JsonMeta.MINUTES]) * 60)
|
||||
interval += abs(min([seconds, JsonMeta.SECONDS]))
|
||||
diff = (end - self.start_time).seconds
|
||||
|
||||
if diff > interval:
|
||||
output_name = self.start_time.strftime('%Y-%m-%d %H-%M-%S') + '.json'
|
||||
if not exists(output_dir):
|
||||
makedirs(output_dir)
|
||||
output = join(output_dir, output_name)
|
||||
self.json_output(output_name=output)
|
||||
self.frames = {}
|
||||
self.start_time = datetime.now()
|
||||
|
||||
def schedule_output_by_frames(self, frames_quota, frame_counter, output_dir=JsonMeta.PATH_TO_SAVE):
|
||||
"""
|
||||
saves as the number of frames quota increases higher.
|
||||
:param frames_quota:
|
||||
:param frame_counter:
|
||||
:param output_dir:
|
||||
:return:
|
||||
"""
|
||||
pass
|
||||
|
||||
def flush(self, output_dir):
|
||||
"""
|
||||
Notes:
|
||||
We use this function to output jsons whenever possible.
|
||||
like the time that we exit the while loop of opencv.
|
||||
|
||||
Args:
|
||||
output_dir:
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
"""
|
||||
filename = self.start_time.strftime('%Y-%m-%d %H-%M-%S') + '-remaining.json'
|
||||
output = join(output_dir, filename)
|
||||
self.json_output(output_name=output)
|
@ -0,0 +1,17 @@
|
||||
import logging
|
||||
|
||||
|
||||
def get_logger(name='root'):
|
||||
formatter = logging.Formatter(
|
||||
# fmt='%(asctime)s [%(levelname)s]: %(filename)s(%(funcName)s:%(lineno)s) >> %(message)s')
|
||||
fmt='%(asctime)s [%(levelname)s]: %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
|
||||
|
||||
handler = logging.StreamHandler()
|
||||
handler.setFormatter(formatter)
|
||||
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.addHandler(handler)
|
||||
return logger
|
||||
|
||||
|
@ -0,0 +1,38 @@
|
||||
import os
|
||||
import yaml
|
||||
from easydict import EasyDict as edict
|
||||
|
||||
class YamlParser(edict):
|
||||
"""
|
||||
This is yaml parser based on EasyDict.
|
||||
"""
|
||||
def __init__(self, cfg_dict=None, config_file=None):
|
||||
if cfg_dict is None:
|
||||
cfg_dict = {}
|
||||
|
||||
if config_file is not None:
|
||||
assert(os.path.isfile(config_file))
|
||||
with open(config_file, 'r') as fo:
|
||||
cfg_dict.update(yaml.safe_load(fo.read()))
|
||||
|
||||
super(YamlParser, self).__init__(cfg_dict)
|
||||
|
||||
|
||||
def merge_from_file(self, config_file):
|
||||
with open(config_file, 'r') as fo:
|
||||
self.update(yaml.safe_load(fo.read()))
|
||||
|
||||
|
||||
def merge_from_dict(self, config_dict):
|
||||
self.update(config_dict)
|
||||
|
||||
|
||||
def get_config(config_file=None):
|
||||
return YamlParser(config_file=config_file)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cfg = YamlParser(config_file="../configs/yolov3.yaml")
|
||||
cfg.merge_from_file("../configs/deep_sort.yaml")
|
||||
|
||||
import ipdb; ipdb.set_trace()
|
@ -0,0 +1,39 @@
|
||||
from functools import wraps
|
||||
from time import time
|
||||
|
||||
|
||||
def is_video(ext: str):
|
||||
"""
|
||||
Returns true if ext exists in
|
||||
allowed_exts for video files.
|
||||
|
||||
Args:
|
||||
ext:
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
|
||||
allowed_exts = ('.mp4', '.webm', '.ogg', '.avi', '.wmv', '.mkv', '.3gp')
|
||||
return any((ext.endswith(x) for x in allowed_exts))
|
||||
|
||||
|
||||
def tik_tok(func):
|
||||
"""
|
||||
keep track of time for each process.
|
||||
Args:
|
||||
func:
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
@wraps(func)
|
||||
def _time_it(*args, **kwargs):
|
||||
start = time()
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
finally:
|
||||
end_ = time()
|
||||
print("time: {:.03f}s, fps: {:.03f}".format(end_ - start, 1 / (end_ - start)))
|
||||
|
||||
return _time_it
|
@ -0,0 +1,55 @@
|
||||
---
|
||||
name: "\U0001F41BBug report"
|
||||
about: Create a report to help us improve
|
||||
title: ''
|
||||
labels: bug
|
||||
assignees: ''
|
||||
|
||||
---
|
||||
|
||||
Before submitting a bug report, please be aware that your issue **must be reproducible** with all of the following, otherwise it is non-actionable, and we can not help you:
|
||||
- **Current repo**: run `git fetch && git status -uno` to check and `git pull` to update repo
|
||||
- **Common dataset**: coco.yaml or coco128.yaml
|
||||
- **Common environment**: Colab, Google Cloud, or Docker image. See https://github.com/ultralytics/yolov5#environments
|
||||
|
||||
If this is a custom dataset/training question you **must include** your `train*.jpg`, `test*.jpg` and `results.png` figures, or we can not help you. You can generate these with `utils.plot_results()`.
|
||||
|
||||
|
||||
## 🐛 Bug
|
||||
A clear and concise description of what the bug is.
|
||||
|
||||
|
||||
## To Reproduce (REQUIRED)
|
||||
|
||||
Input:
|
||||
```
|
||||
import torch
|
||||
|
||||
a = torch.tensor([5])
|
||||
c = a / 0
|
||||
```
|
||||
|
||||
Output:
|
||||
```
|
||||
Traceback (most recent call last):
|
||||
File "/Users/glennjocher/opt/anaconda3/envs/env1/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code
|
||||
exec(code_obj, self.user_global_ns, self.user_ns)
|
||||
File "<ipython-input-5-be04c762b799>", line 5, in <module>
|
||||
c = a / 0
|
||||
RuntimeError: ZeroDivisionError
|
||||
```
|
||||
|
||||
|
||||
## Expected behavior
|
||||
A clear and concise description of what you expected to happen.
|
||||
|
||||
|
||||
## Environment
|
||||
If applicable, add screenshots to help explain your problem.
|
||||
|
||||
- OS: [e.g. Ubuntu]
|
||||
- GPU [e.g. 2080 Ti]
|
||||
|
||||
|
||||
## Additional context
|
||||
Add any other context about the problem here.
|
@ -0,0 +1,27 @@
|
||||
---
|
||||
name: "\U0001F680Feature request"
|
||||
about: Suggest an idea for this project
|
||||
title: ''
|
||||
labels: enhancement
|
||||
assignees: ''
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Feature
|
||||
<!-- A clear and concise description of the feature proposal -->
|
||||
|
||||
## Motivation
|
||||
|
||||
<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
|
||||
|
||||
## Pitch
|
||||
|
||||
<!-- A clear and concise description of what you want to happen. -->
|
||||
|
||||
## Alternatives
|
||||
|
||||
<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
|
||||
|
||||
## Additional context
|
||||
|
||||
<!-- Add any other context or screenshots about the feature request here. -->
|
@ -0,0 +1,13 @@
|
||||
---
|
||||
name: "❓Question"
|
||||
about: Ask a general question
|
||||
title: ''
|
||||
labels: question
|
||||
assignees: ''
|
||||
|
||||
---
|
||||
|
||||
## ❔Question
|
||||
|
||||
|
||||
## Additional context
|
@ -0,0 +1,75 @@
|
||||
name: CI CPU testing
|
||||
|
||||
on: # https://help.github.com/en/actions/reference/events-that-trigger-workflows
|
||||
push:
|
||||
pull_request:
|
||||
schedule:
|
||||
- cron: "0 0 * * *"
|
||||
|
||||
jobs:
|
||||
cpu-tests:
|
||||
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, macos-latest, windows-latest]
|
||||
python-version: [3.8]
|
||||
model: ['yolov5s'] # models to test
|
||||
|
||||
# Timeout: https://stackoverflow.com/a/59076067/4521646
|
||||
timeout-minutes: 50
|
||||
steps:
|
||||
- uses: actions/checkout@v2
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v2
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
# Note: This uses an internal pip API and may not always work
|
||||
# https://github.com/actions/cache/blob/master/examples.md#multiple-oss-in-a-workflow
|
||||
- name: Get pip cache
|
||||
id: pip-cache
|
||||
run: |
|
||||
python -c "from pip._internal.locations import USER_CACHE_DIR; print('::set-output name=dir::' + USER_CACHE_DIR)"
|
||||
|
||||
- name: Cache pip
|
||||
uses: actions/cache@v1
|
||||
with:
|
||||
path: ${{ steps.pip-cache.outputs.dir }}
|
||||
key: ${{ runner.os }}-${{ matrix.python-version }}-pip-${{ hashFiles('requirements.txt') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-${{ matrix.python-version }}-pip-
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -qr requirements.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html
|
||||
pip install -q onnx
|
||||
python --version
|
||||
pip --version
|
||||
pip list
|
||||
shell: bash
|
||||
|
||||
- name: Download data
|
||||
run: |
|
||||
python -c "from utils.google_utils import * ; gdrive_download('1n_oKgR81BJtqk75b00eAjdv03qVCQn2f', 'coco128.zip')"
|
||||
mv ./coco128 ../
|
||||
|
||||
- name: Tests workflow
|
||||
run: |
|
||||
export PYTHONPATH="$PWD" # to run *.py. files in subdirectories
|
||||
di=cpu # inference devices # define device
|
||||
|
||||
# train
|
||||
python train.py --img 256 --batch 8 --weights weights/${{ matrix.model }}.pt --cfg models/${{ matrix.model }}.yaml --epochs 1 --device $di
|
||||
# detect
|
||||
python detect.py --weights weights/${{ matrix.model }}.pt --device $di
|
||||
python detect.py --weights runs/exp0/weights/last.pt --device $di
|
||||
# test
|
||||
python test.py --img 256 --batch 8 --weights weights/${{ matrix.model }}.pt --device $di
|
||||
python test.py --img 256 --batch 8 --weights runs/exp0/weights/last.pt --device $di
|
||||
|
||||
python models/yolo.py --cfg models/${{ matrix.model }}.yaml # inspect
|
||||
python models/export.py --img 256 --batch 1 --weights weights/${{ matrix.model }}.pt # export
|
||||
shell: bash
|
@ -0,0 +1,178 @@
|
||||
import argparse
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
from numpy import random
|
||||
|
||||
from models.experimental import attempt_load
|
||||
from utils.datasets import LoadStreams, LoadImages
|
||||
from utils.general import (
|
||||
check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box, strip_optimizer)
|
||||
from utils.torch_utils import select_device, load_classifier, time_synchronized
|
||||
|
||||
|
||||
def detect(save_img=False):
|
||||
out, source, weights, view_img, save_txt, imgsz = \
|
||||
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
|
||||
webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')
|
||||
|
||||
# Initialize
|
||||
device = select_device(opt.device)
|
||||
if os.path.exists(out):
|
||||
shutil.rmtree(out) # delete output folder
|
||||
os.makedirs(out) # make new output folder
|
||||
half = device.type != 'cpu' # half precision only supported on CUDA
|
||||
|
||||
# Load model
|
||||
model = attempt_load(weights, map_location=device) # load FP32 model
|
||||
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
|
||||
if half:
|
||||
model.half() # to FP16
|
||||
|
||||
# Second-stage classifier
|
||||
classify = False
|
||||
if classify:
|
||||
modelc = load_classifier(name='resnet101', n=2) # initialize
|
||||
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights
|
||||
modelc.to(device).eval()
|
||||
|
||||
# Set Dataloader
|
||||
vid_path, vid_writer = None, None
|
||||
if webcam:
|
||||
view_img = True
|
||||
cudnn.benchmark = True # set True to speed up constant image size inference
|
||||
dataset = LoadStreams(source, img_size=imgsz)
|
||||
else:
|
||||
save_img = True
|
||||
dataset = LoadImages(source, img_size=imgsz)
|
||||
|
||||
# Get names and colors
|
||||
names = model.module.names if hasattr(model, 'module') else model.names
|
||||
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
|
||||
|
||||
# Run inference
|
||||
t0 = time.time()
|
||||
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
|
||||
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
|
||||
for path, img, im0s, vid_cap in dataset:
|
||||
img = torch.from_numpy(img).to(device)
|
||||
img = img.half() if half else img.float() # uint8 to fp16/32
|
||||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||||
if img.ndimension() == 3:
|
||||
img = img.unsqueeze(0)
|
||||
|
||||
# detection time ********************************************************************************
|
||||
# Inference
|
||||
t1 = time_synchronized()
|
||||
pred = model(img, augment=opt.augment)[0] # list, len = batch_size
|
||||
|
||||
# Apply NMS
|
||||
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
|
||||
t2 = time_synchronized()
|
||||
# ********************************************************************************
|
||||
|
||||
# Apply Classifier
|
||||
if classify:
|
||||
pred = apply_classifier(pred, modelc, img, im0s)
|
||||
|
||||
# Process detections
|
||||
for i, det in enumerate(pred): # detections per image. i is batch num
|
||||
# det: (#obj, 6)
|
||||
|
||||
if webcam: # batch_size >= 1
|
||||
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
|
||||
else:
|
||||
p, s, im0 = path, '', im0s
|
||||
|
||||
save_path = str(Path(out) / Path(p).name)
|
||||
txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '')
|
||||
s += '%gx%g ' % img.shape[2:] # print string
|
||||
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
|
||||
|
||||
if det is not None and len(det):
|
||||
# Rescale boxes from img_size to im0 size
|
||||
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
|
||||
|
||||
# Print results
|
||||
for c in det[:, -1].unique():
|
||||
n = (det[:, -1] == c).sum() # detections per class
|
||||
s += '%g %ss, ' % (n, names[int(c)]) # add to string
|
||||
|
||||
# Write results
|
||||
for *xyxy, conf, cls in det:
|
||||
# x1,y1,x2,y2
|
||||
|
||||
if save_txt: # Write to file
|
||||
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
||||
with open(txt_path + '.txt', 'a') as f:
|
||||
f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format
|
||||
|
||||
if save_img or view_img: # Add bbox to image
|
||||
label = '%s %.2f' % (names[int(cls)], conf)
|
||||
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
|
||||
|
||||
# Print time (inference + NMS)
|
||||
print('%sDone. (%.3fs)' % (s, t2 - t1))
|
||||
|
||||
# Stream results
|
||||
if view_img:
|
||||
cv2.imshow(p, im0)
|
||||
if cv2.waitKey(1) == ord('q'): # q to quit
|
||||
raise StopIteration
|
||||
|
||||
# Save results (image with detections)
|
||||
if save_img:
|
||||
if dataset.mode == 'images':
|
||||
cv2.imwrite(save_path, im0)
|
||||
else:
|
||||
if vid_path != save_path: # new video
|
||||
vid_path = save_path
|
||||
if isinstance(vid_writer, cv2.VideoWriter):
|
||||
vid_writer.release() # release previous video writer
|
||||
|
||||
fourcc = 'mp4v' # output video codec
|
||||
fps = vid_cap.get(cv2.CAP_PROP_FPS)
|
||||
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
|
||||
vid_writer.write(im0)
|
||||
|
||||
if save_txt or save_img:
|
||||
print('Results saved to %s' % os.getcwd() + os.sep + out)
|
||||
if platform == 'darwin' and not opt.update: # MacOS
|
||||
os.system('open ' + save_path)
|
||||
|
||||
print('Done. (%.3fs)' % (time.time() - t0))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
|
||||
parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam
|
||||
parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder
|
||||
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
||||
parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
|
||||
parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--view-img', action='store_true', help='display results')
|
||||
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
||||
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
|
||||
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
|
||||
parser.add_argument('--augment', action='store_true', help='augmented inference')
|
||||
parser.add_argument('--update', action='store_true', help='update all models')
|
||||
opt = parser.parse_args()
|
||||
print(opt)
|
||||
|
||||
with torch.no_grad():
|
||||
if opt.update: # update all models (to fix SourceChangeWarning)
|
||||
for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
|
||||
detect()
|
||||
strip_optimizer(opt.weights)
|
||||
else:
|
||||
detect()
|
@ -0,0 +1,99 @@
|
||||
"""File for accessing YOLOv5 via PyTorch Hub https://pytorch.org/hub/
|
||||
|
||||
Usage:
|
||||
import torch
|
||||
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, channels=3, classes=80)
|
||||
"""
|
||||
|
||||
dependencies = ['torch', 'yaml']
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
from models.yolo import Model
|
||||
from utils.google_utils import attempt_download
|
||||
|
||||
|
||||
def create(name, pretrained, channels, classes):
|
||||
"""Creates a specified YOLOv5 model
|
||||
|
||||
Arguments:
|
||||
name (str): name of model, i.e. 'yolov5s'
|
||||
pretrained (bool): load pretrained weights into the model
|
||||
channels (int): number of input channels
|
||||
classes (int): number of model classes
|
||||
|
||||
Returns:
|
||||
pytorch model
|
||||
"""
|
||||
config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) # model.yaml path
|
||||
try:
|
||||
model = Model(config, channels, classes)
|
||||
if pretrained:
|
||||
ckpt = '%s.pt' % name # checkpoint filename
|
||||
attempt_download(ckpt) # download if not found locally
|
||||
state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32
|
||||
state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
|
||||
model.load_state_dict(state_dict, strict=False) # load
|
||||
return model
|
||||
|
||||
except Exception as e:
|
||||
help_url = 'https://github.com/ultralytics/yolov5/issues/36'
|
||||
s = 'Cache maybe be out of date, deleting cache and retrying may solve this. See %s for help.' % help_url
|
||||
raise Exception(s) from e
|
||||
|
||||
|
||||
def yolov5s(pretrained=False, channels=3, classes=80):
|
||||
"""YOLOv5-small model from https://github.com/ultralytics/yolov5
|
||||
|
||||
Arguments:
|
||||
pretrained (bool): load pretrained weights into the model, default=False
|
||||
channels (int): number of input channels, default=3
|
||||
classes (int): number of model classes, default=80
|
||||
|
||||
Returns:
|
||||
pytorch model
|
||||
"""
|
||||
return create('yolov5s', pretrained, channels, classes)
|
||||
|
||||
|
||||
def yolov5m(pretrained=False, channels=3, classes=80):
|
||||
"""YOLOv5-medium model from https://github.com/ultralytics/yolov5
|
||||
|
||||
Arguments:
|
||||
pretrained (bool): load pretrained weights into the model, default=False
|
||||
channels (int): number of input channels, default=3
|
||||
classes (int): number of model classes, default=80
|
||||
|
||||
Returns:
|
||||
pytorch model
|
||||
"""
|
||||
return create('yolov5m', pretrained, channels, classes)
|
||||
|
||||
|
||||
def yolov5l(pretrained=False, channels=3, classes=80):
|
||||
"""YOLOv5-large model from https://github.com/ultralytics/yolov5
|
||||
|
||||
Arguments:
|
||||
pretrained (bool): load pretrained weights into the model, default=False
|
||||
channels (int): number of input channels, default=3
|
||||
classes (int): number of model classes, default=80
|
||||
|
||||
Returns:
|
||||
pytorch model
|
||||
"""
|
||||
return create('yolov5l', pretrained, channels, classes)
|
||||
|
||||
|
||||
def yolov5x(pretrained=False, channels=3, classes=80):
|
||||
"""YOLOv5-xlarge model from https://github.com/ultralytics/yolov5
|
||||
|
||||
Arguments:
|
||||
pretrained (bool): load pretrained weights into the model, default=False
|
||||
channels (int): number of input channels, default=3
|
||||
classes (int): number of model classes, default=80
|
||||
|
||||
Returns:
|
||||
pytorch model
|
||||
"""
|
||||
return create('yolov5x', pretrained, channels, classes)
|
@ -0,0 +1,118 @@
|
||||
# This file contains modules common to various models
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def DWConv(c1, c2, k=1, s=1, act=True):
|
||||
# Depthwise convolution
|
||||
return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
|
||||
|
||||
|
||||
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)
|
||||
self.act = nn.LeakyReLU(0.1, inplace=True) if act else 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
|
||||
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):
|
||||
return torch.cat(x, self.d)
|
||||
|
||||
|
||||
class Flatten(nn.Module):
|
||||
# Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
return x.view(x.size(0), -1)
|
||||
|
||||
|
||||
class Classify(nn.Module):
|
||||
# Classification head, i.e. x(b,c1,20,20) to x(b,c2)
|
||||
def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
|
||||
super(Classify, self).__init__()
|
||||
self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)
|
||||
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) # to x(b,c2,1,1)
|
||||
self.flat = Flatten()
|
||||
|
||||
def forward(self, x):
|
||||
z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list
|
||||
return self.flat(self.conv(z)) # flatten to x(b,c2)
|
@ -0,0 +1,145 @@
|
||||
# This file contains experimental modules
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from models.common import Conv, DWConv
|
||||
from utils.google_utils import attempt_download
|
||||
|
||||
|
||||
class CrossConv(nn.Module):
|
||||
# Cross Convolution Downsample
|
||||
def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
|
||||
# ch_in, ch_out, kernel, stride, groups, expansion, shortcut
|
||||
super(CrossConv, self).__init__()
|
||||
c_ = int(c2 * e) # hidden channels
|
||||
self.cv1 = Conv(c1, c_, (1, k), (1, s))
|
||||
self.cv2 = Conv(c_, c2, (k, 1), (s, 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 C3(nn.Module):
|
||||
# Cross Convolution CSP
|
||||
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
||||
super(C3, 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(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) 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 Sum(nn.Module):
|
||||
# Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
|
||||
def __init__(self, n, weight=False): # n: number of inputs
|
||||
super(Sum, self).__init__()
|
||||
self.weight = weight # apply weights boolean
|
||||
self.iter = range(n - 1) # iter object
|
||||
if weight:
|
||||
self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights
|
||||
|
||||
def forward(self, x):
|
||||
y = x[0] # no weight
|
||||
if self.weight:
|
||||
w = torch.sigmoid(self.w) * 2
|
||||
for i in self.iter:
|
||||
y = y + x[i + 1] * w[i]
|
||||
else:
|
||||
for i in self.iter:
|
||||
y = y + x[i + 1]
|
||||
return y
|
||||
|
||||
|
||||
class GhostConv(nn.Module):
|
||||
# Ghost Convolution https://github.com/huawei-noah/ghostnet
|
||||
def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
|
||||
super(GhostConv, self).__init__()
|
||||
c_ = c2 // 2 # hidden channels
|
||||
self.cv1 = Conv(c1, c_, k, s, g, act)
|
||||
self.cv2 = Conv(c_, c_, 5, 1, c_, act)
|
||||
|
||||
def forward(self, x):
|
||||
y = self.cv1(x)
|
||||
return torch.cat([y, self.cv2(y)], 1)
|
||||
|
||||
|
||||
class GhostBottleneck(nn.Module):
|
||||
# Ghost Bottleneck https://github.com/huawei-noah/ghostnet
|
||||
def __init__(self, c1, c2, k, s):
|
||||
super(GhostBottleneck, self).__init__()
|
||||
c_ = c2 // 2
|
||||
self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
|
||||
DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
|
||||
GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
|
||||
self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),
|
||||
Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(x) + self.shortcut(x)
|
||||
|
||||
|
||||
class MixConv2d(nn.Module):
|
||||
# Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
|
||||
def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
|
||||
super(MixConv2d, self).__init__()
|
||||
groups = len(k)
|
||||
if equal_ch: # equal c_ per group
|
||||
i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices
|
||||
c_ = [(i == g).sum() for g in range(groups)] # intermediate channels
|
||||
else: # equal weight.numel() per group
|
||||
b = [c2] + [0] * groups
|
||||
a = np.eye(groups + 1, groups, k=-1)
|
||||
a -= np.roll(a, 1, axis=1)
|
||||
a *= np.array(k) ** 2
|
||||
a[0] = 1
|
||||
c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
|
||||
|
||||
self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
|
||||
self.bn = nn.BatchNorm2d(c2)
|
||||
self.act = nn.LeakyReLU(0.1, inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
|
||||
|
||||
|
||||
class Ensemble(nn.ModuleList):
|
||||
# Ensemble of models
|
||||
def __init__(self):
|
||||
super(Ensemble, self).__init__()
|
||||
|
||||
def forward(self, x, augment=False):
|
||||
y = []
|
||||
for module in self:
|
||||
y.append(module(x, augment)[0])
|
||||
# y = torch.stack(y).max(0)[0] # max ensemble
|
||||
# y = torch.cat(y, 1) # nms ensemble
|
||||
y = torch.stack(y).mean(0) # mean ensemble
|
||||
return y, None # inference, train output
|
||||
|
||||
|
||||
def attempt_load(weights, map_location=None):
|
||||
# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
|
||||
model = Ensemble()
|
||||
for w in weights if isinstance(weights, list) else [weights]:
|
||||
attempt_download(w)
|
||||
model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model
|
||||
|
||||
if len(model) == 1:
|
||||
return model[-1] # return model
|
||||
else:
|
||||
print('Ensemble created with %s\n' % weights)
|
||||
for k in ['names', 'stride']:
|
||||
setattr(model, k, getattr(model[-1], k))
|
||||
return model # return ensemble
|
@ -0,0 +1,74 @@
|
||||
"""Exports a YOLOv5 *.pt model to ONNX and TorchScript formats
|
||||
|
||||
Usage:
|
||||
$ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
|
||||
from utils.google_utils import attempt_download
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')
|
||||
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')
|
||||
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
||||
opt = parser.parse_args()
|
||||
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
|
||||
print(opt)
|
||||
|
||||
# Input
|
||||
img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection
|
||||
|
||||
# Load PyTorch model
|
||||
attempt_download(opt.weights)
|
||||
model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float()
|
||||
model.eval()
|
||||
model.model[-1].export = True # set Detect() layer export=True
|
||||
y = model(img) # dry run
|
||||
|
||||
# TorchScript export
|
||||
try:
|
||||
print('\nStarting TorchScript export with torch %s...' % torch.__version__)
|
||||
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
|
||||
ts = torch.jit.trace(model, img)
|
||||
ts.save(f)
|
||||
print('TorchScript export success, saved as %s' % f)
|
||||
except Exception as e:
|
||||
print('TorchScript export failure: %s' % e)
|
||||
|
||||
# ONNX export
|
||||
try:
|
||||
import onnx
|
||||
|
||||
print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
|
||||
f = opt.weights.replace('.pt', '.onnx') # filename
|
||||
model.fuse() # only for ONNX
|
||||
torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
|
||||
output_names=['classes', 'boxes'] if y is None else ['output'])
|
||||
|
||||
# Checks
|
||||
onnx_model = onnx.load(f) # load onnx model
|
||||
onnx.checker.check_model(onnx_model) # check onnx model
|
||||
print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model
|
||||
print('ONNX export success, saved as %s' % f)
|
||||
except Exception as e:
|
||||
print('ONNX export failure: %s' % e)
|
||||
|
||||
# CoreML export
|
||||
try:
|
||||
import coremltools as ct
|
||||
|
||||
print('\nStarting CoreML export with coremltools %s...' % ct.__version__)
|
||||
# convert model from torchscript and apply pixel scaling as per detect.py
|
||||
model = ct.convert(ts, inputs=[ct.ImageType(name='images', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
|
||||
f = opt.weights.replace('.pt', '.mlmodel') # filename
|
||||
model.save(f)
|
||||
print('CoreML export success, saved as %s' % f)
|
||||
except Exception as e:
|
||||
print('CoreML export failure: %s' % e)
|
||||
|
||||
# Finish
|
||||
print('\nExport complete. Visualize with https://github.com/lutzroeder/netron.')
|
@ -0,0 +1,51 @@
|
||||
# parameters
|
||||
nc: 80 # number of classes
|
||||
depth_multiple: 1.0 # model depth multiple
|
||||
width_multiple: 1.0 # layer channel multiple
|
||||
|
||||
# anchors
|
||||
anchors:
|
||||
- [10,13, 16,30, 33,23] # P3/8
|
||||
- [30,61, 62,45, 59,119] # P4/16
|
||||
- [116,90, 156,198, 373,326] # P5/32
|
||||
|
||||
# darknet53 backbone
|
||||
backbone:
|
||||
# [from, number, module, args]
|
||||
[[-1, 1, Conv, [32, 3, 1]], # 0
|
||||
[-1, 1, Conv, [64, 3, 2]], # 1-P1/2
|
||||
[-1, 1, Bottleneck, [64]],
|
||||
[-1, 1, Conv, [128, 3, 2]], # 3-P2/4
|
||||
[-1, 2, Bottleneck, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 5-P3/8
|
||||
[-1, 8, Bottleneck, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 7-P4/16
|
||||
[-1, 8, Bottleneck, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 9-P5/32
|
||||
[-1, 4, Bottleneck, [1024]], # 10
|
||||
]
|
||||
|
||||
# YOLOv3-SPP head
|
||||
head:
|
||||
[[-1, 1, Bottleneck, [1024, False]],
|
||||
[-1, 1, SPP, [512, [5, 9, 13]]],
|
||||
[-1, 1, Conv, [1024, 3, 1]],
|
||||
[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large)
|
||||
|
||||
[-2, 1, Conv, [256, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 8], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 1, Bottleneck, [512, False]],
|
||||
[-1, 1, Bottleneck, [512, False]],
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium)
|
||||
|
||||
[-2, 1, Conv, [128, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 1, Bottleneck, [256, False]],
|
||||
[-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small)
|
||||
|
||||
[[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
]
|
@ -0,0 +1,42 @@
|
||||
# parameters
|
||||
nc: 80 # number of classes
|
||||
depth_multiple: 1.0 # model depth multiple
|
||||
width_multiple: 1.0 # layer channel multiple
|
||||
|
||||
# anchors
|
||||
anchors:
|
||||
- [10,13, 16,30, 33,23] # P3/8
|
||||
- [30,61, 62,45, 59,119] # P4/16
|
||||
- [116,90, 156,198, 373,326] # P5/32
|
||||
|
||||
# YOLOv5 backbone
|
||||
backbone:
|
||||
# [from, number, module, args]
|
||||
[[-1, 1, Focus, [64, 3]], # 0-P1/2
|
||||
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
|
||||
[-1, 3, Bottleneck, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 6, BottleneckCSP, [1024]], # 9
|
||||
]
|
||||
|
||||
# YOLOv5 FPN head
|
||||
head:
|
||||
[[-1, 3, BottleneckCSP, [1024, False]], # 10 (P5/32-large)
|
||||
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 14 (P4/16-medium)
|
||||
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 3, BottleneckCSP, [256, False]], # 18 (P3/8-small)
|
||||
|
||||
[[18, 14, 10], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
]
|
@ -0,0 +1,48 @@
|
||||
# parameters
|
||||
nc: 80 # number of classes
|
||||
depth_multiple: 1.0 # model depth multiple
|
||||
width_multiple: 1.0 # layer channel multiple
|
||||
|
||||
# anchors
|
||||
anchors:
|
||||
- [116,90, 156,198, 373,326] # P5/32
|
||||
- [30,61, 62,45, 59,119] # P4/16
|
||||
- [10,13, 16,30, 33,23] # P3/8
|
||||
|
||||
# YOLOv5 backbone
|
||||
backbone:
|
||||
# [from, number, module, args]
|
||||
[[-1, 1, Focus, [64, 3]], # 0-P1/2
|
||||
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
|
||||
[-1, 3, BottleneckCSP, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 9
|
||||
]
|
||||
|
||||
# YOLOv5 PANet head
|
||||
head:
|
||||
[[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 13
|
||||
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small)
|
||||
|
||||
[-1, 1, Conv, [256, 3, 2]],
|
||||
[[-1, 14], 1, Concat, [1]], # cat head P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium)
|
||||
|
||||
[-1, 1, Conv, [512, 3, 2]],
|
||||
[[-1, 10], 1, Concat, [1]], # cat head P5
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large)
|
||||
|
||||
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P5, P4, P3)
|
||||
]
|
@ -0,0 +1,259 @@
|
||||
import argparse
|
||||
import math
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat
|
||||
from models.experimental import MixConv2d, CrossConv, C3
|
||||
from utils.general import check_anchor_order, make_divisible, check_file
|
||||
from utils.torch_utils import (
|
||||
time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, select_device)
|
||||
|
||||
|
||||
class Detect(nn.Module):
|
||||
def __init__(self, nc=80, anchors=(), ch=()): # detection layer
|
||||
super(Detect, self).__init__()
|
||||
self.stride = None # strides computed during build
|
||||
self.nc = nc # number of classes
|
||||
self.no = nc + 5 # number of outputs per anchor
|
||||
self.nl = len(anchors) # number of detection layers
|
||||
self.na = len(anchors[0]) // 2 # number of anchors
|
||||
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
|
||||
self.export = False # onnx export
|
||||
|
||||
def forward(self, x):
|
||||
# x = x.copy() # for profiling
|
||||
z = [] # inference output
|
||||
self.training |= self.export
|
||||
for i in range(self.nl):
|
||||
x[i] = self.m[i](x[i]) # conv
|
||||
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
|
||||
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
|
||||
|
||||
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()
|
||||
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
|
||||
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): # model, input channels, number of classes
|
||||
super(Model, self).__init__()
|
||||
if isinstance(cfg, dict):
|
||||
self.yaml = cfg # model dict
|
||||
else: # is *.yaml
|
||||
import yaml # for torch hub
|
||||
self.yaml_file = Path(cfg).name
|
||||
with open(cfg) as f:
|
||||
self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
||||
|
||||
# Define model
|
||||
if nc and nc != self.yaml['nc']:
|
||||
print('Overriding %s nc=%g with nc=%g' % (cfg, self.yaml['nc'], nc))
|
||||
self.yaml['nc'] = nc # override yaml value
|
||||
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out
|
||||
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
||||
|
||||
# Build strides, anchors
|
||||
m = self.model[-1] # Detect()
|
||||
if isinstance(m, Detect):
|
||||
s = 128 # 2x min stride
|
||||
m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward
|
||||
m.anchors /= m.stride.view(-1, 1, 1)
|
||||
check_anchor_order(m)
|
||||
self.stride = m.stride
|
||||
self._initialize_biases() # only run once
|
||||
# print('Strides: %s' % m.stride.tolist())
|
||||
|
||||
# Init weights, biases
|
||||
initialize_weights(self)
|
||||
self.info()
|
||||
print('')
|
||||
|
||||
def forward(self, x, augment=False, profile=False):
|
||||
if augment:
|
||||
img_size = x.shape[-2:] # height, width
|
||||
s = [1, 0.83, 0.67] # scales
|
||||
f = [None, 3, None] # flips (2-ud, 3-lr)
|
||||
y = [] # outputs
|
||||
for si, fi in zip(s, f):
|
||||
xi = scale_img(x.flip(fi) if fi else x, si)
|
||||
yi = self.forward_once(xi)[0] # forward
|
||||
# cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) # save
|
||||
yi[..., :4] /= si # de-scale
|
||||
if fi == 2:
|
||||
yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud
|
||||
elif fi == 3:
|
||||
yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr
|
||||
y.append(yi)
|
||||
return torch.cat(y, 1), None # augmented inference, train
|
||||
else:
|
||||
return self.forward_once(x, profile) # single-scale inference, train
|
||||
|
||||
def forward_once(self, x, profile=False):
|
||||
y, dt = [], [] # outputs
|
||||
for m in self.model:
|
||||
if m.f != -1: # if not from previous layer
|
||||
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
|
||||
|
||||
if profile:
|
||||
try:
|
||||
import thop
|
||||
o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS
|
||||
except:
|
||||
o = 0
|
||||
t = time_synchronized()
|
||||
for _ in range(10):
|
||||
_ = m(x)
|
||||
dt.append((time_synchronized() - t) * 100)
|
||||
print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
|
||||
|
||||
x = m(x) # run
|
||||
y.append(x if m.i in self.save else None) # save output
|
||||
|
||||
if profile:
|
||||
print('%.1fms total' % sum(dt))
|
||||
return x
|
||||
|
||||
def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
|
||||
# cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
|
||||
m = self.model[-1] # Detect() module
|
||||
for mi, s in zip(m.m, m.stride): # from
|
||||
b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85)
|
||||
b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image)
|
||||
b[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls
|
||||
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
|
||||
|
||||
def _print_biases(self):
|
||||
m = self.model[-1] # Detect() module
|
||||
for mi in m.m: # from
|
||||
b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85)
|
||||
print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
|
||||
|
||||
# def _print_weights(self):
|
||||
# for m in self.model.modules():
|
||||
# if type(m) is Bottleneck:
|
||||
# print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights
|
||||
|
||||
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
|
||||
print('Fusing layers... ', end='')
|
||||
for m in self.model.modules():
|
||||
if type(m) is Conv:
|
||||
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
|
||||
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
|
||||
m.bn = None # remove batchnorm
|
||||
m.forward = m.fuseforward # update forward
|
||||
self.info()
|
||||
return self
|
||||
|
||||
def info(self): # print model information
|
||||
model_info(self)
|
||||
|
||||
|
||||
def parse_model(d, ch): # model_dict, input_channels(3)
|
||||
print('\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 # number of anchors
|
||||
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
|
||||
|
||||
layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
|
||||
for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args
|
||||
m = eval(m) if isinstance(m, str) else m # eval strings
|
||||
for j, a in enumerate(args):
|
||||
try:
|
||||
args[j] = eval(a) if isinstance(a, str) else a # eval strings
|
||||
except:
|
||||
pass
|
||||
|
||||
n = max(round(n * gd), 1) if n > 1 else n # depth gain
|
||||
if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]:
|
||||
c1, c2 = ch[f], args[0]
|
||||
|
||||
# Normal
|
||||
# if i > 0 and args[0] != no: # channel expansion factor
|
||||
# ex = 1.75 # exponential (default 2.0)
|
||||
# e = math.log(c2 / ch[1]) / math.log(2)
|
||||
# c2 = int(ch[1] * ex ** e)
|
||||
# if m != Focus:
|
||||
|
||||
c2 = make_divisible(c2 * gw, 8) if c2 != no else c2
|
||||
|
||||
# Experimental
|
||||
# if i > 0 and args[0] != no: # channel expansion factor
|
||||
# ex = 1 + gw # exponential (default 2.0)
|
||||
# ch1 = 32 # ch[1]
|
||||
# e = math.log(c2 / ch1) / math.log(2) # level 1-n
|
||||
# c2 = int(ch1 * ex ** e)
|
||||
# if m != Focus:
|
||||
# c2 = make_divisible(c2, 8) if c2 != no else c2
|
||||
|
||||
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[-1 if x == -1 else x + 1] for x in f])
|
||||
elif m is Detect:
|
||||
args.append([ch[x + 1] for x in f])
|
||||
if isinstance(args[1], int): # number of anchors
|
||||
args[1] = [list(range(args[1] * 2))] * len(f)
|
||||
else:
|
||||
c2 = ch[f]
|
||||
|
||||
m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module
|
||||
t = str(m)[8:-2].replace('__main__.', '') # module type
|
||||
np = sum([x.numel() for x in m_.parameters()]) # number params
|
||||
m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
|
||||
print('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print
|
||||
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
|
||||
layers.append(m_)
|
||||
ch.append(c2)
|
||||
return nn.Sequential(*layers), sorted(save)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
opt = parser.parse_args()
|
||||
opt.cfg = check_file(opt.cfg) # check file
|
||||
device = select_device(opt.device)
|
||||
|
||||
# Create model
|
||||
model = Model(opt.cfg).to(device)
|
||||
model.train()
|
||||
|
||||
# Profile
|
||||
# img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
|
||||
# y = model(img, profile=True)
|
||||
|
||||
# ONNX export
|
||||
# model.model[-1].export = True
|
||||
# torch.onnx.export(model, img, opt.cfg.replace('.yaml', '.onnx'), verbose=True, opset_version=11)
|
||||
|
||||
# Tensorboard
|
||||
# from torch.utils.tensorboard import SummaryWriter
|
||||
# tb_writer = SummaryWriter()
|
||||
# print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/")
|
||||
# tb_writer.add_graph(model.model, img) # add model to tensorboard
|
||||
# tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard
|
@ -0,0 +1,48 @@
|
||||
# parameters
|
||||
nc: 80 # number of classes
|
||||
depth_multiple: 1.0 # model depth multiple
|
||||
width_multiple: 1.0 # layer channel multiple
|
||||
|
||||
# anchors
|
||||
anchors:
|
||||
- [10,13, 16,30, 33,23] # P3/8
|
||||
- [30,61, 62,45, 59,119] # P4/16
|
||||
- [116,90, 156,198, 373,326] # P5/32
|
||||
|
||||
# YOLOv5 backbone
|
||||
backbone:
|
||||
# [from, number, module, args]
|
||||
[[-1, 1, Focus, [64, 3]], # 0-P1/2
|
||||
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
|
||||
[-1, 3, BottleneckCSP, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 9
|
||||
]
|
||||
|
||||
# YOLOv5 head
|
||||
head:
|
||||
[[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 13
|
||||
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small)
|
||||
|
||||
[-1, 1, Conv, [256, 3, 2]],
|
||||
[[-1, 14], 1, Concat, [1]], # cat head P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium)
|
||||
|
||||
[-1, 1, Conv, [512, 3, 2]],
|
||||
[[-1, 10], 1, Concat, [1]], # cat head P5
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large)
|
||||
|
||||
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
]
|
@ -0,0 +1,48 @@
|
||||
# parameters
|
||||
nc: 80 # number of classes
|
||||
depth_multiple: 0.67 # model depth multiple
|
||||
width_multiple: 0.75 # layer channel multiple
|
||||
|
||||
# anchors
|
||||
anchors:
|
||||
- [10,13, 16,30, 33,23] # P3/8
|
||||
- [30,61, 62,45, 59,119] # P4/16
|
||||
- [116,90, 156,198, 373,326] # P5/32
|
||||
|
||||
# YOLOv5 backbone
|
||||
backbone:
|
||||
# [from, number, module, args]
|
||||
[[-1, 1, Focus, [64, 3]], # 0-P1/2
|
||||
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
|
||||
[-1, 3, BottleneckCSP, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 9
|
||||
]
|
||||
|
||||
# YOLOv5 head
|
||||
head:
|
||||
[[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 13
|
||||
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small)
|
||||
|
||||
[-1, 1, Conv, [256, 3, 2]],
|
||||
[[-1, 14], 1, Concat, [1]], # cat head P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium)
|
||||
|
||||
[-1, 1, Conv, [512, 3, 2]],
|
||||
[[-1, 10], 1, Concat, [1]], # cat head P5
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large)
|
||||
|
||||
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
]
|
@ -0,0 +1,48 @@
|
||||
# parameters
|
||||
nc: 80 # number of classes
|
||||
depth_multiple: 0.33 # model depth multiple
|
||||
width_multiple: 0.50 # layer channel multiple
|
||||
|
||||
# anchors
|
||||
anchors:
|
||||
- [10,13, 16,30, 33,23] # P3/8
|
||||
- [30,61, 62,45, 59,119] # P4/16
|
||||
- [116,90, 156,198, 373,326] # P5/32
|
||||
|
||||
# YOLOv5 backbone
|
||||
backbone:
|
||||
# [from, number, module, args]
|
||||
[[-1, 1, Focus, [64, 3]], # 0-P1/2
|
||||
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
|
||||
[-1, 3, BottleneckCSP, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 9
|
||||
]
|
||||
|
||||
# YOLOv5 head
|
||||
head:
|
||||
[[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 13
|
||||
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small)
|
||||
|
||||
[-1, 1, Conv, [256, 3, 2]],
|
||||
[[-1, 14], 1, Concat, [1]], # cat head P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium)
|
||||
|
||||
[-1, 1, Conv, [512, 3, 2]],
|
||||
[[-1, 10], 1, Concat, [1]], # cat head P5
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large)
|
||||
|
||||
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
]
|
@ -0,0 +1,48 @@
|
||||
# parameters
|
||||
nc: 80 # number of classes
|
||||
depth_multiple: 1.33 # model depth multiple
|
||||
width_multiple: 1.25 # layer channel multiple
|
||||
|
||||
# anchors
|
||||
anchors:
|
||||
- [10,13, 16,30, 33,23] # P3/8
|
||||
- [30,61, 62,45, 59,119] # P4/16
|
||||
- [116,90, 156,198, 373,326] # P5/32
|
||||
|
||||
# YOLOv5 backbone
|
||||
backbone:
|
||||
# [from, number, module, args]
|
||||
[[-1, 1, Focus, [64, 3]], # 0-P1/2
|
||||
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
|
||||
[-1, 3, BottleneckCSP, [128]],
|
||||
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
|
||||
[-1, 9, BottleneckCSP, [256]],
|
||||
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
|
||||
[-1, 9, BottleneckCSP, [512]],
|
||||
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
|
||||
[-1, 1, SPP, [1024, [5, 9, 13]]],
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 9
|
||||
]
|
||||
|
||||
# YOLOv5 head
|
||||
head:
|
||||
[[-1, 1, Conv, [512, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 6], 1, Concat, [1]], # cat backbone P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 13
|
||||
|
||||
[-1, 1, Conv, [256, 1, 1]],
|
||||
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
|
||||
[[-1, 4], 1, Concat, [1]], # cat backbone P3
|
||||
[-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small)
|
||||
|
||||
[-1, 1, Conv, [256, 3, 2]],
|
||||
[[-1, 14], 1, Concat, [1]], # cat head P4
|
||||
[-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium)
|
||||
|
||||
[-1, 1, Conv, [512, 3, 2]],
|
||||
[[-1, 10], 1, Concat, [1]], # cat head P5
|
||||
[-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large)
|
||||
|
||||
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
|
||||
]
|
@ -0,0 +1,291 @@
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
from tqdm import tqdm
|
||||
|
||||
from models.experimental import attempt_load
|
||||
from utils.datasets import create_dataloader
|
||||
from utils.general import (
|
||||
coco80_to_coco91_class, check_file, check_img_size, compute_loss, non_max_suppression,
|
||||
scale_coords, xyxy2xywh, clip_coords, plot_images, xywh2xyxy, box_iou, output_to_target, ap_per_class)
|
||||
from utils.torch_utils import select_device, time_synchronized
|
||||
|
||||
|
||||
def test(data,
|
||||
weights=None,
|
||||
batch_size=16,
|
||||
imgsz=640,
|
||||
conf_thres=0.001,
|
||||
iou_thres=0.6, # for NMS
|
||||
save_json=False,
|
||||
single_cls=False,
|
||||
augment=False,
|
||||
verbose=False,
|
||||
model=None,
|
||||
dataloader=None,
|
||||
save_dir='',
|
||||
merge=False,
|
||||
save_txt=False):
|
||||
# Initialize/load model and set device
|
||||
training = model is not None
|
||||
if training: # called by train.py
|
||||
device = next(model.parameters()).device # get model device
|
||||
|
||||
else: # called directly
|
||||
device = select_device(opt.device, batch_size=batch_size)
|
||||
merge, save_txt = opt.merge, opt.save_txt # use Merge NMS, save *.txt labels
|
||||
if save_txt:
|
||||
out = Path('inference/output')
|
||||
if os.path.exists(out):
|
||||
shutil.rmtree(out) # delete output folder
|
||||
os.makedirs(out) # make new output folder
|
||||
|
||||
# Remove previous
|
||||
for f in glob.glob(str(Path(save_dir) / 'test_batch*.jpg')):
|
||||
os.remove(f)
|
||||
|
||||
# Load model
|
||||
model = attempt_load(weights, map_location=device) # load FP32 model
|
||||
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
|
||||
|
||||
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
|
||||
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
|
||||
# model = nn.DataParallel(model)
|
||||
|
||||
# Half
|
||||
half = device.type != 'cpu' # half precision only supported on CUDA
|
||||
if half:
|
||||
model.half()
|
||||
|
||||
# Configure
|
||||
model.eval()
|
||||
with open(data) as f:
|
||||
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
||||
nc = 1 if single_cls else int(data['nc']) # number of classes
|
||||
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
|
||||
niou = iouv.numel()
|
||||
|
||||
# Dataloader
|
||||
if not training:
|
||||
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
|
||||
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
|
||||
path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
|
||||
dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt,
|
||||
hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0]
|
||||
|
||||
seen = 0
|
||||
names = model.names if hasattr(model, 'names') else model.module.names
|
||||
coco91class = coco80_to_coco91_class()
|
||||
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
|
||||
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
|
||||
loss = torch.zeros(3, device=device)
|
||||
jdict, stats, ap, ap_class = [], [], [], []
|
||||
for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
|
||||
img = img.to(device, non_blocking=True)
|
||||
img = img.half() if half else img.float() # uint8 to fp16/32
|
||||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||||
targets = targets.to(device)
|
||||
nb, _, height, width = img.shape # batch size, channels, height, width
|
||||
whwh = torch.Tensor([width, height, width, height]).to(device)
|
||||
|
||||
# Disable gradients
|
||||
with torch.no_grad():
|
||||
# Run model
|
||||
t = time_synchronized()
|
||||
inf_out, train_out = model(img, augment=augment) # inference and training outputs
|
||||
t0 += time_synchronized() - t
|
||||
|
||||
# Compute loss
|
||||
if training: # if model has loss hyperparameters
|
||||
loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # GIoU, obj, cls
|
||||
|
||||
# Run NMS
|
||||
t = time_synchronized()
|
||||
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, merge=merge)
|
||||
t1 += time_synchronized() - t
|
||||
|
||||
# Statistics per image
|
||||
for si, pred in enumerate(output):
|
||||
labels = targets[targets[:, 0] == si, 1:]
|
||||
nl = len(labels)
|
||||
tcls = labels[:, 0].tolist() if nl else [] # target class
|
||||
seen += 1
|
||||
|
||||
if pred is None:
|
||||
if nl:
|
||||
stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
|
||||
continue
|
||||
|
||||
# Append to text file
|
||||
if save_txt:
|
||||
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh
|
||||
txt_path = str(out / Path(paths[si]).stem)
|
||||
pred[:, :4] = scale_coords(img[si].shape[1:], pred[:, :4], shapes[si][0], shapes[si][1]) # to original
|
||||
for *xyxy, conf, cls in pred:
|
||||
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
||||
with open(txt_path + '.txt', 'a') as f:
|
||||
f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format
|
||||
|
||||
# Clip boxes to image bounds
|
||||
clip_coords(pred, (height, width))
|
||||
|
||||
# Append to pycocotools JSON dictionary
|
||||
if save_json:
|
||||
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
|
||||
image_id = Path(paths[si]).stem
|
||||
box = pred[:, :4].clone() # xyxy
|
||||
scale_coords(img[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape
|
||||
box = xyxy2xywh(box) # xywh
|
||||
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
|
||||
for p, b in zip(pred.tolist(), box.tolist()):
|
||||
jdict.append({'image_id': int(image_id) if image_id.isnumeric() else image_id,
|
||||
'category_id': coco91class[int(p[5])],
|
||||
'bbox': [round(x, 3) for x in b],
|
||||
'score': round(p[4], 5)})
|
||||
|
||||
# Assign all predictions as incorrect
|
||||
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
|
||||
if nl:
|
||||
detected = [] # target indices
|
||||
tcls_tensor = labels[:, 0]
|
||||
|
||||
# target boxes
|
||||
tbox = xywh2xyxy(labels[:, 1:5]) * whwh
|
||||
|
||||
# Per target class
|
||||
for cls in torch.unique(tcls_tensor):
|
||||
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # prediction indices
|
||||
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # target indices
|
||||
|
||||
# Search for detections
|
||||
if pi.shape[0]:
|
||||
# Prediction to target ious
|
||||
ious, i = box_iou(pred[pi, :4], tbox[ti]).max(1) # best ious, indices
|
||||
|
||||
# Append detections
|
||||
for j in (ious > iouv[0]).nonzero(as_tuple=False):
|
||||
d = ti[i[j]] # detected target
|
||||
if d not in detected:
|
||||
detected.append(d)
|
||||
correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn
|
||||
if len(detected) == nl: # all targets already located in image
|
||||
break
|
||||
|
||||
# Append statistics (correct, conf, pcls, tcls)
|
||||
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
|
||||
|
||||
# Plot images
|
||||
if batch_i < 1:
|
||||
f = Path(save_dir) / ('test_batch%g_gt.jpg' % batch_i) # filename
|
||||
plot_images(img, targets, paths, str(f), names) # ground truth
|
||||
f = Path(save_dir) / ('test_batch%g_pred.jpg' % batch_i)
|
||||
plot_images(img, output_to_target(output, width, height), paths, str(f), names) # predictions
|
||||
|
||||
# Compute statistics
|
||||
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
|
||||
if len(stats) and stats[0].any():
|
||||
p, r, ap, f1, ap_class = ap_per_class(*stats)
|
||||
p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, AP@0.5, AP@0.5:0.95]
|
||||
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
|
||||
nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
|
||||
else:
|
||||
nt = torch.zeros(1)
|
||||
|
||||
# Print results
|
||||
pf = '%20s' + '%12.3g' * 6 # print format
|
||||
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
|
||||
|
||||
# Print results per class
|
||||
if verbose and nc > 1 and len(stats):
|
||||
for i, c in enumerate(ap_class):
|
||||
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
|
||||
|
||||
# Print speeds
|
||||
t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
|
||||
if not training:
|
||||
print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t)
|
||||
|
||||
# Save JSON
|
||||
if save_json and len(jdict):
|
||||
f = 'detections_val2017_%s_results.json' % \
|
||||
(weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename
|
||||
print('\nCOCO mAP with pycocotools... saving %s...' % f)
|
||||
with open(f, 'w') as file:
|
||||
json.dump(jdict, file)
|
||||
|
||||
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
|
||||
from pycocotools.coco import COCO
|
||||
from pycocotools.cocoeval import COCOeval
|
||||
|
||||
imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files]
|
||||
cocoGt = COCO(glob.glob('../coco/annotations/instances_val*.json')[0]) # initialize COCO ground truth api
|
||||
cocoDt = cocoGt.loadRes(f) # initialize COCO pred api
|
||||
cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
|
||||
cocoEval.params.imgIds = imgIds # image IDs to evaluate
|
||||
cocoEval.evaluate()
|
||||
cocoEval.accumulate()
|
||||
cocoEval.summarize()
|
||||
map, map50 = cocoEval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
|
||||
except Exception as e:
|
||||
print('ERROR: pycocotools unable to run: %s' % e)
|
||||
|
||||
# Return results
|
||||
model.float() # for training
|
||||
maps = np.zeros(nc) + map
|
||||
for i, c in enumerate(ap_class):
|
||||
maps[c] = ap[i]
|
||||
return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(prog='test.py')
|
||||
parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
|
||||
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
|
||||
parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
|
||||
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
||||
parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
|
||||
parser.add_argument('--iou-thres', type=float, default=0.65, help='IOU threshold for NMS')
|
||||
parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
|
||||
parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
|
||||
parser.add_argument('--augment', action='store_true', help='augmented inference')
|
||||
parser.add_argument('--merge', action='store_true', help='use Merge NMS')
|
||||
parser.add_argument('--verbose', action='store_true', help='report mAP by class')
|
||||
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
||||
opt = parser.parse_args()
|
||||
opt.save_json |= opt.data.endswith('coco.yaml')
|
||||
opt.data = check_file(opt.data) # check file
|
||||
print(opt)
|
||||
|
||||
if opt.task in ['val', 'test']: # run normally
|
||||
test(opt.data,
|
||||
opt.weights,
|
||||
opt.batch_size,
|
||||
opt.img_size,
|
||||
opt.conf_thres,
|
||||
opt.iou_thres,
|
||||
opt.save_json,
|
||||
opt.single_cls,
|
||||
opt.augment,
|
||||
opt.verbose)
|
||||
|
||||
elif opt.task == 'study': # run over a range of settings and save/plot
|
||||
for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov3-spp.pt']:
|
||||
f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem) # filename to save to
|
||||
x = list(range(352, 832, 64)) # x axis
|
||||
y = [] # y axis
|
||||
for i in x: # img-size
|
||||
print('\nRunning %s point %s...' % (f, i))
|
||||
r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json)
|
||||
y.append(r + t) # results and times
|
||||
np.savetxt(f, y, fmt='%10.4g') # save
|
||||
os.system('zip -r study.zip study_*.txt')
|
||||
# plot_study_txt(f, x) # plot
|
@ -0,0 +1,553 @@
|
||||
import argparse
|
||||
import glob
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
import torch.optim.lr_scheduler as lr_scheduler
|
||||
import torch.utils.data
|
||||
import yaml
|
||||
from torch.cuda import amp
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
from tqdm import tqdm
|
||||
|
||||
import test # import test.py to get mAP after each epoch
|
||||
from models.yolo import Model
|
||||
from utils.datasets import create_dataloader
|
||||
from utils.general import (
|
||||
check_img_size, torch_distributed_zero_first, labels_to_class_weights, plot_labels, check_anchors,
|
||||
labels_to_image_weights, compute_loss, plot_images, fitness, strip_optimizer, plot_results,
|
||||
get_latest_run, check_git_status, check_file, increment_dir, print_mutation, plot_evolution)
|
||||
from utils.google_utils import attempt_download
|
||||
from utils.torch_utils import init_seeds, ModelEMA, select_device
|
||||
|
||||
# Hyperparameters
|
||||
hyp = {'lr0': 0.01, # initial learning rate (SGD=1E-2, Adam=1E-3)
|
||||
'momentum': 0.937, # SGD momentum/Adam beta1
|
||||
'weight_decay': 5e-4, # optimizer weight decay
|
||||
'giou': 0.05, # GIoU loss gain
|
||||
'cls': 0.5, # cls loss gain
|
||||
'cls_pw': 1.0, # cls BCELoss positive_weight
|
||||
'obj': 1.0, # obj loss gain (scale with pixels)
|
||||
'obj_pw': 1.0, # obj BCELoss positive_weight
|
||||
'iou_t': 0.20, # IoU training threshold
|
||||
'anchor_t': 4.0, # anchor-multiple threshold
|
||||
'fl_gamma': 0.0, # focal loss gamma (efficientDet default gamma=1.5)
|
||||
'hsv_h': 0.015, # image HSV-Hue augmentation (fraction)
|
||||
'hsv_s': 0.7, # image HSV-Saturation augmentation (fraction)
|
||||
'hsv_v': 0.4, # image HSV-Value augmentation (fraction)
|
||||
'degrees': 0.0, # image rotation (+/- deg)
|
||||
'translate': 0.5, # image translation (+/- fraction)
|
||||
'scale': 0.5, # image scale (+/- gain)
|
||||
'shear': 0.0, # image shear (+/- deg)
|
||||
'perspective': 0.0, # image perspective (+/- fraction), range 0-0.001
|
||||
'flipud': 0.0, # image flip up-down (probability)
|
||||
'fliplr': 0.5, # image flip left-right (probability)
|
||||
'mixup': 0.0} # image mixup (probability)
|
||||
|
||||
|
||||
def train(hyp, opt, device, tb_writer=None):
|
||||
print(f'Hyperparameters {hyp}')
|
||||
log_dir = tb_writer.log_dir if tb_writer else 'runs/evolve' # run directory
|
||||
wdir = str(Path(log_dir) / 'weights') + os.sep # weights directory
|
||||
os.makedirs(wdir, exist_ok=True)
|
||||
last = wdir + 'last.pt'
|
||||
best = wdir + 'best.pt'
|
||||
results_file = log_dir + os.sep + 'results.txt'
|
||||
epochs, batch_size, total_batch_size, weights, rank = \
|
||||
opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.local_rank
|
||||
# TODO: Use DDP logging. Only the first process is allowed to log.
|
||||
|
||||
# Save run settings
|
||||
with open(Path(log_dir) / 'hyp.yaml', 'w') as f:
|
||||
yaml.dump(hyp, f, sort_keys=False)
|
||||
with open(Path(log_dir) / 'opt.yaml', 'w') as f:
|
||||
yaml.dump(vars(opt), f, sort_keys=False)
|
||||
|
||||
# Configure
|
||||
cuda = device.type != 'cpu'
|
||||
init_seeds(2 + rank)
|
||||
with open(opt.data) as f:
|
||||
data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
|
||||
train_path = data_dict['train']
|
||||
test_path = data_dict['val']
|
||||
nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names']) # number classes, names
|
||||
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check
|
||||
|
||||
# Remove previous results
|
||||
if rank in [-1, 0]:
|
||||
for f in glob.glob('*_batch*.jpg') + glob.glob(results_file):
|
||||
os.remove(f)
|
||||
|
||||
# Create model
|
||||
model = Model(opt.cfg, nc=nc).to(device)
|
||||
|
||||
# Image sizes
|
||||
gs = int(max(model.stride)) # grid size (max stride)
|
||||
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples
|
||||
|
||||
# Optimizer
|
||||
nbs = 64 # nominal batch size
|
||||
# default DDP implementation is slow for accumulation according to: https://pytorch.org/docs/stable/notes/ddp.html
|
||||
# all-reduce operation is carried out during loss.backward().
|
||||
# Thus, there would be redundant all-reduce communications in a accumulation procedure,
|
||||
# which means, the result is still right but the training speed gets slower.
|
||||
# TODO: If acceleration is needed, there is an implementation of allreduce_post_accumulation
|
||||
# in https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/BERT/run_pretraining.py
|
||||
accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing
|
||||
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
|
||||
|
||||
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
|
||||
for k, v in model.named_parameters():
|
||||
if v.requires_grad:
|
||||
if '.bias' in k:
|
||||
pg2.append(v) # biases
|
||||
elif '.weight' in k and '.bn' not in k:
|
||||
pg1.append(v) # apply weight decay
|
||||
else:
|
||||
pg0.append(v) # all else
|
||||
|
||||
if opt.adam:
|
||||
optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
|
||||
else:
|
||||
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
|
||||
|
||||
optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay
|
||||
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
|
||||
print('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
|
||||
del pg0, pg1, pg2
|
||||
|
||||
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
|
||||
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
|
||||
lf = lambda x: (((1 + math.cos(x * math.pi / epochs)) / 2) ** 1.0) * 0.8 + 0.2 # cosine
|
||||
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
|
||||
# plot_lr_scheduler(optimizer, scheduler, epochs)
|
||||
|
||||
# Load Model
|
||||
with torch_distributed_zero_first(rank):
|
||||
attempt_download(weights)
|
||||
start_epoch, best_fitness = 0, 0.0
|
||||
if weights.endswith('.pt'): # pytorch format
|
||||
ckpt = torch.load(weights, map_location=device) # load checkpoint
|
||||
|
||||
# load model
|
||||
try:
|
||||
exclude = ['anchor'] # exclude keys
|
||||
ckpt['model'] = {k: v for k, v in ckpt['model'].float().state_dict().items()
|
||||
if k in model.state_dict() and not any(x in k for x in exclude)
|
||||
and model.state_dict()[k].shape == v.shape}
|
||||
model.load_state_dict(ckpt['model'], strict=False)
|
||||
print('Transferred %g/%g items from %s' % (len(ckpt['model']), len(model.state_dict()), weights))
|
||||
except KeyError as e:
|
||||
s = "%s is not compatible with %s. This may be due to model differences or %s may be out of date. " \
|
||||
"Please delete or update %s and try again, or use --weights '' to train from scratch." \
|
||||
% (weights, opt.cfg, weights, weights)
|
||||
raise KeyError(s) from e
|
||||
|
||||
# load optimizer
|
||||
if ckpt['optimizer'] is not None:
|
||||
optimizer.load_state_dict(ckpt['optimizer'])
|
||||
best_fitness = ckpt['best_fitness']
|
||||
|
||||
# load results
|
||||
if ckpt.get('training_results') is not None:
|
||||
with open(results_file, 'w') as file:
|
||||
file.write(ckpt['training_results']) # write results.txt
|
||||
|
||||
# epochs
|
||||
start_epoch = ckpt['epoch'] + 1
|
||||
if epochs < start_epoch:
|
||||
print('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
|
||||
(weights, ckpt['epoch'], epochs))
|
||||
epochs += ckpt['epoch'] # finetune additional epochs
|
||||
|
||||
del ckpt
|
||||
|
||||
# DP mode
|
||||
if cuda and rank == -1 and torch.cuda.device_count() > 1:
|
||||
model = torch.nn.DataParallel(model)
|
||||
|
||||
# SyncBatchNorm
|
||||
if opt.sync_bn and cuda and rank != -1:
|
||||
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
|
||||
print('Using SyncBatchNorm()')
|
||||
|
||||
# Exponential moving average
|
||||
ema = ModelEMA(model) if rank in [-1, 0] else None
|
||||
|
||||
# DDP mode
|
||||
if cuda and rank != -1:
|
||||
model = DDP(model, device_ids=[rank], output_device=rank)
|
||||
|
||||
# Trainloader
|
||||
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, hyp=hyp, augment=True,
|
||||
cache=opt.cache_images, rect=opt.rect, local_rank=rank,
|
||||
world_size=opt.world_size)
|
||||
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
|
||||
nb = len(dataloader) # number of batches
|
||||
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
|
||||
|
||||
# Testloader
|
||||
if rank in [-1, 0]:
|
||||
# local_rank is set to -1. Because only the first process is expected to do evaluation.
|
||||
testloader = create_dataloader(test_path, imgsz_test, total_batch_size, gs, opt, hyp=hyp, augment=False,
|
||||
cache=opt.cache_images, rect=True, local_rank=-1, world_size=opt.world_size)[0]
|
||||
|
||||
# Model parameters
|
||||
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
|
||||
model.nc = nc # attach number of classes to model
|
||||
model.hyp = hyp # attach hyperparameters to model
|
||||
model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou)
|
||||
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights
|
||||
model.names = names
|
||||
|
||||
# Class frequency
|
||||
if rank in [-1, 0]:
|
||||
labels = np.concatenate(dataset.labels, 0)
|
||||
c = torch.tensor(labels[:, 0]) # classes
|
||||
# cf = torch.bincount(c.long(), minlength=nc) + 1.
|
||||
# model._initialize_biases(cf.to(device))
|
||||
plot_labels(labels, save_dir=log_dir)
|
||||
if tb_writer:
|
||||
# tb_writer.add_hparams(hyp, {}) # causes duplicate https://github.com/ultralytics/yolov5/pull/384
|
||||
tb_writer.add_histogram('classes', c, 0)
|
||||
|
||||
# Check anchors
|
||||
if not opt.noautoanchor:
|
||||
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
|
||||
|
||||
# Start training
|
||||
t0 = time.time()
|
||||
nw = max(3 * nb, 1e3) # number of warmup iterations, max(3 epochs, 1k iterations)
|
||||
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
|
||||
maps = np.zeros(nc) # mAP per class
|
||||
results = (0, 0, 0, 0, 0, 0, 0) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
|
||||
scheduler.last_epoch = start_epoch - 1 # do not move
|
||||
scaler = amp.GradScaler(enabled=cuda)
|
||||
if rank in [0, -1]:
|
||||
print('Image sizes %g train, %g test' % (imgsz, imgsz_test))
|
||||
print('Using %g dataloader workers' % dataloader.num_workers)
|
||||
print('Starting training for %g epochs...' % epochs)
|
||||
# torch.autograd.set_detect_anomaly(True)
|
||||
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
|
||||
model.train()
|
||||
|
||||
# Update image weights (optional)
|
||||
if dataset.image_weights:
|
||||
# Generate indices
|
||||
if rank in [-1, 0]:
|
||||
w = model.class_weights.cpu().numpy() * (1 - maps) ** 2 # class weights
|
||||
image_weights = labels_to_image_weights(dataset.labels, nc=nc, class_weights=w)
|
||||
dataset.indices = random.choices(range(dataset.n), weights=image_weights,
|
||||
k=dataset.n) # rand weighted idx
|
||||
# Broadcast if DDP
|
||||
if rank != -1:
|
||||
indices = torch.zeros([dataset.n], dtype=torch.int)
|
||||
if rank == 0:
|
||||
indices[:] = torch.from_tensor(dataset.indices, dtype=torch.int)
|
||||
dist.broadcast(indices, 0)
|
||||
if rank != 0:
|
||||
dataset.indices = indices.cpu().numpy()
|
||||
|
||||
# Update mosaic border
|
||||
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
|
||||
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
|
||||
|
||||
mloss = torch.zeros(4, device=device) # mean losses
|
||||
if rank != -1:
|
||||
dataloader.sampler.set_epoch(epoch)
|
||||
pbar = enumerate(dataloader)
|
||||
if rank in [-1, 0]:
|
||||
print(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'GIoU', 'obj', 'cls', 'total', 'targets', 'img_size'))
|
||||
pbar = tqdm(pbar, total=nb) # progress bar
|
||||
optimizer.zero_grad()
|
||||
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
|
||||
ni = i + nb * epoch # number integrated batches (since train start)
|
||||
imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
|
||||
|
||||
# Warmup
|
||||
if ni <= nw:
|
||||
xi = [0, nw] # x interp
|
||||
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou)
|
||||
accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round())
|
||||
for j, x in enumerate(optimizer.param_groups):
|
||||
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
|
||||
x['lr'] = np.interp(ni, xi, [0.1 if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
|
||||
if 'momentum' in x:
|
||||
x['momentum'] = np.interp(ni, xi, [0.9, hyp['momentum']])
|
||||
|
||||
# Multi-scale
|
||||
if opt.multi_scale:
|
||||
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
|
||||
sf = sz / max(imgs.shape[2:]) # scale factor
|
||||
if sf != 1:
|
||||
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
|
||||
imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
|
||||
|
||||
# Autocast
|
||||
with amp.autocast(enabled=cuda):
|
||||
# Forward
|
||||
pred = model(imgs)
|
||||
|
||||
# Loss
|
||||
loss, loss_items = compute_loss(pred, targets.to(device), model) # scaled by batch_size
|
||||
if rank != -1:
|
||||
loss *= opt.world_size # gradient averaged between devices in DDP mode
|
||||
# if not torch.isfinite(loss):
|
||||
# print('WARNING: non-finite loss, ending training ', loss_items)
|
||||
# return results
|
||||
|
||||
# Backward
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
# Optimize
|
||||
if ni % accumulate == 0:
|
||||
scaler.step(optimizer) # optimizer.step
|
||||
scaler.update()
|
||||
optimizer.zero_grad()
|
||||
if ema is not None:
|
||||
ema.update(model)
|
||||
|
||||
# Print
|
||||
if rank in [-1, 0]:
|
||||
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
|
||||
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB)
|
||||
s = ('%10s' * 2 + '%10.4g' * 6) % (
|
||||
'%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1])
|
||||
pbar.set_description(s)
|
||||
|
||||
# Plot
|
||||
if ni < 3:
|
||||
f = str(Path(log_dir) / ('train_batch%g.jpg' % ni)) # filename
|
||||
result = plot_images(images=imgs, targets=targets, paths=paths, fname=f)
|
||||
if tb_writer and result is not None:
|
||||
tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
|
||||
# tb_writer.add_graph(model, imgs) # add model to tensorboard
|
||||
|
||||
# end batch ------------------------------------------------------------------------------------------------
|
||||
|
||||
# Scheduler
|
||||
scheduler.step()
|
||||
|
||||
# DDP process 0 or single-GPU
|
||||
if rank in [-1, 0]:
|
||||
# mAP
|
||||
if ema is not None:
|
||||
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride'])
|
||||
final_epoch = epoch + 1 == epochs
|
||||
if not opt.notest or final_epoch: # Calculate mAP
|
||||
results, maps, times = test.test(opt.data,
|
||||
batch_size=total_batch_size,
|
||||
imgsz=imgsz_test,
|
||||
save_json=final_epoch and opt.data.endswith(os.sep + 'coco.yaml'),
|
||||
model=ema.ema.module if hasattr(ema.ema, 'module') else ema.ema,
|
||||
single_cls=opt.single_cls,
|
||||
dataloader=testloader,
|
||||
save_dir=log_dir)
|
||||
|
||||
# Write
|
||||
with open(results_file, 'a') as f:
|
||||
f.write(s + '%10.4g' * 7 % results + '\n') # P, R, mAP, F1, test_losses=(GIoU, obj, cls)
|
||||
if len(opt.name) and opt.bucket:
|
||||
os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name))
|
||||
|
||||
# Tensorboard
|
||||
if tb_writer:
|
||||
tags = ['train/giou_loss', 'train/obj_loss', 'train/cls_loss',
|
||||
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
|
||||
'val/giou_loss', 'val/obj_loss', 'val/cls_loss']
|
||||
for x, tag in zip(list(mloss[:-1]) + list(results), tags):
|
||||
tb_writer.add_scalar(tag, x, epoch)
|
||||
|
||||
# Update best mAP
|
||||
fi = fitness(np.array(results).reshape(1, -1)) # fitness_i = weighted combination of [P, R, mAP, F1]
|
||||
if fi > best_fitness:
|
||||
best_fitness = fi
|
||||
|
||||
# Save model
|
||||
save = (not opt.nosave) or (final_epoch and not opt.evolve)
|
||||
if save:
|
||||
with open(results_file, 'r') as f: # create checkpoint
|
||||
ckpt = {'epoch': epoch,
|
||||
'best_fitness': best_fitness,
|
||||
'training_results': f.read(),
|
||||
'model': ema.ema.module if hasattr(ema, 'module') else ema.ema,
|
||||
'optimizer': None if final_epoch else optimizer.state_dict()}
|
||||
|
||||
# Save last, best and delete
|
||||
torch.save(ckpt, last)
|
||||
if best_fitness == fi:
|
||||
torch.save(ckpt, best)
|
||||
del ckpt
|
||||
# end epoch ----------------------------------------------------------------------------------------------------
|
||||
# end training
|
||||
|
||||
if rank in [-1, 0]:
|
||||
# Strip optimizers
|
||||
n = ('_' if len(opt.name) and not opt.name.isnumeric() else '') + opt.name
|
||||
fresults, flast, fbest = 'results%s.txt' % n, wdir + 'last%s.pt' % n, wdir + 'best%s.pt' % n
|
||||
for f1, f2 in zip([wdir + 'last.pt', wdir + 'best.pt', 'results.txt'], [flast, fbest, fresults]):
|
||||
if os.path.exists(f1):
|
||||
os.rename(f1, f2) # rename
|
||||
ispt = f2.endswith('.pt') # is *.pt
|
||||
strip_optimizer(f2) if ispt else None # strip optimizer
|
||||
os.system('gsutil cp %s gs://%s/weights' % (f2, opt.bucket)) if opt.bucket and ispt else None # upload
|
||||
# Finish
|
||||
if not opt.evolve:
|
||||
plot_results(save_dir=log_dir) # save as results.png
|
||||
print('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
|
||||
|
||||
dist.destroy_process_group() if rank not in [-1, 0] else None
|
||||
torch.cuda.empty_cache()
|
||||
return results
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
|
||||
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
|
||||
parser.add_argument('--hyp', type=str, default='', help='hyp.yaml path (optional)')
|
||||
parser.add_argument('--epochs', type=int, default=300)
|
||||
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
|
||||
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes')
|
||||
parser.add_argument('--rect', action='store_true', help='rectangular training')
|
||||
parser.add_argument('--resume', nargs='?', const='get_last', default=False,
|
||||
help='resume from given path/last.pt, or most recent run if blank')
|
||||
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
|
||||
parser.add_argument('--notest', action='store_true', help='only test final epoch')
|
||||
parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
|
||||
parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
|
||||
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
|
||||
parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
|
||||
parser.add_argument('--weights', type=str, default='', help='initial weights path')
|
||||
parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
|
||||
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
|
||||
parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
|
||||
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
|
||||
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
|
||||
opt = parser.parse_args()
|
||||
|
||||
# Resume
|
||||
last = get_latest_run() if opt.resume == 'get_last' else opt.resume # resume from most recent run
|
||||
if last and not opt.weights:
|
||||
print(f'Resuming training from {last}')
|
||||
opt.weights = last if opt.resume and not opt.weights else opt.weights
|
||||
|
||||
if opt.local_rank in [-1, 0]:
|
||||
check_git_status()
|
||||
opt.cfg = check_file(opt.cfg) # check file
|
||||
opt.data = check_file(opt.data) # check file
|
||||
if opt.hyp: # update hyps
|
||||
opt.hyp = check_file(opt.hyp) # check file
|
||||
with open(opt.hyp) as f:
|
||||
hyp.update(yaml.load(f, Loader=yaml.FullLoader)) # update hyps
|
||||
opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size))) # extend to 2 sizes (train, test)
|
||||
device = select_device(opt.device, batch_size=opt.batch_size)
|
||||
opt.total_batch_size = opt.batch_size
|
||||
opt.world_size = 1
|
||||
|
||||
# DDP mode
|
||||
if opt.local_rank != -1:
|
||||
assert torch.cuda.device_count() > opt.local_rank
|
||||
torch.cuda.set_device(opt.local_rank)
|
||||
device = torch.device('cuda', opt.local_rank)
|
||||
dist.init_process_group(backend='nccl', init_method='env://') # distributed backend
|
||||
opt.world_size = dist.get_world_size()
|
||||
assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count'
|
||||
opt.batch_size = opt.total_batch_size // opt.world_size
|
||||
|
||||
print(opt)
|
||||
|
||||
# Train
|
||||
if not opt.evolve:
|
||||
tb_writer = None
|
||||
if opt.local_rank in [-1, 0]:
|
||||
print('Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/')
|
||||
tb_writer = SummaryWriter(log_dir=increment_dir('runs/exp', opt.name))
|
||||
|
||||
train(hyp, opt, device, tb_writer)
|
||||
|
||||
# Evolve hyperparameters (optional)
|
||||
else:
|
||||
# Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
|
||||
meta = {'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
|
||||
'momentum': (0.1, 0.6, 0.98), # SGD momentum/Adam beta1
|
||||
'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
|
||||
'giou': (1, 0.02, 0.2), # GIoU loss gain
|
||||
'cls': (1, 0.2, 4.0), # cls loss gain
|
||||
'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
|
||||
'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
|
||||
'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
|
||||
'iou_t': (0, 0.1, 0.7), # IoU training threshold
|
||||
'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
|
||||
'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
|
||||
'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
|
||||
'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)
|
||||
'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
|
||||
'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
|
||||
'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
|
||||
'scale': (1, 0.0, 0.9), # image scale (+/- gain)
|
||||
'shear': (1, 0.0, 10.0), # image shear (+/- deg)
|
||||
'perspective': (1, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
|
||||
'flipud': (0, 0.0, 1.0), # image flip up-down (probability)
|
||||
'fliplr': (1, 0.0, 1.0), # image flip left-right (probability)
|
||||
'mixup': (1, 0.0, 1.0)} # image mixup (probability)
|
||||
|
||||
assert opt.local_rank == -1, 'DDP mode not implemented for --evolve'
|
||||
opt.notest, opt.nosave = True, True # only test/save final epoch
|
||||
# ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
|
||||
yaml_file = Path('runs/evolve/hyp_evolved.yaml') # save best result here
|
||||
if opt.bucket:
|
||||
os.system('gsutil cp gs://%s/evolve.txt .' % opt.bucket) # download evolve.txt if exists
|
||||
|
||||
for _ in range(100): # generations to evolve
|
||||
if os.path.exists('evolve.txt'): # if evolve.txt exists: select best hyps and mutate
|
||||
# Select parent(s)
|
||||
parent = 'single' # parent selection method: 'single' or 'weighted'
|
||||
x = np.loadtxt('evolve.txt', ndmin=2)
|
||||
n = min(5, len(x)) # number of previous results to consider
|
||||
x = x[np.argsort(-fitness(x))][:n] # top n mutations
|
||||
w = fitness(x) - fitness(x).min() # weights
|
||||
if parent == 'single' or len(x) == 1:
|
||||
# x = x[random.randint(0, n - 1)] # random selection
|
||||
x = x[random.choices(range(n), weights=w)[0]] # weighted selection
|
||||
elif parent == 'weighted':
|
||||
x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination
|
||||
|
||||
# Mutate
|
||||
mp, s = 0.9, 0.2 # mutation probability, sigma
|
||||
npr = np.random
|
||||
npr.seed(int(time.time()))
|
||||
g = np.array([x[0] for x in meta.values()]) # gains 0-1
|
||||
ng = len(meta)
|
||||
v = np.ones(ng)
|
||||
while all(v == 1): # mutate until a change occurs (prevent duplicates)
|
||||
v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
|
||||
for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
|
||||
hyp[k] = float(x[i + 7] * v[i]) # mutate
|
||||
|
||||
# Constrain to limits
|
||||
for k, v in meta.items():
|
||||
hyp[k] = max(hyp[k], v[1]) # lower limit
|
||||
hyp[k] = min(hyp[k], v[2]) # upper limit
|
||||
hyp[k] = round(hyp[k], 5) # significant digits
|
||||
|
||||
# Train mutation
|
||||
results = train(hyp.copy(), opt, device)
|
||||
|
||||
# Write mutation results
|
||||
print_mutation(hyp.copy(), results, yaml_file, opt.bucket)
|
||||
|
||||
# Plot results
|
||||
plot_evolution(yaml_file)
|
||||
print('Hyperparameter evolution complete. Best results saved as: %s\nCommand to train a new model with these '
|
||||
'hyperparameters: $ python train.py --hyp %s' % (yaml_file, yaml_file))
|
File diff suppressed because one or more lines are too long
@ -0,0 +1,69 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
# Swish https://arxiv.org/pdf/1905.02244.pdf ---------------------------------------------------------------------------
|
||||
class Swish(nn.Module): #
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
class HardSwish(nn.Module):
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
return x * F.hardtanh(x + 3, 0., 6., True) / 6.
|
||||
|
||||
|
||||
class MemoryEfficientSwish(nn.Module):
|
||||
class F(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, x):
|
||||
ctx.save_for_backward(x)
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
x = ctx.saved_tensors[0]
|
||||
sx = torch.sigmoid(x)
|
||||
return grad_output * (sx * (1 + x * (1 - sx)))
|
||||
|
||||
def forward(self, x):
|
||||
return self.F.apply(x)
|
||||
|
||||
|
||||
# Mish https://github.com/digantamisra98/Mish --------------------------------------------------------------------------
|
||||
class Mish(nn.Module):
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
return x * F.softplus(x).tanh()
|
||||
|
||||
|
||||
class MemoryEfficientMish(nn.Module):
|
||||
class F(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, x):
|
||||
ctx.save_for_backward(x)
|
||||
return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x)))
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
x = ctx.saved_tensors[0]
|
||||
sx = torch.sigmoid(x)
|
||||
fx = F.softplus(x).tanh()
|
||||
return grad_output * (fx + x * sx * (1 - fx * fx))
|
||||
|
||||
def forward(self, x):
|
||||
return self.F.apply(x)
|
||||
|
||||
|
||||
# FReLU https://arxiv.org/abs/2007.11824 -------------------------------------------------------------------------------
|
||||
class FReLU(nn.Module):
|
||||
def __init__(self, c1, k=3): # ch_in, kernel
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1)
|
||||
self.bn = nn.BatchNorm2d(c1)
|
||||
|
||||
def forward(self, x):
|
||||
return torch.max(x, self.bn(self.conv(x)))
|
@ -0,0 +1,907 @@
|
||||
import glob
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
from threading import Thread
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image, ExifTags
|
||||
from torch.utils.data import Dataset
|
||||
from tqdm import tqdm
|
||||
|
||||
from yolov5.utils.general import xyxy2xywh, xywh2xyxy, torch_distributed_zero_first
|
||||
|
||||
help_url = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
|
||||
img_formats = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.dng']
|
||||
vid_formats = ['.mov', '.avi', '.mp4', '.mpg', '.mpeg', '.m4v', '.wmv', '.mkv']
|
||||
|
||||
# Get orientation exif tag
|
||||
for orientation in ExifTags.TAGS.keys():
|
||||
if ExifTags.TAGS[orientation] == 'Orientation':
|
||||
break
|
||||
|
||||
|
||||
def get_hash(files):
|
||||
# Returns a single hash value of a list of files
|
||||
return sum(os.path.getsize(f) for f in files if os.path.isfile(f))
|
||||
|
||||
|
||||
def exif_size(img):
|
||||
# Returns exif-corrected PIL size
|
||||
s = img.size # (width, height)
|
||||
try:
|
||||
rotation = dict(img._getexif().items())[orientation]
|
||||
if rotation == 6: # rotation 270
|
||||
s = (s[1], s[0])
|
||||
elif rotation == 8: # rotation 90
|
||||
s = (s[1], s[0])
|
||||
except:
|
||||
pass
|
||||
|
||||
return s
|
||||
|
||||
|
||||
def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False,
|
||||
local_rank=-1, world_size=1):
|
||||
# Make sure only the first process in DDP process the dataset first, and the following others can use the cache.
|
||||
with torch_distributed_zero_first(local_rank):
|
||||
dataset = LoadImagesAndLabels(path, imgsz, batch_size,
|
||||
augment=augment, # augment images
|
||||
hyp=hyp, # augmentation hyperparameters
|
||||
rect=rect, # rectangular training
|
||||
cache_images=cache,
|
||||
single_cls=opt.single_cls,
|
||||
stride=int(stride),
|
||||
pad=pad)
|
||||
|
||||
batch_size = min(batch_size, len(dataset))
|
||||
nw = min([os.cpu_count() // world_size, batch_size if batch_size > 1 else 0, 8]) # number of workers
|
||||
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) if local_rank != -1 else None
|
||||
dataloader = torch.utils.data.DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
num_workers=nw,
|
||||
sampler=train_sampler,
|
||||
pin_memory=True,
|
||||
collate_fn=LoadImagesAndLabels.collate_fn)
|
||||
return dataloader, dataset
|
||||
|
||||
|
||||
class LoadImages: # for inference
|
||||
def __init__(self, path, img_size=640):
|
||||
p = str(Path(path)) # os-agnostic
|
||||
p = os.path.abspath(p) # absolute path
|
||||
if '*' in p:
|
||||
files = sorted(glob.glob(p)) # glob
|
||||
elif os.path.isdir(p):
|
||||
files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir
|
||||
elif os.path.isfile(p):
|
||||
files = [p] # files
|
||||
else:
|
||||
raise Exception('ERROR: %s does not exist' % p)
|
||||
|
||||
images = [x for x in files if os.path.splitext(x)[-1].lower() in img_formats]
|
||||
videos = [x for x in files if os.path.splitext(x)[-1].lower() in vid_formats]
|
||||
ni, nv = len(images), len(videos)
|
||||
|
||||
self.img_size = img_size
|
||||
self.files = images + videos
|
||||
self.nf = ni + nv # number of files
|
||||
self.video_flag = [False] * ni + [True] * nv
|
||||
self.mode = 'images'
|
||||
if any(videos):
|
||||
self.new_video(videos[0]) # new video
|
||||
else:
|
||||
self.cap = None
|
||||
assert self.nf > 0, 'No images or videos found in %s. Supported formats are:\nimages: %s\nvideos: %s' % \
|
||||
(p, img_formats, vid_formats)
|
||||
|
||||
def __iter__(self):
|
||||
self.count = 0
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self.count == self.nf:
|
||||
raise StopIteration
|
||||
path = self.files[self.count]
|
||||
|
||||
if self.video_flag[self.count]:
|
||||
# Read video
|
||||
self.mode = 'video'
|
||||
ret_val, img0 = self.cap.read()
|
||||
if not ret_val:
|
||||
self.count += 1
|
||||
self.cap.release()
|
||||
if self.count == self.nf: # last video
|
||||
raise StopIteration
|
||||
else:
|
||||
path = self.files[self.count]
|
||||
self.new_video(path)
|
||||
ret_val, img0 = self.cap.read()
|
||||
|
||||
self.frame += 1
|
||||
print('video %g/%g (%g/%g) %s: ' % (self.count + 1, self.nf, self.frame, self.nframes, path), end='')
|
||||
|
||||
else:
|
||||
# Read image
|
||||
self.count += 1
|
||||
img0 = cv2.imread(path) # BGR
|
||||
assert img0 is not None, 'Image Not Found ' + path
|
||||
print('image %g/%g %s: ' % (self.count, self.nf, path), end='')
|
||||
|
||||
# Padded resize
|
||||
img = letterbox(img0, new_shape=self.img_size)[0]
|
||||
|
||||
# Convert
|
||||
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
||||
img = np.ascontiguousarray(img)
|
||||
|
||||
# cv2.imwrite(path + '.letterbox.jpg', 255 * img.transpose((1, 2, 0))[:, :, ::-1]) # save letterbox image
|
||||
return path, img, img0, self.cap
|
||||
|
||||
def new_video(self, path):
|
||||
self.frame = 0
|
||||
self.cap = cv2.VideoCapture(path)
|
||||
self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
def __len__(self):
|
||||
return self.nf # number of files
|
||||
|
||||
|
||||
class LoadWebcam: # for inference
|
||||
def __init__(self, pipe=0, img_size=640):
|
||||
self.img_size = img_size
|
||||
|
||||
if pipe == '0':
|
||||
pipe = 0 # local camera
|
||||
# pipe = 'rtsp://192.168.1.64/1' # IP camera
|
||||
# pipe = 'rtsp://username:password@192.168.1.64/1' # IP camera with login
|
||||
# pipe = 'rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa' # IP traffic camera
|
||||
# pipe = 'http://wmccpinetop.axiscam.net/mjpg/video.mjpg' # IP golf camera
|
||||
|
||||
# https://answers.opencv.org/question/215996/changing-gstreamer-pipeline-to-opencv-in-pythonsolved/
|
||||
# pipe = '"rtspsrc location="rtsp://username:password@192.168.1.64/1" latency=10 ! appsink' # GStreamer
|
||||
|
||||
# https://answers.opencv.org/question/200787/video-acceleration-gstremer-pipeline-in-videocapture/
|
||||
# https://stackoverflow.com/questions/54095699/install-gstreamer-support-for-opencv-python-package # install help
|
||||
# pipe = "rtspsrc location=rtsp://root:root@192.168.0.91:554/axis-media/media.amp?videocodec=h264&resolution=3840x2160 protocols=GST_RTSP_LOWER_TRANS_TCP ! rtph264depay ! queue ! vaapih264dec ! videoconvert ! appsink" # GStreamer
|
||||
|
||||
self.pipe = pipe
|
||||
self.cap = cv2.VideoCapture(pipe) # video capture object
|
||||
self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer size
|
||||
|
||||
def __iter__(self):
|
||||
self.count = -1
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
self.count += 1
|
||||
if cv2.waitKey(1) == ord('q'): # q to quit
|
||||
self.cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
raise StopIteration
|
||||
|
||||
# Read frame
|
||||
if self.pipe == 0: # local camera
|
||||
ret_val, img0 = self.cap.read()
|
||||
img0 = cv2.flip(img0, 1) # flip left-right
|
||||
else: # IP camera
|
||||
n = 0
|
||||
while True:
|
||||
n += 1
|
||||
self.cap.grab()
|
||||
if n % 30 == 0: # skip frames
|
||||
ret_val, img0 = self.cap.retrieve()
|
||||
if ret_val:
|
||||
break
|
||||
|
||||
# Print
|
||||
assert ret_val, 'Camera Error %s' % self.pipe
|
||||
img_path = 'webcam.jpg'
|
||||
print('webcam %g: ' % self.count, end='')
|
||||
|
||||
# Padded resize
|
||||
img = letterbox(img0, new_shape=self.img_size)[0]
|
||||
|
||||
# Convert
|
||||
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
||||
img = np.ascontiguousarray(img)
|
||||
|
||||
return img_path, img, img0, None
|
||||
|
||||
def __len__(self):
|
||||
return 0
|
||||
|
||||
|
||||
class LoadStreams: # multiple IP or RTSP cameras
|
||||
def __init__(self, sources='streams.txt', img_size=640):
|
||||
self.mode = 'images'
|
||||
self.img_size = img_size
|
||||
|
||||
if os.path.isfile(sources):
|
||||
with open(sources, 'r') as f:
|
||||
sources = [x.strip() for x in f.read().splitlines() if len(x.strip())]
|
||||
else:
|
||||
sources = [sources]
|
||||
|
||||
n = len(sources)
|
||||
self.imgs = [None] * n
|
||||
self.sources = sources
|
||||
for i, s in enumerate(sources):
|
||||
# Start the thread to read frames from the video stream
|
||||
print('%g/%g: %s... ' % (i + 1, n, s), end='')
|
||||
cap = cv2.VideoCapture(0 if s == '0' else s)
|
||||
assert cap.isOpened(), 'Failed to open %s' % s
|
||||
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
fps = cap.get(cv2.CAP_PROP_FPS) % 100
|
||||
_, self.imgs[i] = cap.read() # guarantee first frame
|
||||
thread = Thread(target=self.update, args=([i, cap]), daemon=True)
|
||||
print(' success (%gx%g at %.2f FPS).' % (w, h, fps))
|
||||
thread.start()
|
||||
print('') # newline
|
||||
|
||||
# check for common shapes
|
||||
s = np.stack([letterbox(x, new_shape=self.img_size)[0].shape for x in self.imgs], 0) # inference shapes
|
||||
self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal
|
||||
if not self.rect:
|
||||
print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.')
|
||||
|
||||
def update(self, index, cap):
|
||||
# Read next stream frame in a daemon thread
|
||||
n = 0
|
||||
while cap.isOpened():
|
||||
n += 1
|
||||
# _, self.imgs[index] = cap.read()
|
||||
cap.grab()
|
||||
if n == 4: # read every 4th frame
|
||||
_, self.imgs[index] = cap.retrieve()
|
||||
n = 0
|
||||
time.sleep(0.01) # wait time
|
||||
|
||||
def __iter__(self):
|
||||
self.count = -1
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
self.count += 1
|
||||
img0 = self.imgs.copy()
|
||||
if cv2.waitKey(1) == ord('q'): # q to quit
|
||||
cv2.destroyAllWindows()
|
||||
raise StopIteration
|
||||
|
||||
# Letterbox
|
||||
img = [letterbox(x, new_shape=self.img_size, auto=self.rect)[0] for x in img0]
|
||||
|
||||
# Stack
|
||||
img = np.stack(img, 0)
|
||||
|
||||
# Convert
|
||||
img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416
|
||||
img = np.ascontiguousarray(img)
|
||||
|
||||
return self.sources, img, img0, None
|
||||
|
||||
def __len__(self):
|
||||
return 0 # 1E12 frames = 32 streams at 30 FPS for 30 years
|
||||
|
||||
|
||||
class LoadImagesAndLabels(Dataset): # for training/testing
|
||||
def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
|
||||
cache_images=False, single_cls=False, stride=32, pad=0.0):
|
||||
try:
|
||||
f = [] # image files
|
||||
for p in path if isinstance(path, list) else [path]:
|
||||
p = str(Path(p)) # os-agnostic
|
||||
parent = str(Path(p).parent) + os.sep
|
||||
if os.path.isfile(p): # file
|
||||
with open(p, 'r') as t:
|
||||
t = t.read().splitlines()
|
||||
f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
|
||||
elif os.path.isdir(p): # folder
|
||||
f += glob.iglob(p + os.sep + '*.*')
|
||||
else:
|
||||
raise Exception('%s does not exist' % p)
|
||||
self.img_files = sorted(
|
||||
[x.replace('/', os.sep) for x in f if os.path.splitext(x)[-1].lower() in img_formats])
|
||||
except Exception as e:
|
||||
raise Exception('Error loading data from %s: %s\nSee %s' % (path, e, help_url))
|
||||
|
||||
n = len(self.img_files)
|
||||
assert n > 0, 'No images found in %s. See %s' % (path, help_url)
|
||||
bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index
|
||||
nb = bi[-1] + 1 # number of batches
|
||||
|
||||
self.n = n # number of images
|
||||
self.batch = bi # batch index of image
|
||||
self.img_size = img_size
|
||||
self.augment = augment
|
||||
self.hyp = hyp
|
||||
self.image_weights = image_weights
|
||||
self.rect = False if image_weights else rect
|
||||
self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)
|
||||
self.mosaic_border = [-img_size // 2, -img_size // 2]
|
||||
self.stride = stride
|
||||
|
||||
# Define labels
|
||||
self.label_files = [x.replace('images', 'labels').replace(os.path.splitext(x)[-1], '.txt') for x in
|
||||
self.img_files]
|
||||
|
||||
# Check cache
|
||||
cache_path = str(Path(self.label_files[0]).parent) + '.cache' # cached labels
|
||||
if os.path.isfile(cache_path):
|
||||
cache = torch.load(cache_path) # load
|
||||
if cache['hash'] != get_hash(self.label_files + self.img_files): # dataset changed
|
||||
cache = self.cache_labels(cache_path) # re-cache
|
||||
else:
|
||||
cache = self.cache_labels(cache_path) # cache
|
||||
|
||||
# Get labels
|
||||
labels, shapes = zip(*[cache[x] for x in self.img_files])
|
||||
self.shapes = np.array(shapes, dtype=np.float64)
|
||||
self.labels = list(labels)
|
||||
|
||||
# Rectangular Training https://github.com/ultralytics/yolov3/issues/232
|
||||
if self.rect:
|
||||
# Sort by aspect ratio
|
||||
s = self.shapes # wh
|
||||
ar = s[:, 1] / s[:, 0] # aspect ratio
|
||||
irect = ar.argsort()
|
||||
self.img_files = [self.img_files[i] for i in irect]
|
||||
self.label_files = [self.label_files[i] for i in irect]
|
||||
self.labels = [self.labels[i] for i in irect]
|
||||
self.shapes = s[irect] # wh
|
||||
ar = ar[irect]
|
||||
|
||||
# Set training image shapes
|
||||
shapes = [[1, 1]] * nb
|
||||
for i in range(nb):
|
||||
ari = ar[bi == i]
|
||||
mini, maxi = ari.min(), ari.max()
|
||||
if maxi < 1:
|
||||
shapes[i] = [maxi, 1]
|
||||
elif mini > 1:
|
||||
shapes[i] = [1, 1 / mini]
|
||||
|
||||
self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
|
||||
|
||||
# Cache labels
|
||||
create_datasubset, extract_bounding_boxes, labels_loaded = False, False, False
|
||||
nm, nf, ne, ns, nd = 0, 0, 0, 0, 0 # number missing, found, empty, datasubset, duplicate
|
||||
pbar = tqdm(self.label_files)
|
||||
for i, file in enumerate(pbar):
|
||||
l = self.labels[i] # label
|
||||
if l.shape[0]:
|
||||
assert l.shape[1] == 5, '> 5 label columns: %s' % file
|
||||
assert (l >= 0).all(), 'negative labels: %s' % file
|
||||
assert (l[:, 1:] <= 1).all(), 'non-normalized or out of bounds coordinate labels: %s' % file
|
||||
if np.unique(l, axis=0).shape[0] < l.shape[0]: # duplicate rows
|
||||
nd += 1 # print('WARNING: duplicate rows in %s' % self.label_files[i]) # duplicate rows
|
||||
if single_cls:
|
||||
l[:, 0] = 0 # force dataset into single-class mode
|
||||
self.labels[i] = l
|
||||
nf += 1 # file found
|
||||
|
||||
# Create subdataset (a smaller dataset)
|
||||
if create_datasubset and ns < 1E4:
|
||||
if ns == 0:
|
||||
create_folder(path='./datasubset')
|
||||
os.makedirs('./datasubset/images')
|
||||
exclude_classes = 43
|
||||
if exclude_classes not in l[:, 0]:
|
||||
ns += 1
|
||||
# shutil.copy(src=self.img_files[i], dst='./datasubset/images/') # copy image
|
||||
with open('./datasubset/images.txt', 'a') as f:
|
||||
f.write(self.img_files[i] + '\n')
|
||||
|
||||
# Extract object detection boxes for a second stage classifier
|
||||
if extract_bounding_boxes:
|
||||
p = Path(self.img_files[i])
|
||||
img = cv2.imread(str(p))
|
||||
h, w = img.shape[:2]
|
||||
for j, x in enumerate(l):
|
||||
f = '%s%sclassifier%s%g_%g_%s' % (p.parent.parent, os.sep, os.sep, x[0], j, p.name)
|
||||
if not os.path.exists(Path(f).parent):
|
||||
os.makedirs(Path(f).parent) # make new output folder
|
||||
|
||||
b = x[1:] * [w, h, w, h] # box
|
||||
b[2:] = b[2:].max() # rectangle to square
|
||||
b[2:] = b[2:] * 1.3 + 30 # pad
|
||||
b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(np.int)
|
||||
|
||||
b[[0, 2]] = np.clip(b[[0, 2]], 0, w) # clip boxes outside of image
|
||||
b[[1, 3]] = np.clip(b[[1, 3]], 0, h)
|
||||
assert cv2.imwrite(f, img[b[1]:b[3], b[0]:b[2]]), 'Failure extracting classifier boxes'
|
||||
else:
|
||||
ne += 1 # print('empty labels for image %s' % self.img_files[i]) # file empty
|
||||
# os.system("rm '%s' '%s'" % (self.img_files[i], self.label_files[i])) # remove
|
||||
|
||||
pbar.desc = 'Scanning labels %s (%g found, %g missing, %g empty, %g duplicate, for %g images)' % (
|
||||
cache_path, nf, nm, ne, nd, n)
|
||||
if nf == 0:
|
||||
s = 'WARNING: No labels found in %s. See %s' % (os.path.dirname(file) + os.sep, help_url)
|
||||
print(s)
|
||||
assert not augment, '%s. Can not train without labels.' % s
|
||||
|
||||
# Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)
|
||||
self.imgs = [None] * n
|
||||
if cache_images:
|
||||
gb = 0 # Gigabytes of cached images
|
||||
pbar = tqdm(range(len(self.img_files)), desc='Caching images')
|
||||
self.img_hw0, self.img_hw = [None] * n, [None] * n
|
||||
for i in pbar: # max 10k images
|
||||
self.imgs[i], self.img_hw0[i], self.img_hw[i] = load_image(self, i) # img, hw_original, hw_resized
|
||||
gb += self.imgs[i].nbytes
|
||||
pbar.desc = 'Caching images (%.1fGB)' % (gb / 1E9)
|
||||
|
||||
def cache_labels(self, path='labels.cache'):
|
||||
# Cache dataset labels, check images and read shapes
|
||||
x = {} # dict
|
||||
pbar = tqdm(zip(self.img_files, self.label_files), desc='Scanning images', total=len(self.img_files))
|
||||
for (img, label) in pbar:
|
||||
try:
|
||||
l = []
|
||||
image = Image.open(img)
|
||||
image.verify() # PIL verify
|
||||
# _ = io.imread(img) # skimage verify (from skimage import io)
|
||||
shape = exif_size(image) # image size
|
||||
assert (shape[0] > 9) & (shape[1] > 9), 'image size <10 pixels'
|
||||
if os.path.isfile(label):
|
||||
with open(label, 'r') as f:
|
||||
l = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32) # labels
|
||||
if len(l) == 0:
|
||||
l = np.zeros((0, 5), dtype=np.float32)
|
||||
x[img] = [l, shape]
|
||||
except Exception as e:
|
||||
x[img] = None
|
||||
print('WARNING: %s: %s' % (img, e))
|
||||
|
||||
x['hash'] = get_hash(self.label_files + self.img_files)
|
||||
torch.save(x, path) # save for next time
|
||||
return x
|
||||
|
||||
def __len__(self):
|
||||
return len(self.img_files)
|
||||
|
||||
# def __iter__(self):
|
||||
# self.count = -1
|
||||
# print('ran dataset iter')
|
||||
# #self.shuffled_vector = np.random.permutation(self.nF) if self.augment else np.arange(self.nF)
|
||||
# return self
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.image_weights:
|
||||
index = self.indices[index]
|
||||
|
||||
hyp = self.hyp
|
||||
if self.mosaic:
|
||||
# Load mosaic
|
||||
img, labels = load_mosaic(self, index)
|
||||
shapes = None
|
||||
|
||||
# MixUp https://arxiv.org/pdf/1710.09412.pdf
|
||||
if random.random() < hyp['mixup']:
|
||||
img2, labels2 = load_mosaic(self, random.randint(0, len(self.labels) - 1))
|
||||
r = np.random.beta(8.0, 8.0) # mixup ratio, alpha=beta=8.0
|
||||
img = (img * r + img2 * (1 - r)).astype(np.uint8)
|
||||
labels = np.concatenate((labels, labels2), 0)
|
||||
|
||||
else:
|
||||
# Load image
|
||||
img, (h0, w0), (h, w) = load_image(self, index)
|
||||
|
||||
# Letterbox
|
||||
shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
|
||||
img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
|
||||
shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
|
||||
|
||||
# Load labels
|
||||
labels = []
|
||||
x = self.labels[index]
|
||||
if x.size > 0:
|
||||
# Normalized xywh to pixel xyxy format
|
||||
labels = x.copy()
|
||||
labels[:, 1] = ratio[0] * w * (x[:, 1] - x[:, 3] / 2) + pad[0] # pad width
|
||||
labels[:, 2] = ratio[1] * h * (x[:, 2] - x[:, 4] / 2) + pad[1] # pad height
|
||||
labels[:, 3] = ratio[0] * w * (x[:, 1] + x[:, 3] / 2) + pad[0]
|
||||
labels[:, 4] = ratio[1] * h * (x[:, 2] + x[:, 4] / 2) + pad[1]
|
||||
|
||||
if self.augment:
|
||||
# Augment imagespace
|
||||
if not self.mosaic:
|
||||
img, labels = random_perspective(img, labels,
|
||||
degrees=hyp['degrees'],
|
||||
translate=hyp['translate'],
|
||||
scale=hyp['scale'],
|
||||
shear=hyp['shear'],
|
||||
perspective=hyp['perspective'])
|
||||
|
||||
# Augment colorspace
|
||||
augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
|
||||
|
||||
# Apply cutouts
|
||||
# if random.random() < 0.9:
|
||||
# labels = cutout(img, labels)
|
||||
|
||||
nL = len(labels) # number of labels
|
||||
if nL:
|
||||
labels[:, 1:5] = xyxy2xywh(labels[:, 1:5]) # convert xyxy to xywh
|
||||
labels[:, [2, 4]] /= img.shape[0] # normalized height 0-1
|
||||
labels[:, [1, 3]] /= img.shape[1] # normalized width 0-1
|
||||
|
||||
if self.augment:
|
||||
# flip up-down
|
||||
if random.random() < hyp['flipud']:
|
||||
img = np.flipud(img)
|
||||
if nL:
|
||||
labels[:, 2] = 1 - labels[:, 2]
|
||||
|
||||
# flip left-right
|
||||
if random.random() < hyp['fliplr']:
|
||||
img = np.fliplr(img)
|
||||
if nL:
|
||||
labels[:, 1] = 1 - labels[:, 1]
|
||||
|
||||
labels_out = torch.zeros((nL, 6))
|
||||
if nL:
|
||||
labels_out[:, 1:] = torch.from_numpy(labels)
|
||||
|
||||
# Convert
|
||||
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
|
||||
img = np.ascontiguousarray(img)
|
||||
|
||||
return torch.from_numpy(img), labels_out, self.img_files[index], shapes
|
||||
|
||||
@staticmethod
|
||||
def collate_fn(batch):
|
||||
img, label, path, shapes = zip(*batch) # transposed
|
||||
for i, l in enumerate(label):
|
||||
l[:, 0] = i # add target image index for build_targets()
|
||||
return torch.stack(img, 0), torch.cat(label, 0), path, shapes
|
||||
|
||||
|
||||
# Ancillary functions --------------------------------------------------------------------------------------------------
|
||||
def load_image(self, index):
|
||||
# loads 1 image from dataset, returns img, original hw, resized hw
|
||||
img = self.imgs[index]
|
||||
if img is None: # not cached
|
||||
path = self.img_files[index]
|
||||
img = cv2.imread(path) # BGR
|
||||
assert img is not None, 'Image Not Found ' + path
|
||||
h0, w0 = img.shape[:2] # orig hw
|
||||
r = self.img_size / max(h0, w0) # resize image to img_size
|
||||
if r != 1: # always resize down, only resize up if training with augmentation
|
||||
interp = cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR
|
||||
img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp)
|
||||
return img, (h0, w0), img.shape[:2] # img, hw_original, hw_resized
|
||||
else:
|
||||
return self.imgs[index], self.img_hw0[index], self.img_hw[index] # img, hw_original, hw_resized
|
||||
|
||||
|
||||
def augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):
|
||||
r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1 # random gains
|
||||
hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
|
||||
dtype = img.dtype # uint8
|
||||
|
||||
x = np.arange(0, 256, dtype=np.int16)
|
||||
lut_hue = ((x * r[0]) % 180).astype(dtype)
|
||||
lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
|
||||
lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
|
||||
|
||||
img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype)
|
||||
cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) # no return needed
|
||||
|
||||
# Histogram equalization
|
||||
# if random.random() < 0.2:
|
||||
# for i in range(3):
|
||||
# img[:, :, i] = cv2.equalizeHist(img[:, :, i])
|
||||
|
||||
|
||||
def load_mosaic(self, index):
|
||||
# loads images in a mosaic
|
||||
|
||||
labels4 = []
|
||||
s = self.img_size
|
||||
yc, xc = s, s # mosaic center x, y
|
||||
indices = [index] + [random.randint(0, len(self.labels) - 1) for _ in range(3)] # 3 additional image indices
|
||||
for i, index in enumerate(indices):
|
||||
# Load image
|
||||
img, _, (h, w) = load_image(self, index)
|
||||
|
||||
# place img in img4
|
||||
if i == 0: # top left
|
||||
img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
|
||||
x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)
|
||||
x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)
|
||||
elif i == 1: # top right
|
||||
x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
|
||||
x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
|
||||
elif i == 2: # bottom left
|
||||
x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
|
||||
x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, max(xc, w), min(y2a - y1a, h)
|
||||
elif i == 3: # bottom right
|
||||
x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
|
||||
x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
|
||||
|
||||
img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
|
||||
padw = x1a - x1b
|
||||
padh = y1a - y1b
|
||||
|
||||
# Labels
|
||||
x = self.labels[index]
|
||||
labels = x.copy()
|
||||
if x.size > 0: # Normalized xywh to pixel xyxy format
|
||||
labels[:, 1] = w * (x[:, 1] - x[:, 3] / 2) + padw
|
||||
labels[:, 2] = h * (x[:, 2] - x[:, 4] / 2) + padh
|
||||
labels[:, 3] = w * (x[:, 1] + x[:, 3] / 2) + padw
|
||||
labels[:, 4] = h * (x[:, 2] + x[:, 4] / 2) + padh
|
||||
labels4.append(labels)
|
||||
|
||||
# Concat/clip labels
|
||||
if len(labels4):
|
||||
labels4 = np.concatenate(labels4, 0)
|
||||
# np.clip(labels4[:, 1:] - s / 2, 0, s, out=labels4[:, 1:]) # use with center crop
|
||||
np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:]) # use with random_affine
|
||||
|
||||
# Replicate
|
||||
# img4, labels4 = replicate(img4, labels4)
|
||||
|
||||
# Augment
|
||||
# img4 = img4[s // 2: int(s * 1.5), s // 2:int(s * 1.5)] # center crop (WARNING, requires box pruning)
|
||||
img4, labels4 = random_perspective(img4, labels4,
|
||||
degrees=self.hyp['degrees'],
|
||||
translate=self.hyp['translate'],
|
||||
scale=self.hyp['scale'],
|
||||
shear=self.hyp['shear'],
|
||||
perspective=self.hyp['perspective'],
|
||||
border=self.mosaic_border) # border to remove
|
||||
|
||||
return img4, labels4
|
||||
|
||||
|
||||
def replicate(img, labels):
|
||||
# Replicate labels
|
||||
h, w = img.shape[:2]
|
||||
boxes = labels[:, 1:].astype(int)
|
||||
x1, y1, x2, y2 = boxes.T
|
||||
s = ((x2 - x1) + (y2 - y1)) / 2 # side length (pixels)
|
||||
for i in s.argsort()[:round(s.size * 0.5)]: # smallest indices
|
||||
x1b, y1b, x2b, y2b = boxes[i]
|
||||
bh, bw = y2b - y1b, x2b - x1b
|
||||
yc, xc = int(random.uniform(0, h - bh)), int(random.uniform(0, w - bw)) # offset x, y
|
||||
x1a, y1a, x2a, y2a = [xc, yc, xc + bw, yc + bh]
|
||||
img[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
|
||||
labels = np.append(labels, [[labels[i, 0], x1a, y1a, x2a, y2a]], axis=0)
|
||||
|
||||
return img, labels
|
||||
|
||||
|
||||
def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):
|
||||
# Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
|
||||
shape = img.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 test 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, 64), np.mod(dh, 64) # 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
|
||||
img = cv2.resize(img, 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))
|
||||
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
|
||||
return img, ratio, (dw, dh)
|
||||
|
||||
|
||||
def random_perspective(img, targets=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0, border=(0, 0)):
|
||||
# torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
|
||||
# targets = [cls, xyxy]
|
||||
|
||||
height = img.shape[0] + border[0] * 2 # shape(h,w,c)
|
||||
width = img.shape[1] + border[1] * 2
|
||||
|
||||
# Center
|
||||
C = np.eye(3)
|
||||
C[0, 2] = -img.shape[1] / 2 # x translation (pixels)
|
||||
C[1, 2] = -img.shape[0] / 2 # y translation (pixels)
|
||||
|
||||
# Perspective
|
||||
P = np.eye(3)
|
||||
P[2, 0] = random.uniform(-perspective, perspective) # x perspective (about y)
|
||||
P[2, 1] = random.uniform(-perspective, perspective) # y perspective (about x)
|
||||
|
||||
# Rotation and Scale
|
||||
R = np.eye(3)
|
||||
a = random.uniform(-degrees, degrees)
|
||||
# a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotations
|
||||
s = random.uniform(1 - scale, 1 + scale)
|
||||
# s = 2 ** random.uniform(-scale, scale)
|
||||
R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)
|
||||
|
||||
# Shear
|
||||
S = np.eye(3)
|
||||
S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # x shear (deg)
|
||||
S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # y shear (deg)
|
||||
|
||||
# Translation
|
||||
T = np.eye(3)
|
||||
T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * width # x translation (pixels)
|
||||
T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * height # y translation (pixels)
|
||||
|
||||
# Combined rotation matrix
|
||||
M = T @ S @ R @ P @ C # order of operations (right to left) is IMPORTANT
|
||||
if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changed
|
||||
if perspective:
|
||||
img = cv2.warpPerspective(img, M, dsize=(width, height), borderValue=(114, 114, 114))
|
||||
else: # affine
|
||||
img = cv2.warpAffine(img, M[:2], dsize=(width, height), borderValue=(114, 114, 114))
|
||||
|
||||
# Visualize
|
||||
# import matplotlib.pyplot as plt
|
||||
# ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel()
|
||||
# ax[0].imshow(img[:, :, ::-1]) # base
|
||||
# ax[1].imshow(img2[:, :, ::-1]) # warped
|
||||
|
||||
# Transform label coordinates
|
||||
n = len(targets)
|
||||
if n:
|
||||
# warp points
|
||||
xy = np.ones((n * 4, 3))
|
||||
xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1
|
||||
xy = xy @ M.T # transform
|
||||
if perspective:
|
||||
xy = (xy[:, :2] / xy[:, 2:3]).reshape(n, 8) # rescale
|
||||
else: # affine
|
||||
xy = xy[:, :2].reshape(n, 8)
|
||||
|
||||
# create new boxes
|
||||
x = xy[:, [0, 2, 4, 6]]
|
||||
y = xy[:, [1, 3, 5, 7]]
|
||||
xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
|
||||
|
||||
# # apply angle-based reduction of bounding boxes
|
||||
# radians = a * math.pi / 180
|
||||
# reduction = max(abs(math.sin(radians)), abs(math.cos(radians))) ** 0.5
|
||||
# x = (xy[:, 2] + xy[:, 0]) / 2
|
||||
# y = (xy[:, 3] + xy[:, 1]) / 2
|
||||
# w = (xy[:, 2] - xy[:, 0]) * reduction
|
||||
# h = (xy[:, 3] - xy[:, 1]) * reduction
|
||||
# xy = np.concatenate((x - w / 2, y - h / 2, x + w / 2, y + h / 2)).reshape(4, n).T
|
||||
|
||||
# clip boxes
|
||||
xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)
|
||||
xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)
|
||||
|
||||
# filter candidates
|
||||
i = box_candidates(box1=targets[:, 1:5].T * s, box2=xy.T)
|
||||
targets = targets[i]
|
||||
targets[:, 1:5] = xy[i]
|
||||
|
||||
return img, targets
|
||||
|
||||
|
||||
def box_candidates(box1, box2, wh_thr=2, ar_thr=20, area_thr=0.2): # box1(4,n), box2(4,n)
|
||||
# Compute candidate boxes: box1 before augment, box2 after augment, wh_thr (pixels), aspect_ratio_thr, area_ratio
|
||||
w1, h1 = box1[2] - box1[0], box1[3] - box1[1]
|
||||
w2, h2 = box2[2] - box2[0], box2[3] - box2[1]
|
||||
ar = np.maximum(w2 / (h2 + 1e-16), h2 / (w2 + 1e-16)) # aspect ratio
|
||||
return (w2 > wh_thr) & (h2 > wh_thr) & (w2 * h2 / (w1 * h1 + 1e-16) > area_thr) & (ar < ar_thr) # candidates
|
||||
|
||||
|
||||
def cutout(image, labels):
|
||||
# Applies image cutout augmentation https://arxiv.org/abs/1708.04552
|
||||
h, w = image.shape[:2]
|
||||
|
||||
def bbox_ioa(box1, box2):
|
||||
# Returns the intersection over box2 area given box1, box2. box1 is 4, box2 is nx4. boxes are x1y1x2y2
|
||||
box2 = box2.transpose()
|
||||
|
||||
# Get the coordinates of bounding boxes
|
||||
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]
|
||||
|
||||
# Intersection area
|
||||
inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \
|
||||
(np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).clip(0)
|
||||
|
||||
# box2 area
|
||||
box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + 1e-16
|
||||
|
||||
# Intersection over box2 area
|
||||
return inter_area / box2_area
|
||||
|
||||
# create random masks
|
||||
scales = [0.5] * 1 + [0.25] * 2 + [0.125] * 4 + [0.0625] * 8 + [0.03125] * 16 # image size fraction
|
||||
for s in scales:
|
||||
mask_h = random.randint(1, int(h * s))
|
||||
mask_w = random.randint(1, int(w * s))
|
||||
|
||||
# box
|
||||
xmin = max(0, random.randint(0, w) - mask_w // 2)
|
||||
ymin = max(0, random.randint(0, h) - mask_h // 2)
|
||||
xmax = min(w, xmin + mask_w)
|
||||
ymax = min(h, ymin + mask_h)
|
||||
|
||||
# apply random color mask
|
||||
image[ymin:ymax, xmin:xmax] = [random.randint(64, 191) for _ in range(3)]
|
||||
|
||||
# return unobscured labels
|
||||
if len(labels) and s > 0.03:
|
||||
box = np.array([xmin, ymin, xmax, ymax], dtype=np.float32)
|
||||
ioa = bbox_ioa(box, labels[:, 1:5]) # intersection over area
|
||||
labels = labels[ioa < 0.60] # remove >60% obscured labels
|
||||
|
||||
return labels
|
||||
|
||||
|
||||
def reduce_img_size(path='path/images', img_size=1024): # from utils.datasets import *; reduce_img_size()
|
||||
# creates a new ./images_reduced folder with reduced size images of maximum size img_size
|
||||
path_new = path + '_reduced' # reduced images path
|
||||
create_folder(path_new)
|
||||
for f in tqdm(glob.glob('%s/*.*' % path)):
|
||||
try:
|
||||
img = cv2.imread(f)
|
||||
h, w = img.shape[:2]
|
||||
r = img_size / max(h, w) # size ratio
|
||||
if r < 1.0:
|
||||
img = cv2.resize(img, (int(w * r), int(h * r)), interpolation=cv2.INTER_AREA) # _LINEAR fastest
|
||||
fnew = f.replace(path, path_new) # .replace(Path(f).suffix, '.jpg')
|
||||
cv2.imwrite(fnew, img)
|
||||
except:
|
||||
print('WARNING: image failure %s' % f)
|
||||
|
||||
|
||||
def recursive_dataset2bmp(dataset='path/dataset_bmp'): # from utils.datasets import *; recursive_dataset2bmp()
|
||||
# Converts dataset to bmp (for faster training)
|
||||
formats = [x.lower() for x in img_formats] + [x.upper() for x in img_formats]
|
||||
for a, b, files in os.walk(dataset):
|
||||
for file in tqdm(files, desc=a):
|
||||
p = a + '/' + file
|
||||
s = Path(file).suffix
|
||||
if s == '.txt': # replace text
|
||||
with open(p, 'r') as f:
|
||||
lines = f.read()
|
||||
for f in formats:
|
||||
lines = lines.replace(f, '.bmp')
|
||||
with open(p, 'w') as f:
|
||||
f.write(lines)
|
||||
elif s in formats: # replace image
|
||||
cv2.imwrite(p.replace(s, '.bmp'), cv2.imread(p))
|
||||
if s != '.bmp':
|
||||
os.system("rm '%s'" % p)
|
||||
|
||||
|
||||
def imagelist2folder(path='path/images.txt'): # from utils.datasets import *; imagelist2folder()
|
||||
# Copies all the images in a text file (list of images) into a folder
|
||||
create_folder(path[:-4])
|
||||
with open(path, 'r') as f:
|
||||
for line in f.read().splitlines():
|
||||
os.system('cp "%s" %s' % (line, path[:-4]))
|
||||
print(line)
|
||||
|
||||
|
||||
def create_folder(path='./new'):
|
||||
# Create folder
|
||||
if os.path.exists(path):
|
||||
shutil.rmtree(path) # delete output folder
|
||||
os.makedirs(path) # make new output folder
|
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,99 @@
|
||||
# This file contains google utils: https://cloud.google.com/storage/docs/reference/libraries
|
||||
# pip install --upgrade google-cloud-storage
|
||||
# from google.cloud import storage
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def attempt_download(weights):
|
||||
# Attempt to download pretrained weights if not found locally
|
||||
weights = weights.strip().replace("'", '')
|
||||
msg = weights + ' missing, try downloading from https://drive.google.com/drive/folders/1Drs_Aiu7xx6S-ix95f9kNsA6ueKRpN2J'
|
||||
|
||||
r = 1 # return
|
||||
if len(weights) > 0 and not os.path.isfile(weights):
|
||||
d = {'yolov3-spp.pt': '1mM67oNw4fZoIOL1c8M3hHmj66d8e-ni_', # yolov3-spp.yaml
|
||||
'yolov5s.pt': '1R5T6rIyy3lLwgFXNms8whc-387H0tMQO', # yolov5s.yaml
|
||||
'yolov5m.pt': '1vobuEExpWQVpXExsJ2w-Mbf3HJjWkQJr', # yolov5m.yaml
|
||||
'yolov5l.pt': '1hrlqD1Wdei7UT4OgT785BEk1JwnSvNEV', # yolov5l.yaml
|
||||
'yolov5x.pt': '1mM8aZJlWTxOg7BZJvNUMrTnA2AbeCVzS', # yolov5x.yaml
|
||||
}
|
||||
|
||||
file = Path(weights).name
|
||||
if file in d:
|
||||
r = gdrive_download(id=d[file], name=weights)
|
||||
|
||||
if not (r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6): # weights exist and > 1MB
|
||||
os.remove(weights) if os.path.exists(weights) else None # remove partial downloads
|
||||
s = "curl -L -o %s 'storage.googleapis.com/ultralytics/yolov5/ckpt/%s'" % (weights, file)
|
||||
r = os.system(s) # execute, capture return values
|
||||
|
||||
# Error check
|
||||
if not (r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6): # weights exist and > 1MB
|
||||
os.remove(weights) if os.path.exists(weights) else None # remove partial downloads
|
||||
raise Exception(msg)
|
||||
|
||||
|
||||
def gdrive_download(id='1n_oKgR81BJtqk75b00eAjdv03qVCQn2f', name='coco128.zip'):
|
||||
# Downloads a file from Google Drive, accepting presented query
|
||||
# from utils.google_utils import *; gdrive_download()
|
||||
t = time.time()
|
||||
|
||||
print('Downloading https://drive.google.com/uc?export=download&id=%s as %s... ' % (id, name), end='')
|
||||
os.remove(name) if os.path.exists(name) else None # remove existing
|
||||
os.remove('cookie') if os.path.exists('cookie') else None
|
||||
|
||||
# Attempt file download
|
||||
os.system("curl -c ./cookie -s -L \"drive.google.com/uc?export=download&id=%s\" > /dev/null" % id)
|
||||
if os.path.exists('cookie'): # large file
|
||||
s = "curl -Lb ./cookie \"drive.google.com/uc?export=download&confirm=`awk '/download/ {print $NF}' ./cookie`&id=%s\" -o %s" % (
|
||||
id, name)
|
||||
else: # small file
|
||||
s = 'curl -s -L -o %s "drive.google.com/uc?export=download&id=%s"' % (name, id)
|
||||
r = os.system(s) # execute, capture return values
|
||||
os.remove('cookie') if os.path.exists('cookie') else None
|
||||
|
||||
# Error check
|
||||
if r != 0:
|
||||
os.remove(name) if os.path.exists(name) else None # remove partial
|
||||
print('Download error ') # raise Exception('Download error')
|
||||
return r
|
||||
|
||||
# Unzip if archive
|
||||
if name.endswith('.zip'):
|
||||
print('unzipping... ', end='')
|
||||
os.system('unzip -q %s' % name) # unzip
|
||||
os.remove(name) # remove zip to free space
|
||||
|
||||
print('Done (%.1fs)' % (time.time() - t))
|
||||
return r
|
||||
|
||||
|
||||
# def upload_blob(bucket_name, source_file_name, destination_blob_name):
|
||||
# # Uploads a file to a bucket
|
||||
# # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
|
||||
#
|
||||
# storage_client = storage.Client()
|
||||
# bucket = storage_client.get_bucket(bucket_name)
|
||||
# blob = bucket.blob(destination_blob_name)
|
||||
#
|
||||
# blob.upload_from_filename(source_file_name)
|
||||
#
|
||||
# print('File {} uploaded to {}.'.format(
|
||||
# source_file_name,
|
||||
# destination_blob_name))
|
||||
#
|
||||
#
|
||||
# def download_blob(bucket_name, source_blob_name, destination_file_name):
|
||||
# # Uploads a blob from a bucket
|
||||
# storage_client = storage.Client()
|
||||
# bucket = storage_client.get_bucket(bucket_name)
|
||||
# blob = bucket.blob(source_blob_name)
|
||||
#
|
||||
# blob.download_to_filename(destination_file_name)
|
||||
#
|
||||
# print('Blob {} downloaded to {}.'.format(
|
||||
# source_blob_name,
|
||||
# destination_file_name))
|
@ -0,0 +1,222 @@
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from copy import deepcopy
|
||||
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision.models as models
|
||||
|
||||
|
||||
def init_seeds(seed=0):
|
||||
torch.manual_seed(seed)
|
||||
|
||||
# Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
|
||||
if seed == 0: # slower, more reproducible
|
||||
cudnn.deterministic = True
|
||||
cudnn.benchmark = False
|
||||
else: # faster, less reproducible
|
||||
cudnn.deterministic = False
|
||||
cudnn.benchmark = True
|
||||
|
||||
|
||||
def select_device(device='', batch_size=None):
|
||||
# device = 'cpu' or '0' or '0,1,2,3'
|
||||
cpu_request = device.lower() == 'cpu'
|
||||
if device and not cpu_request: # if device requested other than 'cpu'
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
|
||||
assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device # check availablity
|
||||
|
||||
cuda = False if cpu_request else torch.cuda.is_available()
|
||||
if cuda:
|
||||
c = 1024 ** 2 # bytes to MB
|
||||
ng = torch.cuda.device_count()
|
||||
if ng > 1 and batch_size: # check that batch_size is compatible with device_count
|
||||
assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng)
|
||||
x = [torch.cuda.get_device_properties(i) for i in range(ng)]
|
||||
s = 'Using CUDA '
|
||||
for i in range(0, ng):
|
||||
if i == 1:
|
||||
s = ' ' * len(s)
|
||||
print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
|
||||
(s, i, x[i].name, x[i].total_memory / c))
|
||||
else:
|
||||
print('Using CPU')
|
||||
|
||||
print('') # skip a line
|
||||
return torch.device('cuda:0' if cuda else 'cpu')
|
||||
|
||||
|
||||
def time_synchronized():
|
||||
torch.cuda.synchronize() if torch.cuda.is_available() else None
|
||||
return time.time()
|
||||
|
||||
|
||||
def is_parallel(model):
|
||||
# is model is parallel with DP or DDP
|
||||
return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
|
||||
|
||||
|
||||
def initialize_weights(model):
|
||||
for m in model.modules():
|
||||
t = type(m)
|
||||
if t is nn.Conv2d:
|
||||
pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif t is nn.BatchNorm2d:
|
||||
m.eps = 1e-3
|
||||
m.momentum = 0.03
|
||||
elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
|
||||
m.inplace = True
|
||||
|
||||
|
||||
def find_modules(model, mclass=nn.Conv2d):
|
||||
# finds layer indices matching module class 'mclass'
|
||||
return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
|
||||
|
||||
|
||||
def sparsity(model):
|
||||
# Return global model sparsity
|
||||
a, b = 0., 0.
|
||||
for p in model.parameters():
|
||||
a += p.numel()
|
||||
b += (p == 0).sum()
|
||||
return b / a
|
||||
|
||||
|
||||
def prune(model, amount=0.3):
|
||||
# Prune model to requested global sparsity
|
||||
import torch.nn.utils.prune as prune
|
||||
print('Pruning model... ', end='')
|
||||
for name, m in model.named_modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
prune.l1_unstructured(m, name='weight', amount=amount) # prune
|
||||
prune.remove(m, 'weight') # make permanent
|
||||
print(' %.3g global sparsity' % sparsity(model))
|
||||
|
||||
|
||||
def fuse_conv_and_bn(conv, bn):
|
||||
# https://tehnokv.com/posts/fusing-batchnorm-and-conv/
|
||||
with torch.no_grad():
|
||||
# init
|
||||
fusedconv = nn.Conv2d(conv.in_channels,
|
||||
conv.out_channels,
|
||||
kernel_size=conv.kernel_size,
|
||||
stride=conv.stride,
|
||||
padding=conv.padding,
|
||||
bias=True).to(conv.weight.device)
|
||||
|
||||
# prepare filters
|
||||
w_conv = conv.weight.clone().view(conv.out_channels, -1)
|
||||
w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
|
||||
fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
|
||||
|
||||
# prepare spatial bias
|
||||
b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
|
||||
b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
|
||||
fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
|
||||
|
||||
return fusedconv
|
||||
|
||||
|
||||
def model_info(model, verbose=False):
|
||||
# Plots a line-by-line description of a PyTorch model
|
||||
n_p = sum(x.numel() for x in model.parameters()) # number parameters
|
||||
n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
|
||||
if verbose:
|
||||
print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
|
||||
for i, (name, p) in enumerate(model.named_parameters()):
|
||||
name = name.replace('module_list.', '')
|
||||
print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
|
||||
(i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
|
||||
|
||||
try: # FLOPS
|
||||
from thop import profile
|
||||
flops = profile(deepcopy(model), inputs=(torch.zeros(1, 3, 64, 64),), verbose=False)[0] / 1E9 * 2
|
||||
fs = ', %.1f GFLOPS' % (flops * 100) # 640x640 FLOPS
|
||||
except:
|
||||
fs = ''
|
||||
|
||||
print('Model Summary: %g layers, %g parameters, %g gradients%s' % (len(list(model.parameters())), n_p, n_g, fs))
|
||||
|
||||
|
||||
def load_classifier(name='resnet101', n=2):
|
||||
# Loads a pretrained model reshaped to n-class output
|
||||
model = models.__dict__[name](pretrained=True)
|
||||
|
||||
# Display model properties
|
||||
input_size = [3, 224, 224]
|
||||
input_space = 'RGB'
|
||||
input_range = [0, 1]
|
||||
mean = [0.485, 0.456, 0.406]
|
||||
std = [0.229, 0.224, 0.225]
|
||||
for x in [input_size, input_space, input_range, mean, std]:
|
||||
print(x + ' =', eval(x))
|
||||
|
||||
# Reshape output to n classes
|
||||
filters = model.fc.weight.shape[1]
|
||||
model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True)
|
||||
model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True)
|
||||
model.fc.out_features = n
|
||||
return model
|
||||
|
||||
|
||||
def scale_img(img, ratio=1.0, same_shape=False): # img(16,3,256,416), r=ratio
|
||||
# scales img(bs,3,y,x) by ratio
|
||||
if ratio == 1.0:
|
||||
return img
|
||||
else:
|
||||
h, w = img.shape[2:]
|
||||
s = (int(h * ratio), int(w * ratio)) # new size
|
||||
img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
|
||||
if not same_shape: # pad/crop img
|
||||
gs = 32 # (pixels) grid size
|
||||
h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
|
||||
return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
|
||||
|
||||
|
||||
def copy_attr(a, b, include=(), exclude=()):
|
||||
# Copy attributes from b to a, options to only include [...] and to exclude [...]
|
||||
for k, v in b.__dict__.items():
|
||||
if (len(include) and k not in include) or k.startswith('_') or k in exclude:
|
||||
continue
|
||||
else:
|
||||
setattr(a, k, v)
|
||||
|
||||
|
||||
class ModelEMA:
|
||||
""" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
|
||||
Keep a moving average of everything in the model state_dict (parameters and buffers).
|
||||
This is intended to allow functionality like
|
||||
https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
|
||||
A smoothed version of the weights is necessary for some training schemes to perform well.
|
||||
This class is sensitive where it is initialized in the sequence of model init,
|
||||
GPU assignment and distributed training wrappers.
|
||||
"""
|
||||
|
||||
def __init__(self, model, decay=0.9999, updates=0):
|
||||
# Create EMA
|
||||
self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA
|
||||
# if next(model.parameters()).device.type != 'cpu':
|
||||
# self.ema.half() # FP16 EMA
|
||||
self.updates = updates # number of EMA updates
|
||||
self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs)
|
||||
for p in self.ema.parameters():
|
||||
p.requires_grad_(False)
|
||||
|
||||
def update(self, model):
|
||||
# Update EMA parameters
|
||||
with torch.no_grad():
|
||||
self.updates += 1
|
||||
d = self.decay(self.updates)
|
||||
|
||||
msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict
|
||||
for k, v in self.ema.state_dict().items():
|
||||
if v.dtype.is_floating_point:
|
||||
v *= d
|
||||
v += (1. - d) * msd[k].detach()
|
||||
|
||||
def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
|
||||
# Update EMA attributes
|
||||
copy_attr(self.ema, model, include, exclude)
|
@ -0,0 +1,10 @@
|
||||
#!/bin/bash
|
||||
# Download common models
|
||||
|
||||
python -c "
|
||||
from utils.google_utils import *;
|
||||
attempt_download('weights/yolov5s.pt');
|
||||
attempt_download('weights/yolov5m.pt');
|
||||
attempt_download('weights/yolov5l.pt');
|
||||
attempt_download('weights/yolov5x.pt')
|
||||
"
|
@ -0,0 +1,26 @@
|
||||
|
||||
|
||||
add_library(gps
|
||||
Drivers/src/ashtech.cpp
|
||||
Drivers/src/gps_helper.cpp
|
||||
Drivers/src/mtk.cpp
|
||||
Drivers/src/rtcm.cpp
|
||||
Drivers/src/sbf.cpp
|
||||
Drivers/src/ubx.cpp
|
||||
GPSManager.cc
|
||||
GPSProvider.cc
|
||||
RTCM/RTCMMavlink.cc
|
||||
)
|
||||
|
||||
target_link_libraries(gps
|
||||
Qt6::Core
|
||||
Qt6::Location
|
||||
Qt6::SerialPort
|
||||
Qt6::Svg
|
||||
Qt6::TextToSpeech
|
||||
|
||||
qgc
|
||||
)
|
||||
|
||||
target_include_directories(gps INTERFACE ${CMAKE_CURRENT_SOURCE_DIR})
|
||||
|
@ -0,0 +1,100 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* (c) 2009-2020 QGROUNDCONTROL PROJECT <http://www.qgroundcontrol.org>
|
||||
*
|
||||
* QGroundControl is licensed according to the terms in the file
|
||||
* COPYING.md in the root of the source code directory.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
|
||||
#include "GPSManager.h"
|
||||
#include "QGCLoggingCategory.h"
|
||||
#include "QGCApplication.h"
|
||||
#include "SettingsManager.h"
|
||||
#include "RTKSettings.h"
|
||||
|
||||
GPSManager::GPSManager(QGCApplication* app, QGCToolbox* toolbox)
|
||||
: QGCTool(app, toolbox)
|
||||
{
|
||||
qRegisterMetaType<GPSPositionMessage>();
|
||||
qRegisterMetaType<GPSSatelliteMessage>();
|
||||
}
|
||||
|
||||
GPSManager::~GPSManager()
|
||||
{
|
||||
disconnectGPS();
|
||||
}
|
||||
|
||||
void GPSManager::connectGPS(const QString& device, const QString& gps_type)
|
||||
{
|
||||
RTKSettings* rtkSettings = qgcApp()->toolbox()->settingsManager()->rtkSettings();
|
||||
|
||||
GPSProvider::GPSType type;
|
||||
if (gps_type.contains("trimble", Qt::CaseInsensitive)) {
|
||||
type = GPSProvider::GPSType::trimble;
|
||||
qCDebug(RTKGPSLog) << "Connecting Trimble device";
|
||||
} else if (gps_type.contains("septentrio", Qt::CaseInsensitive)) {
|
||||
type = GPSProvider::GPSType::septentrio;
|
||||
qCDebug(RTKGPSLog) << "Connecting Septentrio device";
|
||||
} else {
|
||||
type = GPSProvider::GPSType::u_blox;
|
||||
qCDebug(RTKGPSLog) << "Connecting U-blox device";
|
||||
}
|
||||
|
||||
disconnectGPS();
|
||||
_requestGpsStop = false;
|
||||
_gpsProvider = new GPSProvider(device,
|
||||
type,
|
||||
true, /* enableSatInfo */
|
||||
rtkSettings->surveyInAccuracyLimit()->rawValue().toDouble(),
|
||||
rtkSettings->surveyInMinObservationDuration()->rawValue().toInt(),
|
||||
rtkSettings->useFixedBasePosition()->rawValue().toBool(),
|
||||
rtkSettings->fixedBasePositionLatitude()->rawValue().toDouble(),
|
||||
rtkSettings->fixedBasePositionLongitude()->rawValue().toDouble(),
|
||||
rtkSettings->fixedBasePositionAltitude()->rawValue().toFloat(),
|
||||
rtkSettings->fixedBasePositionAccuracy()->rawValue().toFloat(),
|
||||
_requestGpsStop);
|
||||
_gpsProvider->start();
|
||||
|
||||
//create RTCM device
|
||||
_rtcmMavlink = new RTCMMavlink(*_toolbox);
|
||||
|
||||
connect(_gpsProvider, &GPSProvider::RTCMDataUpdate, _rtcmMavlink, &RTCMMavlink::RTCMDataUpdate);
|
||||
|
||||
//test: connect to position update
|
||||
connect(_gpsProvider, &GPSProvider::positionUpdate, this, &GPSManager::GPSPositionUpdate);
|
||||
connect(_gpsProvider, &GPSProvider::satelliteInfoUpdate, this, &GPSManager::GPSSatelliteUpdate);
|
||||
connect(_gpsProvider, &GPSProvider::finished, this, &GPSManager::onDisconnect);
|
||||
connect(_gpsProvider, &GPSProvider::surveyInStatus, this, &GPSManager::surveyInStatus);
|
||||
|
||||
emit onConnect();
|
||||
}
|
||||
|
||||
void GPSManager::disconnectGPS(void)
|
||||
{
|
||||
if (_gpsProvider) {
|
||||
_requestGpsStop = true;
|
||||
//Note that we need a relatively high timeout to be sure the GPS thread finished.
|
||||
if (!_gpsProvider->wait(2000)) {
|
||||
qWarning() << "Failed to wait for GPS thread exit. Consider increasing the timeout";
|
||||
}
|
||||
delete(_gpsProvider);
|
||||
}
|
||||
if (_rtcmMavlink) {
|
||||
delete(_rtcmMavlink);
|
||||
}
|
||||
_gpsProvider = nullptr;
|
||||
_rtcmMavlink = nullptr;
|
||||
}
|
||||
|
||||
|
||||
void GPSManager::GPSPositionUpdate(GPSPositionMessage msg)
|
||||
{
|
||||
qCDebug(RTKGPSLog) << QString("GPS: got position update: alt=%1, long=%2, lat=%3").arg(msg.position_data.alt).arg(msg.position_data.lon).arg(msg.position_data.lat);
|
||||
}
|
||||
void GPSManager::GPSSatelliteUpdate(GPSSatelliteMessage msg)
|
||||
{
|
||||
qCDebug(RTKGPSLog) << QString("GPS: got satellite info update, %1 satellites").arg((int)msg.satellite_data.count);
|
||||
emit satelliteUpdate(msg.satellite_data.count);
|
||||
}
|
@ -0,0 +1,50 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* (c) 2009-2020 QGROUNDCONTROL PROJECT <http://www.qgroundcontrol.org>
|
||||
*
|
||||
* QGroundControl is licensed according to the terms in the file
|
||||
* COPYING.md in the root of the source code directory.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "GPSProvider.h"
|
||||
#include "RTCM/RTCMMavlink.h"
|
||||
#include <QGCToolbox.h>
|
||||
|
||||
#include <QString>
|
||||
#include <QObject>
|
||||
|
||||
/**
|
||||
** class GPSManager
|
||||
* handles a GPS provider and RTK
|
||||
*/
|
||||
class GPSManager : public QGCTool
|
||||
{
|
||||
Q_OBJECT
|
||||
public:
|
||||
GPSManager(QGCApplication* app, QGCToolbox* toolbox);
|
||||
~GPSManager();
|
||||
|
||||
void connectGPS (const QString& device, const QString& gps_type);
|
||||
void disconnectGPS (void);
|
||||
bool connected (void) const { return _gpsProvider && _gpsProvider->isRunning(); }
|
||||
|
||||
signals:
|
||||
void onConnect();
|
||||
void onDisconnect();
|
||||
void surveyInStatus(float duration, float accuracyMM, double latitude, double longitude, float altitude, bool valid, bool active);
|
||||
void satelliteUpdate(int numSats);
|
||||
|
||||
private slots:
|
||||
void GPSPositionUpdate(GPSPositionMessage msg);
|
||||
void GPSSatelliteUpdate(GPSSatelliteMessage msg);
|
||||
|
||||
private:
|
||||
GPSProvider* _gpsProvider = nullptr;
|
||||
RTCMMavlink* _rtcmMavlink = nullptr;
|
||||
|
||||
std::atomic_bool _requestGpsStop; ///< signals the thread to quit
|
||||
};
|
@ -0,0 +1,33 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* (c) 2009-2020 QGROUNDCONTROL PROJECT <http://www.qgroundcontrol.org>
|
||||
*
|
||||
* QGroundControl is licensed according to the terms in the file
|
||||
* COPYING.md in the root of the source code directory.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "vehicle_gps_position.h"
|
||||
#include "satellite_info.h"
|
||||
#include <QMetaType>
|
||||
|
||||
/**
|
||||
** struct GPSPositionMessage
|
||||
* wrapper that can be used for Qt signal/slots
|
||||
*/
|
||||
struct GPSPositionMessage
|
||||
{
|
||||
sensor_gps_s position_data;
|
||||
};
|
||||
|
||||
Q_DECLARE_METATYPE(GPSPositionMessage);
|
||||
|
||||
|
||||
struct GPSSatelliteMessage
|
||||
{
|
||||
satellite_info_s satellite_data;
|
||||
};
|
||||
Q_DECLARE_METATYPE(GPSSatelliteMessage);
|
@ -0,0 +1,234 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* (c) 2009-2020 QGROUNDCONTROL PROJECT <http://www.qgroundcontrol.org>
|
||||
*
|
||||
* QGroundControl is licensed according to the terms in the file
|
||||
* COPYING.md in the root of the source code directory.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
|
||||
#include "GPSProvider.h"
|
||||
#include "QGCLoggingCategory.h"
|
||||
#include "QGCApplication.h"
|
||||
#include "SettingsManager.h"
|
||||
|
||||
#define GPS_RECEIVE_TIMEOUT 1200
|
||||
|
||||
#include <QDebug>
|
||||
|
||||
#include "Drivers/src/ubx.h"
|
||||
#include "Drivers/src/sbf.h"
|
||||
#include "Drivers/src/ashtech.h"
|
||||
#include "Drivers/src/base_station.h"
|
||||
#include "definitions.h"
|
||||
|
||||
//#define SIMULATE_RTCM_OUTPUT //if defined, generate simulated RTCM messages
|
||||
//additionally make sure to call connectGPS(""), eg. from QGCToolbox.cc
|
||||
|
||||
|
||||
void GPSProvider::run()
|
||||
{
|
||||
#ifdef SIMULATE_RTCM_OUTPUT
|
||||
const int fakeMsgLengths[3] = { 30, 170, 240 };
|
||||
uint8_t* fakeData = new uint8_t[fakeMsgLengths[2]];
|
||||
while (!_requestStop) {
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
gotRTCMData((uint8_t*) fakeData, fakeMsgLengths[i]);
|
||||
msleep(4);
|
||||
}
|
||||
msleep(100);
|
||||
}
|
||||
delete[] fakeData;
|
||||
#endif /* SIMULATE_RTCM_OUTPUT */
|
||||
|
||||
if (_serial) delete _serial;
|
||||
|
||||
_serial = new QSerialPort();
|
||||
_serial->setPortName(_device);
|
||||
if (!_serial->open(QIODevice::ReadWrite)) {
|
||||
int retries = 60;
|
||||
// Give the device some time to come up. In some cases the device is not
|
||||
// immediately accessible right after startup for some reason. This can take 10-20s.
|
||||
while (retries-- > 0 && _serial->error() == QSerialPort::PermissionError) {
|
||||
qCDebug(RTKGPSLog) << "Cannot open device... retrying";
|
||||
msleep(500);
|
||||
if (_serial->open(QIODevice::ReadWrite)) {
|
||||
_serial->clearError();
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (_serial->error() != QSerialPort::NoError) {
|
||||
qWarning() << "GPS: Failed to open Serial Device" << _device << _serial->errorString();
|
||||
return;
|
||||
}
|
||||
}
|
||||
_serial->setBaudRate(QSerialPort::Baud9600);
|
||||
_serial->setDataBits(QSerialPort::Data8);
|
||||
_serial->setParity(QSerialPort::NoParity);
|
||||
_serial->setStopBits(QSerialPort::OneStop);
|
||||
_serial->setFlowControl(QSerialPort::NoFlowControl);
|
||||
|
||||
unsigned int baudrate;
|
||||
GPSBaseStationSupport* gpsDriver = nullptr;
|
||||
|
||||
while (!_requestStop) {
|
||||
|
||||
if (gpsDriver) {
|
||||
delete gpsDriver;
|
||||
gpsDriver = nullptr;
|
||||
}
|
||||
|
||||
if (_type == GPSType::trimble) {
|
||||
gpsDriver = new GPSDriverAshtech(&callbackEntry, this, &_reportGpsPos, _pReportSatInfo);
|
||||
baudrate = 115200;
|
||||
} else if (_type == GPSType::septentrio) {
|
||||
gpsDriver = new GPSDriverSBF(&callbackEntry, this, &_reportGpsPos, _pReportSatInfo, 5);
|
||||
baudrate = 0; // auto-configure
|
||||
} else {
|
||||
gpsDriver = new GPSDriverUBX(GPSDriverUBX::Interface::UART, &callbackEntry, this, &_reportGpsPos, _pReportSatInfo);
|
||||
baudrate = 0; // auto-configure
|
||||
}
|
||||
gpsDriver->setSurveyInSpecs(_surveyInAccMeters * 10000.0f, _surveryInDurationSecs);
|
||||
|
||||
if (_useFixedBaseLoction) {
|
||||
gpsDriver->setBasePosition(_fixedBaseLatitude, _fixedBaseLongitude, _fixedBaseAltitudeMeters, _fixedBaseAccuracyMeters * 1000.0f);
|
||||
}
|
||||
|
||||
_gpsConfig.output_mode = GPSHelper::OutputMode::RTCM;
|
||||
if (gpsDriver->configure(baudrate, _gpsConfig) == 0) {
|
||||
|
||||
/* reset report */
|
||||
memset(&_reportGpsPos, 0, sizeof(_reportGpsPos));
|
||||
|
||||
//In rare cases it can happen that we get an error from the driver (eg. checksum failure) due to
|
||||
//bus errors or buggy firmware. In this case we want to try multiple times before giving up.
|
||||
int numTries = 0;
|
||||
|
||||
while (!_requestStop && numTries < 3) {
|
||||
int helperRet = gpsDriver->receive(GPS_RECEIVE_TIMEOUT);
|
||||
|
||||
if (helperRet > 0) {
|
||||
numTries = 0;
|
||||
|
||||
if (helperRet & 1) {
|
||||
publishGPSPosition();
|
||||
numTries = 0;
|
||||
}
|
||||
|
||||
if (_pReportSatInfo && (helperRet & 2)) {
|
||||
publishGPSSatellite();
|
||||
numTries = 0;
|
||||
}
|
||||
} else {
|
||||
++numTries;
|
||||
}
|
||||
}
|
||||
if (_serial->error() != QSerialPort::NoError && _serial->error() != QSerialPort::TimeoutError) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
qCDebug(RTKGPSLog) << "Exiting GPS thread";
|
||||
}
|
||||
|
||||
GPSProvider::GPSProvider(const QString& device,
|
||||
GPSType type,
|
||||
bool enableSatInfo,
|
||||
double surveyInAccMeters,
|
||||
int surveryInDurationSecs,
|
||||
bool useFixedBaseLocation,
|
||||
double fixedBaseLatitude,
|
||||
double fixedBaseLongitude,
|
||||
float fixedBaseAltitudeMeters,
|
||||
float fixedBaseAccuracyMeters,
|
||||
const std::atomic_bool& requestStop)
|
||||
: _device (device)
|
||||
, _type (type)
|
||||
, _requestStop (requestStop)
|
||||
, _surveyInAccMeters (surveyInAccMeters)
|
||||
, _surveryInDurationSecs (surveryInDurationSecs)
|
||||
, _useFixedBaseLoction (useFixedBaseLocation)
|
||||
, _fixedBaseLatitude (fixedBaseLatitude)
|
||||
, _fixedBaseLongitude (fixedBaseLongitude)
|
||||
, _fixedBaseAltitudeMeters (fixedBaseAltitudeMeters)
|
||||
, _fixedBaseAccuracyMeters (fixedBaseAccuracyMeters)
|
||||
{
|
||||
qCDebug(RTKGPSLog) << "Survey in accuracy:duration" << surveyInAccMeters << surveryInDurationSecs;
|
||||
if (enableSatInfo) _pReportSatInfo = new satellite_info_s();
|
||||
}
|
||||
|
||||
GPSProvider::~GPSProvider()
|
||||
{
|
||||
if (_pReportSatInfo) delete(_pReportSatInfo);
|
||||
if (_serial) delete _serial;
|
||||
}
|
||||
|
||||
void GPSProvider::publishGPSPosition()
|
||||
{
|
||||
GPSPositionMessage msg;
|
||||
msg.position_data = _reportGpsPos;
|
||||
emit positionUpdate(msg);
|
||||
}
|
||||
|
||||
void GPSProvider::publishGPSSatellite()
|
||||
{
|
||||
GPSSatelliteMessage msg;
|
||||
msg.satellite_data = *_pReportSatInfo;
|
||||
emit satelliteInfoUpdate(msg);
|
||||
}
|
||||
|
||||
void GPSProvider::gotRTCMData(uint8_t* data, size_t len)
|
||||
{
|
||||
QByteArray message((char*)data, static_cast<int>(len));
|
||||
emit RTCMDataUpdate(message);
|
||||
}
|
||||
|
||||
int GPSProvider::callbackEntry(GPSCallbackType type, void *data1, int data2, void *user)
|
||||
{
|
||||
GPSProvider *gps = (GPSProvider *)user;
|
||||
return gps->callback(type, data1, data2);
|
||||
}
|
||||
|
||||
int GPSProvider::callback(GPSCallbackType type, void *data1, int data2)
|
||||
{
|
||||
switch (type) {
|
||||
case GPSCallbackType::readDeviceData: {
|
||||
if (_serial->bytesAvailable() == 0) {
|
||||
int timeout = *((int *) data1);
|
||||
if (!_serial->waitForReadyRead(timeout))
|
||||
return 0; //timeout
|
||||
}
|
||||
return (int)_serial->read((char*) data1, data2);
|
||||
}
|
||||
case GPSCallbackType::writeDeviceData:
|
||||
if (_serial->write((char*) data1, data2) >= 0) {
|
||||
if (_serial->waitForBytesWritten(-1))
|
||||
return data2;
|
||||
}
|
||||
return -1;
|
||||
|
||||
case GPSCallbackType::setBaudrate:
|
||||
return _serial->setBaudRate(data2) ? 0 : -1;
|
||||
|
||||
case GPSCallbackType::gotRTCMMessage:
|
||||
gotRTCMData((uint8_t*) data1, data2);
|
||||
break;
|
||||
|
||||
case GPSCallbackType::surveyInStatus:
|
||||
{
|
||||
SurveyInStatus* status = (SurveyInStatus*)data1;
|
||||
qCDebug(RTKGPSLog) << "Position: " << status->latitude << status->longitude << status->altitude;
|
||||
|
||||
qCDebug(RTKGPSLog) << QString("Survey-in status: %1s cur accuracy: %2mm valid: %3 active: %4").arg(status->duration).arg(status->mean_accuracy).arg((int)(status->flags & 1)).arg((int)((status->flags>>1) & 1));
|
||||
emit surveyInStatus(status->duration, status->mean_accuracy, status->latitude, status->longitude, status->altitude, (int)(status->flags & 1), (int)((status->flags>>1) & 1));
|
||||
}
|
||||
break;
|
||||
|
||||
case GPSCallbackType::setClock:
|
||||
/* do nothing */
|
||||
break;
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
@ -0,0 +1,93 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* (c) 2009-2020 QGROUNDCONTROL PROJECT <http://www.qgroundcontrol.org>
|
||||
*
|
||||
* QGroundControl is licensed according to the terms in the file
|
||||
* COPYING.md in the root of the source code directory.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <QString>
|
||||
#include <QThread>
|
||||
#include <QByteArray>
|
||||
#include <QSerialPort>
|
||||
|
||||
#include <atomic>
|
||||
|
||||
#include "GPSPositionMessage.h"
|
||||
#include "Drivers/src/gps_helper.h"
|
||||
|
||||
|
||||
/**
|
||||
** class GPSProvider
|
||||
* opens a GPS device and handles the protocol
|
||||
*/
|
||||
class GPSProvider : public QThread
|
||||
{
|
||||
Q_OBJECT
|
||||
public:
|
||||
|
||||
enum class GPSType {
|
||||
u_blox,
|
||||
trimble,
|
||||
septentrio
|
||||
};
|
||||
|
||||
GPSProvider(const QString& device,
|
||||
GPSType type,
|
||||
bool enableSatInfo,
|
||||
double surveyInAccMeters,
|
||||
int surveryInDurationSecs,
|
||||
bool useFixedBaseLocation,
|
||||
double fixedBaseLatitude,
|
||||
double fixedBaseLongitude,
|
||||
float fixedBaseAltitudeMeters,
|
||||
float fixedBaseAccuracyMeters,
|
||||
const std::atomic_bool& requestStop);
|
||||
~GPSProvider();
|
||||
|
||||
/**
|
||||
* this is called by the callback method
|
||||
*/
|
||||
void gotRTCMData(uint8_t *data, size_t len);
|
||||
|
||||
signals:
|
||||
void positionUpdate(GPSPositionMessage message);
|
||||
void satelliteInfoUpdate(GPSSatelliteMessage message);
|
||||
void RTCMDataUpdate(QByteArray message);
|
||||
void surveyInStatus(float duration, float accuracyMM, double latitude, double longitude, float altitude, bool valid, bool active);
|
||||
|
||||
protected:
|
||||
void run();
|
||||
|
||||
private:
|
||||
void publishGPSPosition();
|
||||
void publishGPSSatellite();
|
||||
|
||||
/**
|
||||
* callback from the driver for the platform specific stuff
|
||||
*/
|
||||
static int callbackEntry(GPSCallbackType type, void *data1, int data2, void *user);
|
||||
|
||||
int callback(GPSCallbackType type, void *data1, int data2);
|
||||
|
||||
QString _device;
|
||||
GPSType _type;
|
||||
const std::atomic_bool& _requestStop;
|
||||
double _surveyInAccMeters;
|
||||
int _surveryInDurationSecs;
|
||||
bool _useFixedBaseLoction;
|
||||
double _fixedBaseLatitude;
|
||||
double _fixedBaseLongitude;
|
||||
float _fixedBaseAltitudeMeters;
|
||||
float _fixedBaseAccuracyMeters;
|
||||
GPSHelper::GPSConfig _gpsConfig{};
|
||||
|
||||
struct sensor_gps_s _reportGpsPos;
|
||||
struct satellite_info_s *_pReportSatInfo = nullptr;
|
||||
|
||||
QSerialPort *_serial = nullptr;
|
||||
};
|
@ -0,0 +1,83 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* (c) 2009-2020 QGROUNDCONTROL PROJECT <http://www.qgroundcontrol.org>
|
||||
*
|
||||
* QGroundControl is licensed according to the terms in the file
|
||||
* COPYING.md in the root of the source code directory.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
|
||||
#include "RTCMMavlink.h"
|
||||
|
||||
#include "MultiVehicleManager.h"
|
||||
#include "Vehicle.h"
|
||||
|
||||
#include <cstdio>
|
||||
|
||||
RTCMMavlink::RTCMMavlink(QGCToolbox& toolbox)
|
||||
: _toolbox(toolbox)
|
||||
{
|
||||
_bandwidthTimer.start();
|
||||
}
|
||||
|
||||
void RTCMMavlink::RTCMDataUpdate(QByteArray message)
|
||||
{
|
||||
/* statistics */
|
||||
_bandwidthByteCounter += message.size();
|
||||
qint64 elapsed = _bandwidthTimer.elapsed();
|
||||
if (elapsed > 1000) {
|
||||
printf("RTCM bandwidth: %.2f kB/s\n", (float) _bandwidthByteCounter / elapsed * 1000.f / 1024.f);
|
||||
_bandwidthTimer.restart();
|
||||
_bandwidthByteCounter = 0;
|
||||
}
|
||||
|
||||
const qsizetype maxMessageLength = MAVLINK_MSG_GPS_RTCM_DATA_FIELD_DATA_LEN;
|
||||
mavlink_gps_rtcm_data_t mavlinkRtcmData;
|
||||
memset(&mavlinkRtcmData, 0, sizeof(mavlink_gps_rtcm_data_t));
|
||||
|
||||
if (message.size() < maxMessageLength) {
|
||||
mavlinkRtcmData.len = message.size();
|
||||
mavlinkRtcmData.flags = (_sequenceId & 0x1F) << 3;
|
||||
memcpy(&mavlinkRtcmData.data, message.data(), message.size());
|
||||
sendMessageToVehicle(mavlinkRtcmData);
|
||||
} else {
|
||||
// We need to fragment
|
||||
|
||||
uint8_t fragmentId = 0; // Fragment id indicates the fragment within a set
|
||||
int start = 0;
|
||||
while (start < message.size()) {
|
||||
int length = std::min(message.size() - start, maxMessageLength);
|
||||
mavlinkRtcmData.flags = 1; // LSB set indicates message is fragmented
|
||||
mavlinkRtcmData.flags |= fragmentId++ << 1; // Next 2 bits are fragment id
|
||||
mavlinkRtcmData.flags |= (_sequenceId & 0x1F) << 3; // Next 5 bits are sequence id
|
||||
mavlinkRtcmData.len = length;
|
||||
memcpy(&mavlinkRtcmData.data, message.data() + start, length);
|
||||
sendMessageToVehicle(mavlinkRtcmData);
|
||||
start += length;
|
||||
}
|
||||
}
|
||||
++_sequenceId;
|
||||
}
|
||||
|
||||
void RTCMMavlink::sendMessageToVehicle(const mavlink_gps_rtcm_data_t& msg)
|
||||
{
|
||||
QmlObjectListModel& vehicles = *_toolbox.multiVehicleManager()->vehicles();
|
||||
MAVLinkProtocol* mavlinkProtocol = _toolbox.mavlinkProtocol();
|
||||
for (int i = 0; i < vehicles.count(); i++) {
|
||||
Vehicle* vehicle = qobject_cast<Vehicle*>(vehicles[i]);
|
||||
WeakLinkInterfacePtr weakLink = vehicle->vehicleLinkManager()->primaryLink();
|
||||
|
||||
if (!weakLink.expired()) {
|
||||
mavlink_message_t message;
|
||||
SharedLinkInterfacePtr sharedLink = weakLink.lock();
|
||||
|
||||
mavlink_msg_gps_rtcm_data_encode_chan(mavlinkProtocol->getSystemId(),
|
||||
mavlinkProtocol->getComponentId(),
|
||||
sharedLink->mavlinkChannel(),
|
||||
&message,
|
||||
&msg);
|
||||
vehicle->sendMessageOnLinkThreadSafe(sharedLink.get(), message);
|
||||
}
|
||||
}
|
||||
}
|
@ -0,0 +1,40 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* (c) 2009-2020 QGROUNDCONTROL PROJECT <http://www.qgroundcontrol.org>
|
||||
*
|
||||
* QGroundControl is licensed according to the terms in the file
|
||||
* COPYING.md in the root of the source code directory.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <QObject>
|
||||
#include <QElapsedTimer>
|
||||
|
||||
#include "QGCToolbox.h"
|
||||
#include "MAVLinkProtocol.h"
|
||||
|
||||
/**
|
||||
** class RTCMMavlink
|
||||
* Receives RTCM updates and sends them via MAVLINK to the device
|
||||
*/
|
||||
class RTCMMavlink : public QObject
|
||||
{
|
||||
Q_OBJECT
|
||||
public:
|
||||
RTCMMavlink(QGCToolbox& toolbox);
|
||||
//TODO: API to select device(s)?
|
||||
|
||||
public slots:
|
||||
void RTCMDataUpdate(QByteArray message);
|
||||
|
||||
private:
|
||||
void sendMessageToVehicle(const mavlink_gps_rtcm_data_t& msg);
|
||||
|
||||
QGCToolbox& _toolbox;
|
||||
QElapsedTimer _bandwidthTimer;
|
||||
int _bandwidthByteCounter = 0;
|
||||
uint8_t _sequenceId = 0;
|
||||
};
|
@ -0,0 +1,94 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* Copyright (c) 2016 PX4 Development Team. All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions
|
||||
* are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright
|
||||
* notice, this list of conditions and the following disclaimer.
|
||||
* 2. Redistributions in binary form must reproduce the above copyright
|
||||
* notice, this list of conditions and the following disclaimer in
|
||||
* the documentation and/or other materials provided with the
|
||||
* distribution.
|
||||
* 3. Neither the name PX4 nor the names of its contributors may be
|
||||
* used to endorse or promote products derived from this software
|
||||
* without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
|
||||
* FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
|
||||
* COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
|
||||
* INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
|
||||
* OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
|
||||
* AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
|
||||
* LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
|
||||
* ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
* POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
/**
|
||||
* @file definitions.h
|
||||
* common platform-specific definitions & abstractions for gps
|
||||
* @author Beat Küng <beat-kueng@gmx.net>
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <QtGlobal>
|
||||
|
||||
#define GPS_READ_BUFFER_SIZE 1024
|
||||
|
||||
#define GPS_INFO(...) qInfo(__VA_ARGS__)
|
||||
#define GPS_WARN(...) qWarning(__VA_ARGS__)
|
||||
#define GPS_ERR(...) qCritical(__VA_ARGS__)
|
||||
|
||||
#include "vehicle_gps_position.h"
|
||||
#include "satellite_info.h"
|
||||
|
||||
#define M_DEG_TO_RAD (M_PI / 180.0)
|
||||
#define M_RAD_TO_DEG (180.0 / M_PI)
|
||||
#define M_DEG_TO_RAD_F 0.01745329251994f
|
||||
#define M_RAD_TO_DEG_F 57.2957795130823f
|
||||
|
||||
#include <QThread>
|
||||
|
||||
class Sleeper : public QThread
|
||||
{
|
||||
public:
|
||||
static void usleep(unsigned long usecs) { QThread::usleep(usecs); }
|
||||
};
|
||||
|
||||
static inline void gps_usleep(unsigned long usecs) {
|
||||
Sleeper::usleep(usecs);
|
||||
}
|
||||
|
||||
typedef uint64_t gps_abstime;
|
||||
|
||||
#include <QDateTime>
|
||||
/**
|
||||
* Get the current time in us. Function signature:
|
||||
* uint64_t hrt_absolute_time()
|
||||
*/
|
||||
static inline gps_abstime gps_absolute_time() {
|
||||
//FIXME: is there something with microsecond accuracy?
|
||||
return QDateTime::currentMSecsSinceEpoch() * 1000;
|
||||
}
|
||||
|
||||
//timespec is UNIX-specific
|
||||
#ifdef _WIN32
|
||||
#if _MSC_VER < 1900
|
||||
struct timespec
|
||||
{
|
||||
time_t tv_sec;
|
||||
long tv_nsec;
|
||||
};
|
||||
#else
|
||||
#include <time.h>
|
||||
#endif
|
||||
#endif
|
||||
|
@ -0,0 +1,55 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* Copyright (c) 2012-2014 PX4 Development Team. All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions
|
||||
* are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright
|
||||
* notice, this list of conditions and the following disclaimer.
|
||||
* 2. Redistributions in binary form must reproduce the above copyright
|
||||
* notice, this list of conditions and the following disclaimer in
|
||||
* the documentation and/or other materials provided with the
|
||||
* distribution.
|
||||
* 3. Neither the name PX4 nor the names of its contributors may be
|
||||
* used to endorse or promote products derived from this software
|
||||
* without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
|
||||
* FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
|
||||
* COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
|
||||
* INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
|
||||
* OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
|
||||
* AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
|
||||
* LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
|
||||
* ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
* POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
#pragma once
|
||||
#include <stdint.h>
|
||||
|
||||
/*
|
||||
* This file is auto-generated from https://github.com/PX4/Firmware/blob/master/msg/satellite_info.msg
|
||||
* and was manually copied here.
|
||||
*/
|
||||
|
||||
struct satellite_info_s {
|
||||
uint64_t timestamp;
|
||||
uint8_t count;
|
||||
uint8_t svid[20];
|
||||
uint8_t used[20];
|
||||
uint8_t elevation[20];
|
||||
uint8_t azimuth[20];
|
||||
uint8_t snr[20];
|
||||
uint8_t prn[20];
|
||||
#ifdef __cplusplus
|
||||
static const uint8_t SAT_INFO_MAX_SATELLITES = 20;
|
||||
|
||||
#endif
|
||||
};
|
@ -0,0 +1,72 @@
|
||||
/****************************************************************************
|
||||
*
|
||||
* Copyright (c) 2012-2014 PX4 Development Team. All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions
|
||||
* are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright
|
||||
* notice, this list of conditions and the following disclaimer.
|
||||
* 2. Redistributions in binary form must reproduce the above copyright
|
||||
* notice, this list of conditions and the following disclaimer in
|
||||
* the documentation and/or other materials provided with the
|
||||
* distribution.
|
||||
* 3. Neither the name PX4 nor the names of its contributors may be
|
||||
* used to endorse or promote products derived from this software
|
||||
* without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
|
||||
* FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
|
||||
* COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
|
||||
* INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
|
||||
* OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
|
||||
* AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
|
||||
* LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
|
||||
* ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
* POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
****************************************************************************/
|
||||
|
||||
/* Auto-generated by genmsg_cpp from file /home/beat/px4/src/Firmware/msg/vehicle_gps_position.msg */
|
||||
|
||||
|
||||
#pragma once
|
||||
#include <stdint.h>
|
||||
|
||||
/*
|
||||
* This file is auto-generated from https://github.com/PX4/Firmware/blob/master/msg/vehicle_gps_position.msg
|
||||
* and was manually copied here.
|
||||
*/
|
||||
|
||||
struct sensor_gps_s {
|
||||
uint64_t timestamp;
|
||||
uint64_t time_utc_usec;
|
||||
int32_t lat;
|
||||
int32_t lon;
|
||||
int32_t alt;
|
||||
int32_t alt_ellipsoid;
|
||||
uint16_t automatic_gain_control;
|
||||
uint8_t jamming_state;
|
||||
float s_variance_m_s;
|
||||
float c_variance_rad;
|
||||
float eph;
|
||||
float epv;
|
||||
float hdop;
|
||||
float vdop;
|
||||
int32_t noise_per_ms;
|
||||
int32_t jamming_indicator;
|
||||
float vel_m_s;
|
||||
float vel_n_m_s;
|
||||
float vel_e_m_s;
|
||||
float vel_d_m_s;
|
||||
float cog_rad;
|
||||
int32_t timestamp_time_relative;
|
||||
float heading;
|
||||
uint8_t fix_type;
|
||||
bool vel_ned_valid;
|
||||
uint8_t satellites_used;
|
||||
};
|
Loading…
Reference in new issue