forked from wffjwbbf/ComDesignProject
Compare commits
7 Commits
_test_mast
...
yolov7-inf
| Author | SHA1 | Date |
|---|---|---|
|
|
7ff0d6be69 | |
|
|
c880537c4f | |
|
|
8e74f642a3 | |
|
|
6663090d5a | |
|
|
d6e0318f91 | |
|
|
ab64f31b08 | |
|
|
76697fa840 |
BIN
DeepSORT1.png
BIN
DeepSORT1.png
Binary file not shown.
|
Before Width: | Height: | Size: 434 KiB |
124
README.md
124
README.md
|
|
@ -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读取视频的方法进行即可:
|
||||
|
||||

|
||||
```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
|
||||
```
|
||||
|
||||
### 输入图片
|
||||
|
||||

|
||||
|
||||
### 输出图片
|
||||
|
||||

|
||||
|
||||
### 参考链接
|
||||
|
||||
[YOLOv7官方实现-github](https://github.com/WongKinYiu/yolov7)
|
||||
|
||||
[YOLOv7论文-arxiv](https://arxiv.org/abs/2207.02696)
|
||||
|
|
|
|||
|
|
@ -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)))
|
||||
|
||||
|
||||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -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")
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
39
detect.py
39
detect.py
|
|
@ -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
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
||||
|
||||
185
example_reid.py
185
example_reid.py
|
|
@ -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")
|
||||
14
param.py
14
param.py
|
|
@ -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
|
||||
|
||||
|
|
@ -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
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 304 KiB |
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue