218 lines
5.3 KiB
Python
218 lines
5.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Created on Fri May 31 14:50:21 2024
|
|
|
|
@author: ym
|
|
"""
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from scipy.spatial.distance import cdist
|
|
from tracking.trackers.reid.config import config as ReIDConfig
|
|
from tracking.trackers.reid.reid_interface import ReIDInterface
|
|
ReIDEncoder = ReIDInterface(ReIDConfig)
|
|
|
|
def read_data_file(datapath):
|
|
|
|
with open(datapath, 'r') as file:
|
|
lines = file.readlines()
|
|
Videos = []
|
|
FrameBoxes, FrameFeats = [], []
|
|
boxes, feats = [], []
|
|
|
|
bboxes, ffeats = [], []
|
|
timestamp = []
|
|
t1 = None
|
|
for line in lines:
|
|
if line.find('CameraId') >= 0:
|
|
t = int(line.split(',')[1].split(':')[1])
|
|
timestamp.append(t)
|
|
|
|
if len(boxes) and len(feats):
|
|
FrameBoxes.append(np.array(boxes, dtype = np.float32))
|
|
FrameFeats.append(np.array(feats, dtype = np.float32))
|
|
|
|
boxes, feats = [], []
|
|
|
|
if t1 and t - t1 > 1e4:
|
|
Videos.append((FrameBoxes, FrameFeats))
|
|
FrameBoxes, FrameFeats = [], []
|
|
t1 = int(line.split(',')[1].split(':')[1])
|
|
|
|
if line.find('box') >= 0:
|
|
box = line.split(':', )[1].split(',')[:-1]
|
|
boxes.append(box)
|
|
bboxes.append(boxes)
|
|
|
|
|
|
if line.find('feat') >= 0:
|
|
feat = line.split(':', )[1].split(',')[:-1]
|
|
feats.append(feat)
|
|
ffeats.append(feat)
|
|
|
|
|
|
|
|
|
|
FrameBoxes.append(np.array(boxes, dtype = np.float32))
|
|
FrameFeats.append(np.array(feats, dtype = np.float32))
|
|
Videos.append((FrameBoxes, FrameFeats))
|
|
|
|
TimeStamp = np.array(timestamp, dtype = np.float32)
|
|
DimesDiff = np.diff((timestamp))
|
|
|
|
return Videos
|
|
|
|
def inference_image(image, detections):
|
|
H, W, _ = np.shape(image)
|
|
imgs = []
|
|
batch_patches = []
|
|
patches = []
|
|
for d in range(np.size(detections, 0)):
|
|
tlbr = detections[d, :4].astype(np.int_)
|
|
tlbr[0] = max(0, tlbr[0])
|
|
tlbr[1] = max(0, tlbr[1])
|
|
tlbr[2] = min(W - 1, tlbr[2])
|
|
tlbr[3] = min(H - 1, tlbr[3])
|
|
img1 = image[tlbr[1]:tlbr[3], tlbr[0]:tlbr[2], :]
|
|
|
|
img = img1[:, :, ::-1].copy() # the model expects RGB inputs
|
|
patch = ReIDEncoder.transform(img)
|
|
|
|
imgs.append(img1)
|
|
# patch = patch.to(device=self.device).half()
|
|
if str(ReIDEncoder.device) != "cpu":
|
|
patch = patch.to(device=ReIDEncoder.device).half()
|
|
else:
|
|
patch = patch.to(device=ReIDEncoder.device)
|
|
|
|
patches.append(patch)
|
|
if (d + 1) % ReIDEncoder.batch_size == 0:
|
|
patches = torch.stack(patches, dim=0)
|
|
batch_patches.append(patches)
|
|
patches = []
|
|
|
|
if len(patches):
|
|
patches = torch.stack(patches, dim=0)
|
|
batch_patches.append(patches)
|
|
|
|
features = np.zeros((0, ReIDEncoder.embedding_size))
|
|
for patches in batch_patches:
|
|
pred = ReIDEncoder.model(patches)
|
|
pred[torch.isinf(pred)] = 1.0
|
|
feat = pred.cpu().data.numpy()
|
|
features = np.vstack((features, feat))
|
|
|
|
return imgs, features
|
|
|
|
def test_dog():
|
|
|
|
|
|
datapath = r"D:\datasets\ym\Img_ResnetData\dog_224x224\dog_224x224.txt"
|
|
with open(datapath, 'r') as file:
|
|
lines = file.readlines()
|
|
dlist = lines[0].split(',')
|
|
dfloat = [float(d) for d in dlist]
|
|
|
|
afeat = np.array(dfloat).reshape(1, -1)
|
|
|
|
|
|
|
|
imgpath = r"D:\datasets\ym\Img_ResnetData\dog_224x224\dog_224x224.jpg"
|
|
image = cv2.imread(imgpath)
|
|
patches = []
|
|
|
|
img = image[:, :, ::-1].copy() # the model expects RGB inputs
|
|
patch = ReIDEncoder.transform(img)
|
|
|
|
patch = patch.to(device=ReIDEncoder.device)
|
|
|
|
patches.append(patch)
|
|
patches = torch.stack(patches, dim=0)
|
|
pred = ReIDEncoder.model(patches)
|
|
pred[torch.isinf(pred)] = 1.0
|
|
bfeat = pred.cpu().data.numpy()
|
|
|
|
|
|
|
|
|
|
aafeat = afeat / np.linalg.norm(afeat, ord=2, axis=1, keepdims=True)
|
|
bbfeat = bfeat / np.linalg.norm(bfeat, ord=2, axis=1, keepdims=True)
|
|
|
|
|
|
cost_matrix = 1 - np.maximum(0.0, cdist(aafeat, bbfeat, 'cosine'))
|
|
|
|
|
|
|
|
|
|
print("Done!!!")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def main():
|
|
imgpath = r"D:\datasets\ym\Img_ResnetData\20240531-103547_0354b1cb-53fa-48de-86cd-ac3c5b127ada_6921168593576\3568800050000_0.jpeg"
|
|
datapath = r"D:\datasets\ym\Img_ResnetData\0_tracker_inout.data"
|
|
savepath = r"D:\datasets\ym\Img_ResnetData\result"
|
|
|
|
image = cv2.imread(imgpath)
|
|
|
|
|
|
Videos = read_data_file(datapath)
|
|
|
|
bboxes, afeats = Videos[0][0][0], Videos[0][1][0]
|
|
imgs, bfeats = inference_image(image, bboxes)
|
|
|
|
|
|
aafeats = afeats / np.linalg.norm(afeats, ord=2, axis=1, keepdims=True)
|
|
bbfeats = bfeats / np.linalg.norm(bfeats, ord=2, axis=1, keepdims=True)
|
|
|
|
|
|
cost_matrix = 1 - np.maximum(0.0, cdist(aafeats, bbfeats, 'cosine'))
|
|
|
|
|
|
for i, img in enumerate(imgs):
|
|
cv2.imwrite(savepath + f"\{i}.png", img)
|
|
|
|
|
|
|
|
print("Done!!!!")
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|
|
# test_dog()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|