modify 1:1 比对方式
This commit is contained in:
王庆刚
156
track_reid.py
156
track_reid.py
@ -64,6 +64,9 @@ from contrast.feat_extract.config import config as conf
|
||||
from contrast.feat_extract.inference import FeatsInterface
|
||||
ReIDEncoder = FeatsInterface(conf)
|
||||
|
||||
IMG_FORMATS = '.bmp', '.dng', '.jpeg', '.jpg', '.mpo', '.png', '.tif', '.tiff', '.webp', '.pfm' # include image suffixes
|
||||
VID_FORMATS = '.asf', '.avi', '.gif', '.m4v', '.mkv', '.mov', '.mp4', '.mpeg', '.mpg', '.ts', '.wmv' # include video suffixes
|
||||
|
||||
# from tracking.trackers.reid.reid_interface import ReIDInterface
|
||||
# from tracking.trackers.reid.config import config as ReIDConfig
|
||||
# ReIDEncoder = ReIDInterface(ReIDConfig)
|
||||
@ -141,6 +144,9 @@ def yolo_resnet_tracker(
|
||||
name='exp', # save results to project/name
|
||||
save_dir = '',
|
||||
|
||||
is_save_img = False,
|
||||
is_save_video = True,
|
||||
|
||||
tracker_yaml = "./tracking/trackers/cfg/botsort.yaml",
|
||||
imgsz=(640, 640), # inference size (height, width)
|
||||
conf_thres=0.25, # confidence threshold
|
||||
@ -153,18 +159,15 @@ def yolo_resnet_tracker(
|
||||
save_csv=False, # save results in CSV format
|
||||
save_conf=False, # save confidences in --save-txt labels
|
||||
save_crop=False, # save cropped prediction boxes
|
||||
|
||||
nosave=False, # do not save images/videos
|
||||
is_save_img = False,
|
||||
is_save_video = True,
|
||||
|
||||
|
||||
update=False, # update all models
|
||||
exist_ok=False, # existing project/name ok, do not increment
|
||||
|
||||
classes=None, # filter by class: --class 0, or --class 0 2 3
|
||||
agnostic_nms=False, # class-agnostic NMS
|
||||
augment=False, # augmented inference
|
||||
visualize=False, # visualize features
|
||||
update=False, # update all models
|
||||
exist_ok=False, # existing project/name ok, do not increment
|
||||
|
||||
line_thickness=3, # bounding box thickness (pixels)
|
||||
hide_labels=False, # hide labels
|
||||
hide_conf=False, # hide confidencesL
|
||||
@ -179,6 +182,8 @@ def yolo_resnet_tracker(
|
||||
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
|
||||
stride, names, pt = model.stride, model.names, model.pt
|
||||
imgsz = check_img_size(imgsz, s=stride) # check image size
|
||||
|
||||
|
||||
|
||||
# Dataloader
|
||||
bs = 1 # batch_size
|
||||
@ -203,8 +208,8 @@ def yolo_resnet_tracker(
|
||||
|
||||
# Inference
|
||||
with dt[1]:
|
||||
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
|
||||
pred = model(im, augment=augment, visualize=visualize)
|
||||
visualize = increment_path(project / Path(path).stem, mkdir=True) if visualize else False
|
||||
pred = model(im, augment=augment, visualize=False)
|
||||
|
||||
# NMS
|
||||
with dt[2]:
|
||||
@ -273,19 +278,27 @@ def yolo_resnet_tracker(
|
||||
annotator.box_label(xyxy, label, color=color)
|
||||
|
||||
'''====== Save results (image and video) ======'''
|
||||
save_path = str(save_dir / Path(path).name) # 带有后缀名
|
||||
# save_path = str(save_dir / Path(path).name) # 带有后缀名
|
||||
im0 = annotator.result()
|
||||
if is_save_img:
|
||||
save_path_img, ext = os.path.splitext(save_path)
|
||||
save_path_img = str(save_dir / Path(path).stem)
|
||||
if dataset.mode == 'image':
|
||||
imgpath = save_path_img + ".png"
|
||||
else:
|
||||
imgpath = save_path_img + f"_{frameId}.png"
|
||||
cv2.imwrite(Path(imgpath), im0)
|
||||
|
||||
if dataset.mode == 'video' and is_save_video:
|
||||
if vid_path[i] != save_path: # new video
|
||||
vid_path[i] = save_path
|
||||
# if dataset.mode == 'video' and is_save_video:
|
||||
|
||||
if is_save_video:
|
||||
if dataset.mode == 'video':
|
||||
vdieo_path = str(save_dir / Path(path).stem) + '.mp4' # 带有后缀名
|
||||
else:
|
||||
videoname = str(Path(path).stem).split('_')[0] + '.mp4'
|
||||
vdieo_path = str(save_dir / videoname)
|
||||
|
||||
if vid_path[i] != vdieo_path: # new video
|
||||
vid_path[i] = vdieo_path
|
||||
if isinstance(vid_writer[i], cv2.VideoWriter):
|
||||
vid_writer[i].release() # release previous video writer
|
||||
if vid_cap: # video
|
||||
@ -293,9 +306,9 @@ def yolo_resnet_tracker(
|
||||
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
else: # stream
|
||||
fps, w, h = 30, im0.shape[1], im0.shape[0]
|
||||
save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
|
||||
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
|
||||
fps, w, h = 25, im0.shape[1], im0.shape[0]
|
||||
vdieo_path = str(Path(vdieo_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
|
||||
vid_writer[i] = cv2.VideoWriter(vdieo_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
|
||||
vid_writer[i].write(im0)
|
||||
|
||||
# Print time (inference-only)
|
||||
@ -344,6 +357,9 @@ def run(
|
||||
vid_stride=1, # video frame-rate stride
|
||||
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
|
||||
):
|
||||
'''
|
||||
source: 视频文件或图像列表
|
||||
'''
|
||||
source = str(source)
|
||||
# filename = os.path.split(source)[-1]
|
||||
|
||||
@ -355,6 +371,9 @@ def run(
|
||||
if is_url and is_file:
|
||||
source = check_file(source) # download
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# spth = source.split('\\')[-2] + "_" + Path(source).stem
|
||||
save_dir = Path(project) / Path(source.split('\\')[-2] + "_" + str(Path(source).stem))
|
||||
@ -440,8 +459,7 @@ def run(
|
||||
|
||||
|
||||
|
||||
p = Path(p) # to Path
|
||||
save_path = str(save_dir / p.name) # im.jpg
|
||||
|
||||
s += '%gx%g ' % im.shape[2:] # print string
|
||||
|
||||
# im0_ant = im0.copy()
|
||||
@ -552,28 +570,33 @@ def run(
|
||||
|
||||
# Save results (image and video with tracking)
|
||||
im0 = annotator.result()
|
||||
save_path_img, ext = os.path.splitext(save_path)
|
||||
|
||||
p = Path(p) # to Path
|
||||
save_path = str(save_dir / p.name) # im.jpg
|
||||
if save_img:
|
||||
save_path_img, ext = os.path.splitext(save_path)
|
||||
if dataset.mode == 'image':
|
||||
imgpath = save_path_img + f"_{dataset}.png"
|
||||
imgpath = save_path_img + ".png"
|
||||
else:
|
||||
imgpath = save_path_img + f"_{frameId}.png"
|
||||
|
||||
cv2.imwrite(Path(imgpath), im0)
|
||||
|
||||
if vid_path[i] != save_path: # new video
|
||||
vid_path[i] = save_path
|
||||
if isinstance(vid_writer[i], cv2.VideoWriter):
|
||||
vid_writer[i].release() # release previous video writer
|
||||
if vid_cap: # video
|
||||
fps = vid_cap.get(cv2.CAP_PROP_FPS)
|
||||
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
else: # stream
|
||||
fps, w, h = 30, im0.shape[1], im0.shape[0]
|
||||
save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
|
||||
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
|
||||
vid_writer[i].write(im0)
|
||||
if dataset.mode == 'video':
|
||||
|
||||
if vid_path[i] != save_path: # new video
|
||||
vid_path[i] = save_path
|
||||
if isinstance(vid_writer[i], cv2.VideoWriter):
|
||||
vid_writer[i].release() # release previous video writer
|
||||
if vid_cap: # video
|
||||
fps = vid_cap.get(cv2.CAP_PROP_FPS)
|
||||
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
else: # stream
|
||||
fps, w, h = 30, im0.shape[1], im0.shape[0]
|
||||
save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
|
||||
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
|
||||
vid_writer[i].write(im0)
|
||||
|
||||
|
||||
# Print time (inference-only)
|
||||
LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
|
||||
@ -672,40 +695,25 @@ def parse_opt():
|
||||
print_args(vars(opt))
|
||||
return opt
|
||||
|
||||
def find_files_in_nested_dirs(root_dir):
|
||||
def find_video_imgs(root_dir):
|
||||
all_files = []
|
||||
extensions = ['.mp4']
|
||||
for dirpath, dirnames, filenames in os.walk(root_dir):
|
||||
for filename in filenames:
|
||||
file, ext = os.path.splitext(filename)
|
||||
if ext in extensions:
|
||||
if ext in IMG_FORMATS + VID_FORMATS:
|
||||
all_files.append(os.path.join(dirpath, filename))
|
||||
return all_files
|
||||
|
||||
print('=======')
|
||||
|
||||
def main(opt):
|
||||
check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
|
||||
optdict = vars(opt)
|
||||
|
||||
p = r"D:\datasets\ym"
|
||||
p = r"D:\exhibition\images\153112511_0_seek_105.mp4"
|
||||
|
||||
optdict["project"] = r"D:\exhibition\result"
|
||||
|
||||
files = []
|
||||
if os.path.isdir(p):
|
||||
files.extend(sorted(glob.glob(os.path.join(p, '*.*'))))
|
||||
optdict["source"] = files
|
||||
elif os.path.isfile(p):
|
||||
optdict["source"] = p
|
||||
|
||||
run(**optdict)
|
||||
|
||||
|
||||
def main_loop(opt):
|
||||
def main():
|
||||
'''
|
||||
run(): 单张图像或单个视频文件的推理，不支持图像序列，
|
||||
'''
|
||||
|
||||
check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
|
||||
|
||||
opt = parse_opt()
|
||||
optdict = vars(opt)
|
||||
|
||||
# p = r"D:\datasets\ym\永辉测试数据_比对"
|
||||
@ -714,28 +722,22 @@ def main_loop(opt):
|
||||
# p = r"D:\datasets\ym\实验室测试"
|
||||
# p = r"D:\datasets\ym\永辉双摄视频\新建文件夹"
|
||||
# p = r"\\192.168.1.28\share\测试_202406\0723\0723_2\20240723-112522_"
|
||||
p = r"D:\datasets\ym\联华中环"
|
||||
# p = r"D:\datasets\ym\联华中环"
|
||||
p = r"D:\exhibition\images\153112511_0_seek_105.mp4"
|
||||
# p = r"D:\exhibition\images\image"
|
||||
|
||||
k = 0
|
||||
|
||||
optdict["project"] = r"D:\exhibition\result"
|
||||
if os.path.isdir(p):
|
||||
files = find_files_in_nested_dirs(p)
|
||||
|
||||
# files = [r"D:\datasets\ym\广告板遮挡测试\8\6926636301004_20240508-175300_back_addGood_70f754088050_215_17327712807.mp4",
|
||||
# r"D:\datasets\ym\videos\标记视频\test_20240402-173935_6920152400975_back_174037372.mp4",
|
||||
# r"D:\datasets\ym\videos\标记视频\test_20240402-173935_6920152400975_front_174037379.mp4",
|
||||
# r"D:\datasets\ym\广告板遮挡测试\8\2500441577966_20240508-175946_front_addGood_70f75407b7ae_155_17788571404.mp4"
|
||||
# ]
|
||||
|
||||
# files = [r"\\192.168.1.28\share\测试_202406\0723\0723_2\20240723-095838_\1_seek_193.mp4"]
|
||||
|
||||
|
||||
files = find_video_imgs(p)
|
||||
k = 0
|
||||
for file in files:
|
||||
optdict["source"] = file
|
||||
run(**optdict)
|
||||
|
||||
# k += 1
|
||||
# if k == 10:
|
||||
# break
|
||||
k += 1
|
||||
if k == 1:
|
||||
break
|
||||
elif os.path.isfile(p):
|
||||
optdict["source"] = p
|
||||
run(**optdict)
|
||||
@ -744,10 +746,8 @@ def main_loop(opt):
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
opt = parse_opt()
|
||||
|
||||
main(opt)
|
||||
# main_loop(opt)
|
||||
main()
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user