Compare commits

..

7 Commits

28 changed files with 294 additions and 2090 deletions

Binary file not shown.

Before

Width:  |  Height:  |  Size: 434 KiB

124
README.md
View File

@ -1,15 +1,123 @@
# YOLOv7+DeepSORT行人检测与跟踪
# 简化后的YOLOv7推断模块
## 简化后的YOLOv7推断模块
## 程序余API介绍
***见yolov7-inferonly分支***
>主要的程序接口API全部位于detect.py和param.py两个文件中。在输入接口 `detect(self,imgs:torch.Tensor) -> list`的输入变量类型是torch的张量类型可以通过opencv的`cv2.imread()`或者`capture.read()`读取得到numpy数组(array)类型,然后通过`torch.Tensor()`或者`torchvision.transforms`进行类型转换。图片尺寸必须是方形即W=H,如果不是建议先通过resize和pad操作进行变换如果是单张图片输入模型前必须扩充维数可以使用Tensor的`unsqueeze(dim=0)`方法进行(即扩充第一维度),实现的伪代码如下:
## DeepSORT模块
```python
raw_img = cv2.imread(path_src) # 读取单张图片path_src处填写图片路径
raw_transform = transforms.Compose([transforms.ToPILImage(),
transforms.Resize((360,640)),
transforms.Pad((0,(640-360)//2)),
transforms.ToTensor()])
return raw_transform(raw_img).unsqueeze(dim=0)
```
* 效果图
>如果是读取视频只需按opencv读取视频的方法进行即可:
![avatar](DeepSORT1.png)
```python
capture = cv2.VideoCapture(path_src) # 读取摄像头为cv2.VideoCapture(index),index是相机索引通常为0即可读取视频文件时path_src填写路径即可
raw_transform = transforms.Compose([transforms.ToPILImage(),
transforms.Resize((360,640)),
transforms.Pad((0,(640-360)//2)),
transforms.ToTensor()]) # 预先组合好的变换函数
ret,frame = capture.read()
while ret is not None:
frame_tensor = raw_transform(frame).unsqueeze(dim=0) # 由于输入的是单张图片需要在dim=0进行维数扩充由(C,H,W)到(1,C,H,W),总尺寸大小其实不会发生改变
# your coder for detection
# ...
ret,frame = capture.read()
capture.release()
```
## Reference
### 检测模块detect.py
[DeepSORT实现-github](https://github.com/ZQPei/deep_sort_pytorch "DeepSORT实现")
主要是包含*一个类*和*一段示例代码*
* `class YOLOv7(object)`:该类即为本程序提供的主要接口
* `__init__(self,option:param.Parameters) -> None`
>主要是创建一个已加载权重的模型model以供后续推断过程的使用之所以将函数与推断模型分开的原因是权重加载较慢这样做可以使得不必在每一次推断的过程中重复加载权重这对提高实验程序的实时性有很大帮助。尽管使用ONNX模型更加有利于提升速度但很多时候完成课程或毕业设计这样已经足够了
>
>参数:`opt:param.Parameters`是位于param.py中的一个类仅仅是为了避免使用全局变量同时满足各种情形下的传参操作才这么设计。
>
>返回值:无,但内部变量`self.model`是已加载参数的模型其父类为torch.nn.Moudules)。
* `detect(self,imgs:torch.Tensor) -> list`
>参数:`imgs`(类型为torch.Tensor),大小应为(B,3,H,W)。其中B为照片的张数当为单张照片时为B==1即可。3是RGB三通道W、H分别为图片的高和宽**要求为64的倍数**。
>
>返回值:一个包含所有图片的bouding_box的列表列表中每个元素对应每一张图片其类型是二维的torch.Tensor第二维是对应单张图片中的bouding_box;最后一维长度为6分别对应(左上对角点横坐标`x1`,左上对角点纵坐标`y1`,右下角坐标横坐标`x2`,右下角纵坐标`y2`,可信度`conf`,类别`cls`)。*注意元素类型都为float型,后期处理可能涉及到类型转换。*
### 传参模块param.py
>主要是一个简单的类`Parameters`
| 变量名称 | 说明 |
| ------------ | ----------------------------------------------- |
| weights | 权重文件`.pt`文件的路径 |
| conf_thres | 置信度阈值,小于该置信度将被过滤 |
| iou_thres | 非极大值抑制(NMS)阈值,重叠比值超过该值将被合并 |
| device | 所部署到的设备`.to(device)` |
| classes | 选中的类,为`None`时代表全选 |
| agnostic_nms | NMS算法选项通常为`None`即可 |
| augment | 传入YOLOv7模型的参数通常为`None`即可 |
## 示例
### 代码
```python
def get_img_test(path_src:str):
raw_img = cv2.imread(path_src) # get image as numpy arrary
raw_transform = transforms.Compose([transforms.ToPILImage(), # process it as needed and transform it into torch.Tensor
transforms.Resize((360,640)),
transforms.Pad((0,(640-360)//2)),
transforms.ToTensor()])
return raw_transform(raw_img).unsqueeze(dim=0)
```
```python
def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
pred_img = pred_img.astype(np.uint) # transform its type into uint,pay attention to the confidence would be zero as percentage never large than 1.0
print(type(img))
for info in pred_img:
x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)# recover the pixel bias because of padding
print(type(img))
cv2.imshow("test",img)
cv2.waitKey(0)
cv2.imwrite("result.png",img) # save the image
```
```python
# example
if __name__ == '__main__':
my_option = param.Parameters()
model = YOLOv7(my_option) # load model in head
with torch.no_grad():
img = get_img_test("test.png") # get image as torch.Tensor with size of [1,1,640,640]
print(f"img size is {img.size()}")
result = model.detect(img) # get the sequence of result
print(f"result is :{result}")
result = result[0].detach().cpu().numpy() # the single img is index 0
img_src = cv2.imread("test.png") # re-read for imshow
img_src = cv2.resize(img_src,(640,360))
opencv_box_plot(img_src,result) # show the result in picture
```
### 输入图片
![avatar](test.png)
### 输出图片
![avatar](result.png)
### 参考链接
[YOLOv7官方实现-github](https://github.com/WongKinYiu/yolov7)
[YOLOv7论文-arxiv](https://arxiv.org/abs/2207.02696)

View File

View File

@ -1,15 +0,0 @@
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)))

View File

@ -1,91 +0,0 @@
import torch
from torch import nn
import torchvision.transforms as transforms
import numpy as np
import cv2
import logging
from .model import Net,MyNet,RestNet18
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()
class MyExtractor(object):
def __init__(self, model_path, use_cuda=True):
self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
self.net = RestNet18()
self.net.load_state_dict(torch.load(model_path,map_location=torch.device(self.device)))
if self.device=="cuda":
self.net = self.net.to(self.device)
logger = logging.getLogger("root.tracker")
logger.info("Loading weights from {}... Done!".format(model_path))
self.raw_transformer = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((128,64)),# hxw
transforms.ToTensor(),
])
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
"""
im_batch = torch.cat([self.raw_transformer(im).unsqueeze(0) for im in im_crops], dim=0)
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()

View File

@ -1,228 +0,0 @@
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
# Define model
# input:128x256x3
# output: 64
class MyNet(nn.Module):
def __init__(self):
super(MyNet, self).__init__()
self.conv1 = nn.Conv2d(3,8,5,1,2)
self.maxpool1 = nn.MaxPool2d(2) # 64x128x8
self.bn1 = nn.BatchNorm2d(8)
self.conv2 = nn.Conv2d(8,16,5,1,2)
self.maxpool2 = nn.MaxPool2d(2) # 32x64x16
self.bn2 = nn.BatchNorm2d(16)
self.conv3 = nn.Conv2d(16,32,5,1,2)
self.conv4 = nn.Conv2d(32,64,5,1,2)
self.maxpool3 = nn.MaxPool2d(2) # 16x32x64
self.bn3 = nn.BatchNorm2d(64)
self.conv5 = nn.Conv2d(64,32,5,1,2) # 16x32x32
self.maxpool4 = nn.MaxPool2d(2) # 8x16x32
self.bn4 = nn.BatchNorm2d(32)
self.conv6 = nn.Conv2d(32,64,8,8) #1x2x64
self.flat = nn.Flatten()
self.linear = nn.Linear(2*64,64) # 64D-feature ID
self.sigmoid = nn.Sigmoid()
self.lrelu = nn.LeakyReLU()
def forward(self, x:torch.Tensor):
x = self.conv1(x)
x = self.lrelu(x)
x = self.maxpool1(x)
x = self.bn1(x)
x = self.conv2(x)
x = self.lrelu(x)
x = self.maxpool2(x)
x = self.bn2(x)
x = self.conv3(x)
x = self.lrelu(x)
x = self.conv4(x)
x = self.lrelu(x)
x = self.maxpool3(x)
x = self.bn3(x)
x = self.conv5(x)
x = self.lrelu(x)
x = self.maxpool4(x)
x = self.bn4(x)
x = self.conv6(x)
x = self.sigmoid(x)
x = self.flat(x)
x = self.linear(x)
return x
if __name__ == '__main__':
net = Net()
x = torch.randn(4,3,128,64)
y = net(x)
#########ResNet18
class RestNetBasicBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride):
super(RestNetBasicBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
self.bn1 = nn.BatchNorm2d(out_channels)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1)
self.bn2 = nn.BatchNorm2d(out_channels)
def forward(self, x):
output = self.conv1(x)
output = F.relu(self.bn1(output))
output = self.conv2(output)
output = self.bn2(output)
return F.leaky_relu(x + output)
class RestNetDownBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride):
super(RestNetDownBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride[0], padding=1)
self.bn1 = nn.BatchNorm2d(out_channels)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride[1], padding=1)
self.bn2 = nn.BatchNorm2d(out_channels)
self.extra = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride[0], padding=0),
nn.BatchNorm2d(out_channels)
)
def forward(self, x):
extra_x = self.extra(x)
output = self.conv1(x)
out = F.relu(self.bn1(output))
out = self.conv2(out)
out = self.bn2(out)
return F.leaky_relu(extra_x + out)
class RestNet18(nn.Module):
def __init__(self):
super(RestNet18, self).__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=(3, 3), padding=1)
self.bn1 = nn.BatchNorm2d(32)
self.maxpool = nn.MaxPool2d(2)
self.layer1 = nn.Sequential(RestNetBasicBlock(32, 32, 1),
RestNetBasicBlock(32, 32, 1))
self.layer2 = nn.Sequential(RestNetDownBlock(32, 64, [2, 1]),
RestNetBasicBlock(64, 64, 1))
self.layer3 = nn.Sequential(RestNetDownBlock(64, 128, [2, 1]),
RestNetBasicBlock(128, 128, 1))
self.layer4 = nn.Sequential(RestNetDownBlock(128, 256, [2, 1]),
RestNetBasicBlock(256, 256, 1))
self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
self.fc = nn.Linear(32768, 23)
def forward(self, x):
x = self.conv1(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
# x = self.avgpool(x)
x = x.reshape(x.shape[0], -1)
# print(x.size())
x = self.fc(x)
return x

View File

@ -1,77 +0,0 @@
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")

View File

@ -1,189 +0,0 @@
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()

View File

@ -1,130 +0,0 @@
import numpy as np
import torch
from .deep.feature_extractor import Extractor,MyExtractor
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, model_config=None, max_dist=1e-2, min_confidence=0.35, nms_max_overlap=0.9, max_iou_distance=0.40, max_age=10, n_init=3, nn_budget=5000, use_cuda=True):
self.min_confidence = min_confidence
self.nms_max_overlap = nms_max_overlap
if model_config is None:
self.extractor = Extractor(model_path, use_cuda=use_cuda)
else:
self.extractor = MyExtractor(model_path, use_cuda=use_cuda)
max_cosine_distance = max_dist
metric = NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
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):
self.height, self.width = ori_img.shape[:2]
# generate detections
features = self._get_features(bbox_xywh, ori_img)
bbox_tlwh = self._xywh_to_tlwh(bbox_xywh)
detections = [Detection(bbox_tlwh[i], conf, features[i]) for i,conf in enumerate(confidences) if conf>self.min_confidence]
# run on non-maximum supression
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)
detections = [detections[i] for i in indices]
# update tracker
self.tracker.predict()
self.tracker.update(detections)
# output bbox identities
outputs = []
output_features = []
for track in self.tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
box = track.to_tlwh()
x1,y1,x2,y2 = self._tlwh_to_xyxy(box)
track_id = track.track_id
output_features.append(track.features)
outputs.append(np.array([x1,y1,x2,y2,track_id], dtype=np.int))
if len(outputs) > 0:
outputs = np.stack(outputs,axis=0)
return outputs,output_features
"""
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()
if len(bbox_xywh):
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 _xyxy_to_xywh(self, bbox_xyxy):
x1,y1,x2,y2 = bbox_xyxy
t = int((x1+x2)/2)
l = int((y1+y2)/2)
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

View File

@ -1,49 +0,0 @@
# 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)
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

View File

@ -1,90 +0,0 @@
# 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.
"""
length = len(candidates)
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)
# should be consious
gious = []
for i in range(length):
gious.append(float((area_bbox + area_candidates[i] - area_intersection[i])/(max(bbox[0]+bbox[2],candidates[i][0]+candidates[i][2])-min(bbox[0],candidates[i][0]))/(max(bbox[1]+bbox[3],candidates[i][1]+candidates[i][3])-min(bbox[1],candidates[i][1]))) )
# return area_intersection / (area_bbox + area_candidates - area_intersection)
return np.array(gious)
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

View File

@ -1,229 +0,0 @@
# vim: expandtab:ts=4:sw=4
import numpy as np
import scipy.linalg
"""
Table for the 0.95 quantile of the chi-square distribution with N degrees of
freedom (contains values for N=1, ..., 9). Taken from MATLAB/Octave's chi2inv
function and used as Mahalanobis gating threshold.
"""
chi2inv95 = {
1: 3.8415,
2: 5.9915,
3: 7.8147,
4: 9.4877,
5: 11.070,
6: 12.592,
7: 14.067,
8: 15.507,
9: 16.919}
class KalmanFilter(object):
"""
A simple Kalman filter for tracking bounding boxes in image space.
The 8-dimensional state space
x, y, a, h, vx, vy, va, vh
contains the bounding box center position (x, y), aspect ratio a, height h,
and their respective velocities.
Object motion follows a constant velocity model. The bounding box location
(x, y, a, h) is taken as direct observation of the state space (linear
observation model).
"""
def __init__(self):
ndim, dt = 4, 1.
# Create Kalman filter model matrices.
self._motion_mat = np.eye(2 * ndim, 2 * ndim)
for i in range(ndim):
self._motion_mat[i, ndim + i] = dt
self._update_mat = np.eye(ndim, 2 * ndim)
# Motion and observation uncertainty are chosen relative to the current
# state estimate. These weights control the amount of uncertainty in
# the model. This is a bit hacky.
self._std_weight_position = 1. / 20
self._std_weight_velocity = 1. / 160
def initiate(self, measurement):
"""Create track from unassociated measurement.
Parameters
----------
measurement : ndarray
Bounding box coordinates (x, y, a, h) with center position (x, y),
aspect ratio a, and height h.
Returns
-------
(ndarray, ndarray)
Returns the mean vector (8 dimensional) and covariance matrix (8x8
dimensional) of the new track. Unobserved velocities are initialized
to 0 mean.
"""
mean_pos = measurement
mean_vel = np.zeros_like(mean_pos)
mean = np.r_[mean_pos, mean_vel]
std = [
2 * self._std_weight_position * measurement[3],
2 * self._std_weight_position * measurement[3],
1e-2,
2 * self._std_weight_position * measurement[3],
10 * self._std_weight_velocity * measurement[3],
10 * self._std_weight_velocity * measurement[3],
1e-5,
10 * self._std_weight_velocity * measurement[3]]
covariance = np.diag(np.square(std))
return mean, covariance
def predict(self, mean, covariance):
"""Run Kalman filter prediction step.
Parameters
----------
mean : ndarray
The 8 dimensional mean vector of the object state at the previous
time step.
covariance : ndarray
The 8x8 dimensional covariance matrix of the object state at the
previous time step.
Returns
-------
(ndarray, ndarray)
Returns the mean vector and covariance matrix of the predicted
state. Unobserved velocities are initialized to 0 mean.
"""
std_pos = [
self._std_weight_position * mean[3],
self._std_weight_position * mean[3],
1e-2,
self._std_weight_position * mean[3]]
std_vel = [
self._std_weight_velocity * mean[3],
self._std_weight_velocity * mean[3],
1e-5,
self._std_weight_velocity * mean[3]]
motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))
mean = np.dot(self._motion_mat, mean)
covariance = np.linalg.multi_dot((
self._motion_mat, covariance, self._motion_mat.T)) + motion_cov
return mean, covariance
def project(self, mean, covariance):
"""Project state distribution to measurement space.
Parameters
----------
mean : ndarray
The state's mean vector (8 dimensional array).
covariance : ndarray
The state's covariance matrix (8x8 dimensional).
Returns
-------
(ndarray, ndarray)
Returns the projected mean and covariance matrix of the given state
estimate.
"""
std = [
self._std_weight_position * mean[3],
self._std_weight_position * mean[3],
1e-1,
self._std_weight_position * mean[3]]
innovation_cov = np.diag(np.square(std))
mean = np.dot(self._update_mat, mean)
covariance = np.linalg.multi_dot((
self._update_mat, covariance, self._update_mat.T))
return mean, covariance + innovation_cov
def update(self, mean, covariance, measurement):
"""Run Kalman filter correction step.
Parameters
----------
mean : ndarray
The predicted state's mean vector (8 dimensional).
covariance : ndarray
The state's covariance matrix (8x8 dimensional).
measurement : ndarray
The 4 dimensional measurement vector (x, y, a, h), where (x, y)
is the center position, a the aspect ratio, and h the height of the
bounding box.
Returns
-------
(ndarray, ndarray)
Returns the measurement-corrected state distribution.
"""
projected_mean, projected_cov = self.project(mean, covariance)
chol_factor, lower = scipy.linalg.cho_factor(
projected_cov, lower=True, check_finite=False)
kalman_gain = scipy.linalg.cho_solve(
(chol_factor, lower), np.dot(covariance, self._update_mat.T).T,
check_finite=False).T
innovation = measurement - projected_mean
new_mean = mean + np.dot(innovation, kalman_gain.T)
new_covariance = covariance - np.linalg.multi_dot((
kalman_gain, projected_cov, kalman_gain.T))
return new_mean, new_covariance
def gating_distance(self, mean, covariance, measurements,
only_position=False):
"""Compute gating distance between state distribution and measurements.
A suitable distance threshold can be obtained from `chi2inv95`. If
`only_position` is False, the chi-square distribution has 4 degrees of
freedom, otherwise 2.
Parameters
----------
mean : ndarray
Mean vector over the state distribution (8 dimensional).
covariance : ndarray
Covariance of the state distribution (8x8 dimensional).
measurements : ndarray
An Nx4 dimensional matrix of N measurements, each in
format (x, y, a, h) where (x, y) is the bounding box center
position, a the aspect ratio, and h the height.
only_position : Optional[bool]
If True, distance computation is done with respect to the bounding
box center position only.
Returns
-------
ndarray
Returns an array of length N, where the i-th element contains the
squared Mahalanobis distance between (mean, covariance) and
`measurements[i]`.
"""
mean, covariance = self.project(mean, covariance)
if only_position:
mean, covariance = mean[:2], covariance[:2, :2]
measurements = measurements[:, :2]
cholesky_factor = np.linalg.cholesky(covariance)
d = measurements - mean
z = scipy.linalg.solve_triangular(
cholesky_factor, d.T, lower=True, check_finite=False,
overwrite_b=True)
squared_maha = np.sum(z * z, axis=0)
return squared_maha

View File

@ -1,192 +0,0 @@
# 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

View File

@ -1,177 +0,0 @@
# 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

View File

@ -1,73 +0,0 @@
# 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

View File

@ -1,166 +0,0 @@
# 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()
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

View File

@ -1,138 +0,0 @@
# vim: expandtab:ts=4:sw=4
from __future__ import absolute_import
import numpy as np
from . import kalman_filter
from . import linear_assignment
from . import iou_matching
from .track import Track
class Tracker:
"""
This is the multi-target tracker.
Parameters
----------
metric : nn_matching.NearestNeighborDistanceMetric
A distance metric for measurement-to-track association.
max_age : int
Maximum number of missed misses before a track is deleted.
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.
Attributes
----------
metric : nn_matching.NearestNeighborDistanceMetric
The distance metric used for measurement to track association.
max_age : int
Maximum number of missed misses before a track is deleted.
n_init : int
Number of frames that a track remains in initialization phase.
kf : kalman_filter.KalmanFilter
A Kalman filter to filter target trajectories in image space.
tracks : List[Track]
The list of active tracks at the current time step.
"""
def __init__(self, metric, max_iou_distance=0.7, max_age=70, n_init=3):
self.metric = metric
self.max_iou_distance = max_iou_distance
self.max_age = max_age
self.n_init = n_init
self.kf = kalman_filter.KalmanFilter()
self.tracks = []
self._next_id = 1
def predict(self):
"""Propagate track state distributions one time step forward.
This function should be called once every time step, before `update`.
"""
for track in self.tracks:
track.predict(self.kf)
def update(self, detections):
"""Perform measurement update and track management.
Parameters
----------
detections : List[deep_sort.detection.Detection]
A list of detections at the current time step.
"""
# Run matching cascade.
matches, unmatched_tracks, unmatched_detections = \
self._match(detections)
# Update track set.
for track_idx, detection_idx in matches:
self.tracks[track_idx].update(
self.kf, detections[detection_idx])
for track_idx in unmatched_tracks:
self.tracks[track_idx].mark_missed()
for detection_idx in unmatched_detections:
self._initiate_track(detections[detection_idx])
self.tracks = [t for t in self.tracks if not t.is_deleted()]
# Update distance metric.
active_targets = [t.track_id for t in self.tracks if t.is_confirmed()]
features, targets = [], []
for track in self.tracks:
if not track.is_confirmed():
continue
features += track.features
targets += [track.track_id for _ in track.features]
# track.features = []
self.metric.partial_fit(
np.asarray(features), np.asarray(targets), active_targets)
def _match(self, detections):
def gated_metric(tracks, dets, track_indices, detection_indices):
features = np.array([dets[i].feature for i in detection_indices])
targets = np.array([tracks[i].track_id for i in track_indices])
cost_matrix = self.metric.distance(features, targets)
cost_matrix = linear_assignment.gate_cost_matrix(
self.kf, cost_matrix, tracks, dets, track_indices,
detection_indices)
return cost_matrix
# Split track set into confirmed and unconfirmed tracks.
confirmed_tracks = [
i for i, t in enumerate(self.tracks) if t.is_confirmed()]
unconfirmed_tracks = [
i for i, t in enumerate(self.tracks) if not t.is_confirmed()]
# Associate confirmed tracks using appearance features.
matches_a, unmatched_tracks_a, unmatched_detections = \
linear_assignment.matching_cascade(
gated_metric, self.metric.matching_threshold, self.max_age,
self.tracks, detections, confirmed_tracks)
# Associate remaining tracks together with unconfirmed tracks using IOU.
iou_track_candidates = unconfirmed_tracks + [
k for k in unmatched_tracks_a if
self.tracks[k].time_since_update == 1]
unmatched_tracks_a = [
k for k in unmatched_tracks_a if
self.tracks[k].time_since_update != 1]
matches_b, unmatched_tracks_b, unmatched_detections = \
linear_assignment.min_cost_matching(
iou_matching.iou_cost, self.max_iou_distance, self.tracks,
detections, iou_track_candidates, unmatched_detections)
matches = matches_a + matches_b
unmatched_tracks = list(set(unmatched_tracks_a + unmatched_tracks_b))
return matches, unmatched_tracks, unmatched_detections
def _initiate_track(self, detection):
mean, covariance = self.kf.initiate(detection.to_xyah())
self.tracks.append(Track(
mean, covariance, self._next_id, self.n_init, self.max_age,
detection.feature))
self._next_id += 1

View File

@ -1,7 +1,10 @@
import cv2
import torch
import numpy as np
from torchvision import transforms
from models.experimental import attempt_load
from utils.general import non_max_suppression
from utils.general import check_img_size,non_max_suppression
import param
@ -10,7 +13,7 @@ class YOLOv7(object):
def __init__(self,option:param.Parameters) -> None:
# load FP32 model
self.opt = option
self.model = attempt_load(weights=self.opt.weights_yolo, device=self.opt.device)
self.model = attempt_load(weights=self.opt.weights, device=self.opt.device)
def detect(self,imgs:torch.Tensor) -> list:
# img.size should be (Batch,3,H,W),every value should be range of [0,1]
@ -25,6 +28,38 @@ class YOLOv7(object):
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
return pred
def get_img_test(path_src:str):
raw_img = cv2.imread(path_src)
raw_transform = transforms.Compose([transforms.ToPILImage(),
transforms.Resize((360,640)),
transforms.Pad((0,(640-360)//2)),
transforms.ToTensor()])
return raw_transform(raw_img).unsqueeze(dim=0)
def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
pred_img = pred_img.astype(np.uint)
print(type(img))
for info in pred_img:
x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)
print(type(img))
cv2.imshow("test",img)
cv2.waitKey(0)
cv2.imwrite("result.png",img)
# example
if __name__ == '__main__':
my_option = param.Parameters()
model = YOLOv7(my_option) # load model in head
with torch.no_grad():
img = get_img_test("test.png") # get image as torch.Tensor with size of [1,1,640,640]
print(f"img size is {img.size()}")
result = model.detect(img) # get the sequence of result
print(f"result is :{result}")
result = result[0].detach().cpu().numpy() # the single img is index 0
img_src = cv2.imread("test.png") # re-read for imshow
img_src = cv2.resize(img_src,(640,360))
opencv_box_plot(img_src,result) # show the result in picture

136
detector.py Normal file
View File

@ -0,0 +1,136 @@
import os
import cv2
import torch
import numpy as np
from torchvision import transforms
from models.experimental import attempt_load
from utils.general import check_img_size,non_max_suppression
import param
def IoU(box1, box2):
"""
:param box1: list in format [xmin1, ymin1, xmax1, ymax1]
:param box2: list in format [xmin2, ymin2, xamx2, ymax2]
:return: returns IoU ratio (intersection over union) of two boxes
"""
xmin1, ymin1, xmax1, ymax1 = box1
xmin2, ymin2, xmax2, ymax2 = box2[0],box2[1],box2[2],box2[3]
ymin2=ymin2-420
ymax2=ymax2-420
x_overlap = max(0, min(xmax1, xmax2) - max(xmin1, xmin2))
y_overlap = max(0, min(ymax1, ymax2) - max(ymin1, ymin2))
intersection = x_overlap * y_overlap
union = (xmax1 - xmin1) * (ymax1 - ymin1) + (xmax2 - xmin2) * (ymax2 - ymin2) - intersection
IOU=float(intersection) / union
return IOU
def cxywh2xyxy(cxywh):
xyxy = [0, 0, 0, 0]
xyxy[0] = int(cxywh[0] - cxywh[2] / 2)
xyxy[1] = int(cxywh[1] - cxywh[3] / 2)
xyxy[2] = int(cxywh[0] + cxywh[2] / 2)
xyxy[3] = int(cxywh[1] + cxywh[3] / 2)
xyxy = [(num if num > 0 else 0) for num in xyxy]
return xyxy
class YOLOv7(object):
def __init__(self,option:param.Parameters) -> None:
# load FP32 model
self.opt = option
self.model = attempt_load(weights=self.opt.weights, device=self.opt.device)
def detect(self,imgs:torch.Tensor) -> list:
# img.size should be (Batch,3,H,W),every value should be range of [0,1]
# Both H and W % 64 == 0
# Initialize
model = self.model
opt = self.opt
# Inference
with torch.no_grad(): # Calculating gradients would cause a GPU memory leak
pred = model(imgs, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
return pred
def get_img_test(path_src:str):
raw_img = cv2.imread(path_src)
raw_transform = transforms.Compose([transforms.ToPILImage(),
transforms.Resize((1080,1920)),
transforms.Pad((0, (1920-1080)//2)),
transforms.ToTensor()])
return raw_transform(raw_img).unsqueeze(dim=0)
def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
pred_img = pred_img.astype(np.uint)
print(type(img))
for info in pred_img:
x1, y1, x2, y2 = info[0], info[1], info[2], info[3]
if abs(x2-x1)> 32:
cut = img[(y1-420):(y2 - 420), x1:x2]
# cv2.rectangle(img, (x1, y1-420), (x2, y2-420), (255, 0, 0), 2)
cut = cv2.resize(cut, (128, 256))
results_len = len(results)
number_bbox = 0
for i in range(results_len):
id_bbox = results[i]
id = str(id_bbox).split(',')[0][1:]
xxyy = cxywh2xyxy(id_bbox[1:5])
if IoU(xxyy, info) > 0.7:
cv2.imwrite(f"./detector/cam{camera}/" + "d" + "_" + name.split(".")[0] + "_" + str(id) + ".jpg", cut)
break
else:
number_bbox += 1
if number_bbox == results_len:
cv2.imwrite(f"./detector/cam{camera}/" + "d" + "_" + name.split(".")[0] + "_" + "0" + ".jpg", cut)
# print(type(img))
# cv2.imshow("test", img)
cv2.waitKey(0)
# cv2.imwrite("result.png", img)
# example
if __name__ == '__main__':
my_option = param.Parameters()
model = YOLOv7(my_option) # load model in head
count_frame = -1
ret = None
frame = None
with torch.no_grad():
# for i in range(1,7):
camera = 4
with open(f"../UESTC_ReID_Dataset_V3/label/cam{camera}.txt", 'r') as flabel:
camera_path = f"../UESTC_ReID_Dataset_V3/images/img/cam{camera}/"
txt_data = flabel.readlines()
filename = os.listdir(camera_path)
results=[]
# filename = "011.jpg"
count = 0
for name in filename:
count += 1
if count % 12 == 0:
for line in txt_data:
data2 = line.split(',')
data2 = [int(data2) for data2 in data2[0:6]]
# print(data2[0])
if int(data2[0]) == int(name.split(".")[0]):
results.append(data2[1:6])
print(name)
img = get_img_test(f"../UESTC_ReID_Dataset_V3/images/img/cam{camera}/"+name) # get image as torch.Tensor with size of [1,1,640,640]
print(f"img size is {img.size()}")
result = model.detect(img) # get the sequence of result
print(f"result is :{result}")
result = result[0].detach().cpu().numpy() # the single img is index 0
img_src = cv2.imread(f"../UESTC_ReID_Dataset_V3/images/img/cam{camera}/"+name) # re-read for imshow
img_src = cv2.resize(img_src, (1920, 1080))
opencv_box_plot(img_src, result) # show the result in picture
results.clear()

View File

@ -1,185 +0,0 @@
import torch
from torchvision import transforms
import numpy as np
import cv2
import os
import param
from detect import YOLOv7
from deep_sort import deep_sort as dsort
def get_img_test(raw_img:np.array):
raw_transform = transforms.Compose([transforms.ToPILImage(),
transforms.Resize((360,640)),
transforms.Pad((0,(640-360)//2)),])
return raw_transform(raw_img)
def reverse_box_get_from_yolo(detection):
x1 = int(detection[0]*1920/640)
x2 = int(detection[2]*1920/640)
y1 = int((detection[1]-140)*1080/360)
y2 = int((detection[3]-140)*1080/360)
return x1,y1,x2,y2
def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
pred_img = pred_img.astype(np.uint)
for info in pred_img:
x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)
cv2.imshow("test",img)
cv2.waitKey(0)
cv2.imwrite("result.png",img)
def opencv_match_pointer_plot(img:cv2.Mat,start_base:tuple,text:str,color:tuple):
SIZE = 50
brush = [0,0]
brush[0]=start_base[0]
brush[1]=start_base[1]
brush_push = [brush[0]+SIZE,brush[1]-SIZE]
cv2.line(img,brush,brush_push,color,1)
# brush=brush_push
# brush_push[0]+=SIZE*2
# brush_push[1]+=0
# cv2.line(img,brush,brush_push,color,1)
cv2.putText(img,text,brush_push, cv2.FONT_HERSHEY_PLAIN, 1, color, 2)
return img
routine = {}
def opencv_sort_plot(img:cv2.Mat,pred_yolo:np.array,pred_sort:np.array):
for info in pred_yolo:
x1,y1,x2,y2 = int(info[0]),int(info[1]),int(info[2]),int(info[3])
cv2.rectangle(img,(x1,y1),(x2,y2),(0,255,0),1)
for info in pred_sort:
x1,y1,x2,y2,id = info[0],info[1],info[2],info[3],info[4]
cxy = (int((x1+x2)/2),int((y1+y2)/2))
color = (int(id%3*100),int(id%4*75),int(id%5*50))
cv2.putText(img,str(id),cxy, cv2.FONT_HERSHEY_PLAIN, 1.0, color, 2)
if routine.get(id) is None:
routine[id] = [cxy]
else:
for i in range(1,len(routine[id])):
cv2.line(img,routine[id][i-1],routine[id][i],color,1)
cv2.line(img,routine[id][-1],cxy,color,1)
routine[id].append(cxy)
return img
def extractidfeature(id_path:str,extractor):
img_id = cv2.imread(id_path)
tensor_id = extractor([img_id])
return tensor_id
# example
if __name__ == '__main__':
MAX_BUFFLEN = 10
DISTANCE_THRESHOLD = 0.4
SAVE_TXT_FLAG = True
my_option = param.Parameters()
print(torch.__version__)
print(my_option.device)
query_path = f"./mydataset/query/cam{my_option.query_index}/" # query from .jpg photos
gallary_path = f"mydataset/video/cam{my_option.gallary_index}.mp4" # source from mp4 via yolo
output_video_path = f"output/output{my_option.gallary_index}_{my_option.query_index}.mp4" # output file path
output_txt_path = f"output/rank10_detection_{my_option.gallary_index}_{my_option.query_index}.txt"
query_features = []
paths = os.listdir(query_path)
query_match_buff = [[] for p in paths]
query_matched = [-1 for p in paths]
# load detection model and deepsort model
model = YOLOv7(my_option)
deepsort = dsort.DeepSort(model_path=my_option.weights_reid,model_config="self",use_cuda=(torch.device("cuda:0") == my_option.device))
# get query feature id
for p in paths:
query_features.append(extractidfeature(query_path+p,deepsort.extractor))
with torch.no_grad():
sources = cv2.VideoCapture(gallary_path)
target = cv2.VideoWriter(output_video_path,cv2.VideoWriter_fourcc('m', 'p', '4', 'v'),24,(1920,1080))
frame_counter = 0
while True:
ret,frame = sources.read()
if ret is False:
break
img = get_img_test(frame) # get image as torch.Tensor with size of [1,1,640,640]
img_tensor = transforms.ToTensor()(img).unsqueeze(dim=0).to(my_option.device)
detections = model.detect(img_tensor) # get the sequence of result
detections = detections[0].detach().cpu().numpy() # the single img is index 0
for detection in detections:
detection[0],detection[1],detection[2],detection[3] = reverse_box_get_from_yolo(detection[0:4])
bbox_xywhs = []
confs = []
for xyxycc in detections:
xywh = deepsort._xyxy_to_xywh(xyxycc[0:4])
conf = xyxycc[4]
clas = int(xyxycc[5])
bbox_xywhs.append(xywh[:])
confs.append(conf)
dpsort, features = deepsort.update(bbox_xywh=np.array(bbox_xywhs),confidences=np.array(confs),ori_img=np.array(frame))
img_drawing = np.array(frame)
img_drawing = opencv_sort_plot(img_drawing,detections,dpsort)
# ReID: compute features
for i in range(len(dpsort)):
track = dpsort[i]
MOT_id = track[-1]
feature_distances = []
for j in range(len(query_features)):
feature_distance = np.mean(np.power(features[i]-query_features[j],2))
feature_distances.append(feature_distance)
# ReID: rank distances for every feature
feature_distances_sorted = sorted(feature_distances)
for distance in feature_distances_sorted:
if distance < DISTANCE_THRESHOLD:
index_d = feature_distances.index(distance)
query_match_buff[index_d].append((MOT_id,distance,(frame_counter,track[0],track[1],track[2],track[3])))
query_match_buff[index_d].sort(key=lambda x:x[1],reverse=True)
if len(query_match_buff[index_d]) > MAX_BUFFLEN:
query_match_buff[index_d].pop()
count = 0
for term in query_match_buff[index_d]:
if term[0] == MOT_id:
count+=1
if count > MAX_BUFFLEN//3:
cx = int((track[0]+track[2])/2)
cy = int((track[1]+track[3])/2)
color = (int(index_d%2*100),int(index_d%3*75),int(index_d%4*50))
# img_drawing=opencv_match_pointer_plot(img_drawing,(cx,track[1]-10),f"{paths[index_d]}",color)
# cv2.rectangle(img_drawing,(track[i][0],track[i][1]),(track[i][2],track[i][3]),color,1)
cv2.putText(img_drawing,f"q_{paths[index_d]}".split(".")[0],(track[0],track[1]), cv2.FONT_HERSHEY_PLAIN, 1, color, 2)
break
else:
break
# Final plot and show video
cv2.imshow("test",img_drawing)
cv2.waitKey(1)
target.write(img_drawing)
frame_counter += 1
# save the video and end the program
target.release()
# rank10 output as txt
if SAVE_TXT_FLAG is True:
with open(output_txt_path,mode="w+") as f:
result = []
for i in range(len(query_match_buff)):
for raw_info in query_match_buff[i]:
detect_info = raw_info[-1]
result.append((detect_info[0],i,detect_info[1],detect_info[2],detect_info[3],detect_info[4]))
result.sort(key=lambda x:x[0])
for info in result:
f.write(str(info[0])+","+str(info[1])+","+str(info[2])+","+str(info[3])+","+str(info[4])+","+str(info[5])+"\n")

View File

@ -2,16 +2,12 @@ import torch
class Parameters(object):
def __init__(self) -> None:
self.weights_yolo = "./best.pt"
self.weights_reid = "./resnet18.pth"
self.conf_thres = 0.40
self.iou_thres = 0.60
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.nosave = False
self.classes = [0]
self.weights = "./best.pt"
self.conf_thres = 0.25
self.iou_thres = 0.45
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.classes = [0]
self.agnostic_nms = None
self.augment = None
self.query_index = 1
self.gallary_index= 2
pass

View File

@ -1,40 +0,0 @@
# Usage: pip install -r requirements.txt
# Base ----------------------------------------
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.1
Pillow>=7.1.2
PyYAML>=5.3.1
requests>=2.23.0
scipy>=1.4.1
torch>=1.7.0,!=1.12.0
torchvision>=0.8.1,!=0.13.0
tqdm>=4.41.0
protobuf<4.21.3
# Logging -------------------------------------
tensorboard>=2.4.1
# wandb
# Plotting ------------------------------------
pandas>=1.1.4
seaborn>=0.11.0
# Export --------------------------------------
# coremltools>=4.1 # CoreML export
# onnx>=1.9.0 # ONNX export
# onnx-simplifier>=0.3.6 # ONNX simplifier
# scikit-learn==0.19.2 # CoreML quantization
# tensorflow>=2.4.1 # TFLite export
# tensorflowjs>=3.9.0 # TF.js export
# openvino-dev # OpenVINO export
# Extras --------------------------------------
ipython # interactive notebook
psutil # system utilization
thop # FLOPs computation
# albumentations>=1.0.3
# pycocotools>=2.0 # COCO mAP
# roboflow
yacs

BIN
result.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 304 KiB

BIN
test.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.7 MiB

View File

@ -129,7 +129,6 @@ def profile(x, ops, n=100, device=None):
s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters
raise Exception("thop package is discarded")
print(f'{p:12}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}')
@ -214,7 +213,6 @@ def model_info(model, verbose=False, img_size=640):
(i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
try: # FLOPS
raise Exception("package thop is discarded")
from thop import profile
stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input