noham/Stage-2024
1
0
Fork 0
mirror of https://github.com/NohamR/Stage-2024.git synced 2026-01-11 08:48:15 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

clean

Browse Source
This commit is contained in:
√(noham)²
2024-07-18 00:42:59 +02:00
parent 5c9313c4ca
commit 3cf13b815c
180 changed files with 34499 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

180
yolov7-tracker-example/tools/convert_MOT17_to_yolo.py Normal file
View File

@@ -0,0 +1,180 @@
"""
将UAVDT转换为yolo v5格式
class_id, xc_norm, yc_norm, w_norm, h_norm
"""
import os
import os.path as osp
import argparse
import cv2
import glob
import numpy as np
import random
DATA_ROOT = '/data/wujiapeng/datasets/MOT17/'
image_wh_dict = {} # seq->(w,h) 字典 用于归一化
def generate_imgs_and_labels(opts):
"""
产生图片路径的txt文件以及yolo格式真值文件
"""
if opts.split == 'test':
seq_list = os.listdir(osp.join(DATA_ROOT, 'test'))
else:
seq_list = os.listdir(osp.join(DATA_ROOT, 'train'))
seq_list = [item for item in seq_list if 'FRCNN' in item] # 只取一个FRCNN即可
if 'val' in opts.split: opts.half = True # 验证集取训练集的一半
print('--------------------------')
print(f'Total {len(seq_list)} seqs!!')
print(seq_list)
if opts.random:
random.shuffle(seq_list)
# 定义类别 MOT只有一类
CATEGOTY_ID = 0 # pedestrian
# 定义帧数范围
frame_range = {'start': 0.0, 'end': 1.0}
if opts.half: # half 截取一半
frame_range['end'] = 0.5
if opts.split == 'test':
process_train_test(seqs=seq_list, frame_range=frame_range, cat_id=CATEGOTY_ID, split='test')
else:
process_train_test(seqs=seq_list, frame_range=frame_range, cat_id=CATEGOTY_ID, split=opts.split)
def process_train_test(seqs: list, frame_range: dict, cat_id: int = 0, split: str = 'trian') -> None:
"""
处理MOT17的train 或 test
由于操作相似 故另写函数
"""
for seq in seqs:
print(f'Dealing with {split} dataset...')
img_dir = osp.join(DATA_ROOT, 'train', seq, 'img1') if split != 'test' else osp.join(DATA_ROOT, 'test', seq, 'img1') # 图片路径
imgs = sorted(os.listdir(img_dir)) # 所有图片的相对路径
seq_length = len(imgs) # 序列长度
if split != 'test':
# 求解图片高宽
img_eg = cv2.imread(osp.join(img_dir, imgs[0]))
w0, h0 = img_eg.shape[1], img_eg.shape[0] # 原始高宽
ann_of_seq_path = os.path.join(img_dir, '../', 'gt', 'gt.txt') # GT文件路径
ann_of_seq = np.loadtxt(ann_of_seq_path, dtype=np.float32, delimiter=',') # GT内容
gt_to_path = osp.join(DATA_ROOT, 'labels', split, seq) # 要写入的真值文件夹
# 如果不存在就创建
if not osp.exists(gt_to_path):
os.makedirs(gt_to_path)
exist_gts = [] # 初始化该列表 每个元素对应该seq的frame中有无真值框
# 如果没有 就在train.txt产生图片路径
for idx, img in enumerate(imgs):
# img 形如: img000001.jpg
if idx < int(seq_length * frame_range['start']) or idx > int(seq_length * frame_range['end']):
continue
# 第一步 产生图片软链接
# print('step1, creating imgs symlink...')
if opts.generate_imgs:
img_to_path = osp.join(DATA_ROOT, 'images', split, seq) # 该序列图片存储位置
if not osp.exists(img_to_path):
os.makedirs(img_to_path)
os.symlink(osp.join(img_dir, img),
osp.join(img_to_path, img)) # 创建软链接
# 第二步 产生真值文件
# print('step2, generating gt files...')
ann_of_current_frame = ann_of_seq[ann_of_seq[:, 0] == float(idx + 1), :] # 筛选真值文件里本帧的目标信息
exist_gts.append(True if ann_of_current_frame.shape[0] != 0 else False)
gt_to_file = osp.join(gt_to_path, img[: -4] + '.txt')
with open(gt_to_file, 'w') as f_gt:
for i in range(ann_of_current_frame.shape[0]):
if int(ann_of_current_frame[i][6]) == 1 and int(ann_of_current_frame[i][7]) == 1 \
and float(ann_of_current_frame[i][8]) > 0.25:
# bbox xywh
x0, y0 = int(ann_of_current_frame[i][2]), int(ann_of_current_frame[i][3])
x0, y0 = max(x0, 0), max(y0, 0)
w, h = int(ann_of_current_frame[i][4]), int(ann_of_current_frame[i][5])
xc, yc = x0 + w // 2, y0 + h // 2 # 中心点 x y
# 归一化
xc, yc = xc / w0, yc / h0
xc, yc = min(xc, 1.0), min(yc, 1.0)
w, h = w / w0, h / h0
w, h = min(w, 1.0), min(h, 1.0)
assert w <= 1 and h <= 1, f'{w}, {h} must be normed, original size{w0}, {h0}'
assert xc >= 0 and yc >= 0, f'{x0}, {y0} must be positve'
assert xc <= 1 and yc <= 1, f'{x0}, {y0} must be le than 1'
category_id = cat_id
write_line = '{:d} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format(
category_id, xc, yc, w, h)
f_gt.write(write_line)
f_gt.close()
else: # test 只产生图片软链接
for idx, img in enumerate(imgs):
# img 形如: img000001.jpg
if idx < int(seq_length * frame_range['start']) or idx > int(seq_length * frame_range['end']):
continue
# 第一步 产生图片软链接
# print('step1, creating imgs symlink...')
if opts.generate_imgs:
img_to_path = osp.join(DATA_ROOT, 'images', split, seq) # 该序列图片存储位置
if not osp.exists(img_to_path):
os.makedirs(img_to_path)
os.symlink(osp.join(img_dir, img),
osp.join(img_to_path, img)) # 创建软链接
# 第三步 产生图片索引train.txt等
print(f'generating img index file of {seq}')
to_file = os.path.join('./mot17/', split + '.txt')
with open(to_file, 'a') as f:
for idx, img in enumerate(imgs):
if idx < int(seq_length * frame_range['start']) or idx > int(seq_length * frame_range['end']):
continue
if split == 'test' or exist_gts[idx]:
f.write('MOT17/' + 'images/' + split + '/' \
+ seq + '/' + img + '\n')
f.close()
if __name__ == '__main__':
if not osp.exists('./mot17'):
os.system('mkdir mot17')
parser = argparse.ArgumentParser()
parser.add_argument('--split', type=str, default='train', help='train, test or val')
parser.add_argument('--generate_imgs', action='store_true', help='whether generate soft link of imgs')
parser.add_argument('--certain_seqs', action='store_true', help='for debug')
parser.add_argument('--half', action='store_true', help='half frames')
parser.add_argument('--ratio', type=float, default=0.8, help='ratio of test dataset devide train dataset')
parser.add_argument('--random', action='store_true', help='random split train and test')
opts = parser.parse_args()
generate_imgs_and_labels(opts)
# python tools/convert_MOT17_to_yolo.py --split train --generate_imgs

159
yolov7-tracker-example/tools/convert_UAVDT_to_yolo.py Normal file
View File

@@ -0,0 +1,159 @@
"""
将UAVDT转换为yolo v5格式
class_id, xc_norm, yc_norm, w_norm, h_norm
"""
import os
import os.path as osp
import argparse
import cv2
import glob
import numpy as np
import random
DATA_ROOT = '/data/wujiapeng/datasets/UAVDT/'
image_wh_dict = {} # seq->(w,h) 字典 用于归一化
def generate_imgs_and_labels(opts):
"""
产生图片路径的txt文件以及yolo格式真值文件
"""
seq_list = os.listdir(osp.join(DATA_ROOT, 'UAV-benchmark-M'))
print('--------------------------')
print(f'Total {len(seq_list)} seqs!!')
# 划分train test
if opts.random:
random.shuffle(seq_list)
bound = int(opts.ratio * len(seq_list))
train_seq_list = seq_list[: bound]
test_seq_list = seq_list[bound:]
del bound
print(f'train dataset: {train_seq_list}')
print(f'test dataset: {test_seq_list}')
print('--------------------------')
if not osp.exists('./uavdt/'):
os.makedirs('./uavdt/')
# 定义类别 UAVDT只有一类
CATEGOTY_ID = 0 # car
# 定义帧数范围
frame_range = {'start': 0.0, 'end': 1.0}
if opts.half: # half 截取一半
frame_range['end'] = 0.5
# 分别处理train与test
process_train_test(train_seq_list, frame_range, CATEGOTY_ID, 'train')
process_train_test(test_seq_list, {'start': 0.0, 'end': 1.0}, CATEGOTY_ID, 'test')
print('All Done!!')
def process_train_test(seqs: list, frame_range: dict, cat_id: int = 0, split: str = 'trian') -> None:
"""
处理UAVDT的train 或 test
由于操作相似 故另写函数
"""
for seq in seqs:
print('Dealing with train dataset...')
img_dir = osp.join(DATA_ROOT, 'UAV-benchmark-M', seq, 'img1') # 图片路径
imgs = sorted(os.listdir(img_dir)) # 所有图片的相对路径
seq_length = len(imgs) # 序列长度
# 求解图片高宽
img_eg = cv2.imread(osp.join(img_dir, imgs[0]))
w0, h0 = img_eg.shape[1], img_eg.shape[0] # 原始高宽
ann_of_seq_path = os.path.join(img_dir, '../', 'gt', 'gt.txt') # GT文件路径
ann_of_seq = np.loadtxt(ann_of_seq_path, dtype=np.float32, delimiter=',') # GT内容
gt_to_path = osp.join(DATA_ROOT, 'labels', split, seq) # 要写入的真值文件夹
# 如果不存在就创建
if not osp.exists(gt_to_path):
os.makedirs(gt_to_path)
exist_gts = [] # 初始化该列表 每个元素对应该seq的frame中有无真值框
# 如果没有 就在train.txt产生图片路径
for idx, img in enumerate(imgs):
# img 形如: img000001.jpg
if idx < int(seq_length * frame_range['start']) or idx > int(seq_length * frame_range['end']):
continue
# 第一步 产生图片软链接
# print('step1, creating imgs symlink...')
if opts.generate_imgs:
img_to_path = osp.join(DATA_ROOT, 'images', split, seq) # 该序列图片存储位置
if not osp.exists(img_to_path):
os.makedirs(img_to_path)
os.symlink(osp.join(img_dir, img),
osp.join(img_to_path, img)) # 创建软链接
# 第二步 产生真值文件
# print('step2, generating gt files...')
ann_of_current_frame = ann_of_seq[ann_of_seq[:, 0] == float(idx + 1), :] # 筛选真值文件里本帧的目标信息
exist_gts.append(True if ann_of_current_frame.shape[0] != 0 else False)
gt_to_file = osp.join(gt_to_path, img[:-4] + '.txt')
with open(gt_to_file, 'w') as f_gt:
for i in range(ann_of_current_frame.shape[0]):
if int(ann_of_current_frame[i][6]) == 1:
# bbox xywh
x0, y0 = int(ann_of_current_frame[i][2]), int(ann_of_current_frame[i][3])
w, h = int(ann_of_current_frame[i][4]), int(ann_of_current_frame[i][5])
xc, yc = x0 + w // 2, y0 + h // 2 # 中心点 x y
# 归一化
xc, yc = xc / w0, yc / h0
w, h = w / w0, h / h0
category_id = cat_id
write_line = '{:d} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format(
category_id, xc, yc, w, h)
f_gt.write(write_line)
f_gt.close()
# 第三步 产生图片索引train.txt等
print(f'generating img index file of {seq}')
to_file = os.path.join('./uavdt/', split + '.txt')
with open(to_file, 'a') as f:
for idx, img in enumerate(imgs):
if idx < int(seq_length * frame_range['start']) or idx > int(seq_length * frame_range['end']):
continue
if exist_gts[idx]:
f.write('UAVDT/' + 'images/' + split + '/' \
+ seq + '/' + img + '\n')
f.close()
if __name__ == '__main__':
if not osp.exists('./uavdt'):
os.system('mkdir ./uavdt')
else:
os.system('rm -rf ./uavdt/*')
parser = argparse.ArgumentParser()
parser.add_argument('--generate_imgs', action='store_true', help='whether generate soft link of imgs')
parser.add_argument('--certain_seqs', action='store_true', help='for debug')
parser.add_argument('--half', action='store_true', help='half frames')
parser.add_argument('--ratio', type=float, default=0.8, help='ratio of test dataset devide train dataset')
parser.add_argument('--random', action='store_true', help='random split train and test')
opts = parser.parse_args()
generate_imgs_and_labels(opts)
# python tools/convert_UAVDT_to_yolo.py --generate_imgs --half --random

182
yolov7-tracker-example/tools/convert_VisDrone_to_yolo.py Normal file
View File

@@ -0,0 +1,182 @@
"""
将VisDrone转换为yolo v5格式
class_id, xc_norm, yc_norm, w_norm, h_norm
"""
import os
import os.path as osp
import argparse
import cv2
import glob
DATA_ROOT = '/data/wujiapeng/datasets/VisDrone2019/VisDrone2019'
# 以下两个seqs只跟踪车的时候有用
certain_seqs = ['uav0000071_03240_v', 'uav0000072_04488_v','uav0000072_05448_v', 'uav0000072_06432_v','uav0000124_00944_v','uav0000126_00001_v','uav0000138_00000_v','uav0000145_00000_v','uav0000150_02310_v','uav0000222_03150_v','uav0000239_12336_v','uav0000243_00001_v',
'uav0000248_00001_v','uav0000263_03289_v','uav0000266_03598_v','uav0000273_00001_v','uav0000279_00001_v','uav0000281_00460_v','uav0000289_00001_v','uav0000289_06922_v','uav0000307_00000_v',
'uav0000308_00000_v','uav0000308_01380_v','uav0000326_01035_v','uav0000329_04715_v','uav0000361_02323_v','uav0000366_00001_v']
ignored_seqs = ['uav0000013_00000_v', 'uav0000013_01073_v', 'uav0000013_01392_v',
'uav0000020_00406_v', 'uav0000079_00480_v',
'uav0000084_00000_v', 'uav0000099_02109_v', 'uav0000086_00000_v',
'uav0000073_00600_v', 'uav0000073_04464_v', 'uav0000088_00290_v']
image_wh_dict = {} # seq->(w,h) 字典 用于归一化
def generate_imgs(split_name='VisDrone2019-MOT-train', generate_imgs=True, if_certain_seqs=False, car_only=False):
"""
产生图片文件夹 例如 VisDrone/images/VisDrone2019-MOT-train/uav0000076_00720_v/000010.jpg
同时产生序列->高,宽的字典 便于后续
split: str, 'VisDrone2019-MOT-train', 'VisDrone2019-MOT-val' or 'VisDrone2019-MOT-test-dev'
if_certain_seqs: bool, use for debug.
"""
if not if_certain_seqs:
seq_list = os.listdir(osp.join(DATA_ROOT, split_name, 'sequences')) # 所有序列名称
else:
seq_list = certain_seqs
if car_only: # 只跟踪车就忽略行人多的视频
seq_list = [seq for seq in seq_list if seq not in ignored_seqs]
# 遍历所有序列 给图片创建软链接 同时更新seq->(w,h)字典
if_write_txt = True if glob.glob('./visdrone/*.txt') else False
# if_write_txt = True if not osp.exists(f'./visdrone/.txt') else False # 是否需要写txt 用于生成visdrone.train
if not if_write_txt:
for seq in seq_list:
img_dir = osp.join(DATA_ROOT, split_name, 'sequences', seq) # 该序列下所有图片路径
imgs = sorted(os.listdir(img_dir)) # 所有图片
if generate_imgs:
to_path = osp.join(DATA_ROOT, 'images', split_name, seq) # 该序列图片存储位置
if not osp.exists(to_path):
os.makedirs(to_path)
for img in imgs: # 遍历该序列下的图片
os.symlink(osp.join(img_dir, img),
osp.join(to_path, img)) # 创建软链接
img_sample = cv2.imread(osp.join(img_dir, imgs[0])) # 每个序列第一张图片 用于获取w, h
w, h = img_sample.shape[1], img_sample.shape[0] # w, h
image_wh_dict[seq] = (w, h) # 更新seq->(w,h) 字典
# print(image_wh_dict)
# return
else:
with open('./visdrone.txt', 'a') as f:
for seq in seq_list:
img_dir = osp.join(DATA_ROOT, split_name, 'sequences', seq) # 该序列下所有图片路径
imgs = sorted(os.listdir(img_dir)) # 所有图片
if generate_imgs:
to_path = osp.join(DATA_ROOT, 'images', split_name, seq) # 该序列图片存储位置
if not osp.exists(to_path):
os.makedirs(to_path)
for img in imgs: # 遍历该序列下的图片
f.write('VisDrone2019/' + 'VisDrone2019/' + 'images/' + split_name + '/' \
+ seq + '/' + img + '\n')
os.symlink(osp.join(img_dir, img),
osp.join(to_path, img)) # 创建软链接
img_sample = cv2.imread(osp.join(img_dir, imgs[0])) # 每个序列第一张图片 用于获取w, h
w, h = img_sample.shape[1], img_sample.shape[0] # w, h
image_wh_dict[seq] = (w, h) # 更新seq->(w,h) 字典
f.close()
if if_certain_seqs: # for debug
print(image_wh_dict)
def generate_labels(split='VisDrone2019-MOT-train', if_certain_seqs=False, car_only=False):
"""
split: str, 'train', 'val' or 'test'
if_certain_seqs: bool, use for debug.
"""
# from choose_anchors import image_wh_dict
# print(image_wh_dict)
if not if_certain_seqs:
seq_list = os.listdir(osp.join(DATA_ROOT, split, 'sequences')) # 序列列表
else:
seq_list = certain_seqs
if car_only: # 只跟踪车就忽略行人多的视频
seq_list = [seq for seq in seq_list if seq not in ignored_seqs]
category_list = ['4', '5', '6', '9']
else:
category_list = [str(i) for i in range(1, 11)]
# 类别ID 从0开始
category_dict = {category_list[idx]: idx for idx in range(len(category_list))}
# 每张图片分配一个txt
# 要从sequence的txt里分出来
for seq in seq_list:
seq_dir = osp.join(DATA_ROOT, split, 'annotations', seq + '.txt') # 真值文件
with open(seq_dir, 'r') as f:
lines = f.readlines()
for row in lines:
current_line = row.split(',')
frame = current_line[0] # 第几帧
if current_line[6] == '0' or current_line[7] not in category_list:
continue
to_file = osp.join(DATA_ROOT, 'labels', split, seq) # 要写入的文件名
# 如果不存在就创建
if not osp.exists(to_file):
os.makedirs(to_file)
to_file = osp.join(to_file, frame.zfill(7) + '.txt')
category_id = category_dict[current_line[7]]
x0, y0 = int(current_line[2]), int(current_line[3]) # 左上角 x y
w, h = int(current_line[4]), int(current_line[5]) # 宽 高
x_c, y_c = x0 + w // 2, y0 + h // 2 # 中心点 x y
image_w, image_h = image_wh_dict[seq][0], image_wh_dict[seq][1] # 图像高宽
# 归一化
w, h = w / image_w, h / image_h
x_c, y_c = x_c / image_w, y_c / image_h
with open(to_file, 'a') as f_to:
write_line = '{:d} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format(
category_id, x_c, y_c, w, h)
f_to.write(write_line)
f_to.close()
f.close()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--split', type=str, default='VisDrone2019-MOT-train', help='train or test')
parser.add_argument('--generate_imgs', action='store_true', help='whether generate soft link of imgs')
parser.add_argument('--car_only', action='store_true', help='only cars')
parser.add_argument('--if_certain_seqs', action='store_true', help='for debug')
opt = parser.parse_args()
print('generating images...')
generate_imgs(opt.split, opt.generate_imgs, opt.if_certain_seqs, opt.car_only)
print('generating labels...')
generate_labels(opt.split, opt.if_certain_seqs, opt.car_only)
print('Done!')
# python convert_VisDrone_to_yolo.py --split VisDrone2019-MOT-train
# python convert_VisDrone_to_yolo.py --split VisDrone2019-MOT-train --car_only --if_certain_seqs

168
yolov7-tracker-example/tools/convert_VisDrone_to_yolov2.py Normal file
View File

@@ -0,0 +1,168 @@
"""
将VisDrone转换为yolo v5格式
class_id, xc_norm, yc_norm, w_norm, h_norm
改动:
1. 将产生img和label函数合成一个
2. 增加如果无label就不产生当前img路径的功能
3. 增加half选项 每个视频截取一半
"""
import os
import os.path as osp
import argparse
import cv2
import glob
import numpy as np
DATA_ROOT = '/data/wujiapeng/datasets/VisDrone2019/VisDrone2019'
# 以下两个seqs只跟踪车的时候有用
certain_seqs = ['uav0000071_03240_v', 'uav0000072_04488_v','uav0000072_05448_v', 'uav0000072_06432_v','uav0000124_00944_v','uav0000126_00001_v','uav0000138_00000_v','uav0000145_00000_v','uav0000150_02310_v','uav0000222_03150_v','uav0000239_12336_v','uav0000243_00001_v',
'uav0000248_00001_v','uav0000263_03289_v','uav0000266_03598_v','uav0000273_00001_v','uav0000279_00001_v','uav0000281_00460_v','uav0000289_00001_v','uav0000289_06922_v','uav0000307_00000_v',
'uav0000308_00000_v','uav0000308_01380_v','uav0000326_01035_v','uav0000329_04715_v','uav0000361_02323_v','uav0000366_00001_v']
ignored_seqs = ['uav0000013_00000_v', 'uav0000013_01073_v', 'uav0000013_01392_v',
'uav0000020_00406_v', 'uav0000079_00480_v',
'uav0000084_00000_v', 'uav0000099_02109_v', 'uav0000086_00000_v',
'uav0000073_00600_v', 'uav0000073_04464_v', 'uav0000088_00290_v']
image_wh_dict = {} # seq->(w,h) 字典 用于归一化
def generate_imgs_and_labels(opts):
"""
产生图片路径的txt文件以及yolo格式真值文件
"""
if not opts.certain_seqs:
seq_list = os.listdir(osp.join(DATA_ROOT, opts.split_name, 'sequences')) # 所有序列名称
else:
seq_list = certain_seqs
if opts.car_only: # 只跟踪车就忽略行人多的视频
seq_list = [seq for seq in seq_list if seq not in ignored_seqs]
category_list = [4, 5, 6, 9] # 感兴趣的类别编号 List[int]
else:
category_list = [i for i in range(1, 11)]
print(f'Total {len(seq_list)} seqs!!')
if not osp.exists('./visdrone/'):
os.makedirs('./visdrone/')
# 类别ID 从0开始
category_dict = {category_list[idx]: idx for idx in range(len(category_list))}
txt_name_dict = {'VisDrone2019-MOT-train': 'train',
'VisDrone2019-MOT-val': 'val',
'VisDrone2019-MOT-test-dev': 'test'} # 产生txt文件名称对应关系
# 如果已经存在就不写了
write_txt = False if os.path.isfile(os.path.join('./visdrone', txt_name_dict[opts.split_name] + '.txt')) else True
print(f'write txt is {write_txt}')
frame_range = {'start': 0.0, 'end': 1.0}
if opts.half: # VisDrone-half 截取一半
frame_range['end'] = 0.5
# 以序列为单位进行处理
for seq in seq_list:
img_dir = osp.join(DATA_ROOT, opts.split_name, 'sequences', seq) # 该序列下所有图片路径
imgs = sorted(os.listdir(img_dir)) # 所有图片
seq_length = len(imgs) # 序列长度
img_eg = cv2.imread(os.path.join(img_dir, imgs[0])) # 序列的第一张图 用以计算高宽
w0, h0 = img_eg.shape[1], img_eg.shape[0] # 原始高宽
ann_of_seq_path = os.path.join(DATA_ROOT, opts.split_name, 'annotations', seq + '.txt') # GT文件路径
ann_of_seq = np.loadtxt(ann_of_seq_path, dtype=np.float32, delimiter=',') # GT内容
gt_to_path = osp.join(DATA_ROOT, 'labels', opts.split_name, seq) # 要写入的真值文件夹
# 如果不存在就创建
if not osp.exists(gt_to_path):
os.makedirs(gt_to_path)
exist_gts = [] # 初始化该列表 每个元素对应该seq的frame中有无真值框
# 如果没有 就在train.txt产生图片路径
for idx, img in enumerate(imgs):
# img: 相对路径 即 图片名称 0000001.jpg
if idx < int(seq_length * frame_range['start']) or idx > int(seq_length * frame_range['end']):
continue
# 第一步 产生图片软链接
# print('step1, creating imgs symlink...')
if opts.generate_imgs:
img_to_path = osp.join(DATA_ROOT, 'images', opts.split_name, seq) # 该序列图片存储位置
if not osp.exists(img_to_path):
os.makedirs(img_to_path)
os.symlink(osp.join(img_dir, img),
osp.join(img_to_path, img)) # 创建软链接
# print('Done!\n')
# 第二步 产生真值文件
# print('step2, generating gt files...')
# 根据本序列的真值文件读取
# ann_idx = int(ann_of_seq[:, 0]) == idx + 1
ann_of_current_frame = ann_of_seq[ann_of_seq[:, 0] == float(idx + 1), :] # 筛选真值文件里本帧的目标信息
exist_gts.append(True if ann_of_current_frame.shape[0] != 0 else False)
gt_to_file = osp.join(gt_to_path, img[:-4] + '.txt')
with open(gt_to_file, 'a') as f_gt:
for i in range(ann_of_current_frame.shape[0]):
category = int(ann_of_current_frame[i][7])
if int(ann_of_current_frame[i][6]) == 1 and category in category_list:
# bbox xywh
x0, y0 = int(ann_of_current_frame[i][2]), int(ann_of_current_frame[i][3])
w, h = int(ann_of_current_frame[i][4]), int(ann_of_current_frame[i][5])
xc, yc = x0 + w // 2, y0 + h // 2 # 中心点 x y
# 归一化
xc, yc = xc / w0, yc / h0
w, h = w / w0, h / h0
category_id = category_dict[category]
write_line = '{:d} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format(
category_id, xc, yc, w, h)
f_gt.write(write_line)
f_gt.close()
# print('Done!\n')
print(f'img symlink and gt files of seq {seq} Done!')
# 第三步 产生图片索引train.txt等
print(f'generating img index file of {seq}')
if write_txt:
to_file = os.path.join('./visdrone', txt_name_dict[opts.split_name] + '.txt')
with open(to_file, 'a') as f:
for idx, img in enumerate(imgs):
if idx < int(seq_length * frame_range['start']) or idx > int(seq_length * frame_range['end']):
continue
if exist_gts[idx]:
f.write('VisDrone2019/' + 'VisDrone2019/' + 'images/' + opts.split_name + '/' \
+ seq + '/' + img + '\n')
f.close()
print('All done!!')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--split_name', type=str, default='VisDrone2019-MOT-train', help='train or test')
parser.add_argument('--generate_imgs', action='store_true', help='whether generate soft link of imgs')
parser.add_argument('--car_only', action='store_true', help='only cars')
parser.add_argument('--certain_seqs', action='store_true', help='for debug')
parser.add_argument('--half', action='store_true', help='half frames')
opts = parser.parse_args()
generate_imgs_and_labels(opts)
# python tools/convert_VisDrone_to_yolov2.py --split_name VisDrone2019-MOT-train --generate_imgs --car_only --half

479
yolov7-tracker-example/tools/reparameterization.ipynb Normal file
View File

@@ -0,0 +1,479 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "d7cbe5ee",
"metadata": {},
"source": [
"# Reparameterization"
]
},
{
"cell_type": "markdown",
"id": "13393b70",
"metadata": {},
"source": [
"## YOLOv7 reparameterization"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bf53becf",
"metadata": {},
"outputs": [],
"source": [
"# import\n",
"from copy import deepcopy\n",
"from models.yolo import Model\n",
"import torch\n",
"from utils.torch_utils import select_device, is_parallel\n",
"\n",
"device = select_device('0', batch_size=1)\n",
"# model trained by cfg/training/*.yaml\n",
"ckpt = torch.load('cfg/training/yolov7.pt', map_location=device)\n",
"# reparameterized model in cfg/deploy/*.yaml\n",
"model = Model('cfg/deploy/yolov7.yaml', ch=3, nc=80).to(device)\n",
"\n",
"# copy intersect weights\n",
"state_dict = ckpt['model'].float().state_dict()\n",
"exclude = []\n",
"intersect_state_dict = {k: v for k, v in state_dict.items() if k in model.state_dict() and not any(x in k for x in exclude) and v.shape == model.state_dict()[k].shape}\n",
"model.load_state_dict(intersect_state_dict, strict=False)\n",
"model.names = ckpt['model'].names\n",
"model.nc = ckpt['model'].nc\n",
"\n",
"# reparametrized YOLOR\n",
"for i in range(255):\n",
" model.state_dict()['model.105.m.0.weight'].data[i, :, :, :] *= state_dict['model.105.im.0.implicit'].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.105.m.1.weight'].data[i, :, :, :] *= state_dict['model.105.im.1.implicit'].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.105.m.2.weight'].data[i, :, :, :] *= state_dict['model.105.im.2.implicit'].data[:, i, : :].squeeze()\n",
"model.state_dict()['model.105.m.0.bias'].data += state_dict['model.105.m.0.weight'].mul(state_dict['model.105.ia.0.implicit']).sum(1).squeeze()\n",
"model.state_dict()['model.105.m.1.bias'].data += state_dict['model.105.m.1.weight'].mul(state_dict['model.105.ia.1.implicit']).sum(1).squeeze()\n",
"model.state_dict()['model.105.m.2.bias'].data += state_dict['model.105.m.2.weight'].mul(state_dict['model.105.ia.2.implicit']).sum(1).squeeze()\n",
"model.state_dict()['model.105.m.0.bias'].data *= state_dict['model.105.im.0.implicit'].data.squeeze()\n",
"model.state_dict()['model.105.m.1.bias'].data *= state_dict['model.105.im.1.implicit'].data.squeeze()\n",
"model.state_dict()['model.105.m.2.bias'].data *= state_dict['model.105.im.2.implicit'].data.squeeze()\n",
"\n",
"# model to be saved\n",
"ckpt = {'model': deepcopy(model.module if is_parallel(model) else model).half(),\n",
" 'optimizer': None,\n",
" 'training_results': None,\n",
" 'epoch': -1}\n",
"\n",
"# save reparameterized model\n",
"torch.save(ckpt, 'cfg/deploy/yolov7.pt')\n"
]
},
{
"cell_type": "markdown",
"id": "5b396a53",
"metadata": {},
"source": [
"## YOLOv7x reparameterization"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9d54d17f",
"metadata": {},
"outputs": [],
"source": [
"# import\n",
"from copy import deepcopy\n",
"from models.yolo import Model\n",
"import torch\n",
"from utils.torch_utils import select_device, is_parallel\n",
"\n",
"device = select_device('0', batch_size=1)\n",
"# model trained by cfg/training/*.yaml\n",
"ckpt = torch.load('cfg/training/yolov7x.pt', map_location=device)\n",
"# reparameterized model in cfg/deploy/*.yaml\n",
"model = Model('cfg/deploy/yolov7x.yaml', ch=3, nc=80).to(device)\n",
"\n",
"# copy intersect weights\n",
"state_dict = ckpt['model'].float().state_dict()\n",
"exclude = []\n",
"intersect_state_dict = {k: v for k, v in state_dict.items() if k in model.state_dict() and not any(x in k for x in exclude) and v.shape == model.state_dict()[k].shape}\n",
"model.load_state_dict(intersect_state_dict, strict=False)\n",
"model.names = ckpt['model'].names\n",
"model.nc = ckpt['model'].nc\n",
"\n",
"# reparametrized YOLOR\n",
"for i in range(255):\n",
" model.state_dict()['model.121.m.0.weight'].data[i, :, :, :] *= state_dict['model.121.im.0.implicit'].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.121.m.1.weight'].data[i, :, :, :] *= state_dict['model.121.im.1.implicit'].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.121.m.2.weight'].data[i, :, :, :] *= state_dict['model.121.im.2.implicit'].data[:, i, : :].squeeze()\n",
"model.state_dict()['model.121.m.0.bias'].data += state_dict['model.121.m.0.weight'].mul(state_dict['model.121.ia.0.implicit']).sum(1).squeeze()\n",
"model.state_dict()['model.121.m.1.bias'].data += state_dict['model.121.m.1.weight'].mul(state_dict['model.121.ia.1.implicit']).sum(1).squeeze()\n",
"model.state_dict()['model.121.m.2.bias'].data += state_dict['model.121.m.2.weight'].mul(state_dict['model.121.ia.2.implicit']).sum(1).squeeze()\n",
"model.state_dict()['model.121.m.0.bias'].data *= state_dict['model.121.im.0.implicit'].data.squeeze()\n",
"model.state_dict()['model.121.m.1.bias'].data *= state_dict['model.121.im.1.implicit'].data.squeeze()\n",
"model.state_dict()['model.121.m.2.bias'].data *= state_dict['model.121.im.2.implicit'].data.squeeze()\n",
"\n",
"# model to be saved\n",
"ckpt = {'model': deepcopy(model.module if is_parallel(model) else model).half(),\n",
" 'optimizer': None,\n",
" 'training_results': None,\n",
" 'epoch': -1}\n",
"\n",
"# save reparameterized model\n",
"torch.save(ckpt, 'cfg/deploy/yolov7x.pt')\n"
]
},
{
"cell_type": "markdown",
"id": "11a9108e",
"metadata": {},
"source": [
"## YOLOv7-W6 reparameterization"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d032c629",
"metadata": {},
"outputs": [],
"source": [
"# import\n",
"from copy import deepcopy\n",
"from models.yolo import Model\n",
"import torch\n",
"from utils.torch_utils import select_device, is_parallel\n",
"\n",
"device = select_device('0', batch_size=1)\n",
"# model trained by cfg/training/*.yaml\n",
"ckpt = torch.load('cfg/training/yolov7-w6.pt', map_location=device)\n",
"# reparameterized model in cfg/deploy/*.yaml\n",
"model = Model('cfg/deploy/yolov7-w6.yaml', ch=3, nc=80).to(device)\n",
"\n",
"# copy intersect weights\n",
"state_dict = ckpt['model'].float().state_dict()\n",
"exclude = []\n",
"intersect_state_dict = {k: v for k, v in state_dict.items() if k in model.state_dict() and not any(x in k for x in exclude) and v.shape == model.state_dict()[k].shape}\n",
"model.load_state_dict(intersect_state_dict, strict=False)\n",
"model.names = ckpt['model'].names\n",
"model.nc = ckpt['model'].nc\n",
"\n",
"idx = 118\n",
"idx2 = 122\n",
"\n",
"# copy weights of lead head\n",
"model.state_dict()['model.{}.m.0.weight'.format(idx)].data -= model.state_dict()['model.{}.m.0.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.1.weight'.format(idx)].data -= model.state_dict()['model.{}.m.1.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.2.weight'.format(idx)].data -= model.state_dict()['model.{}.m.2.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.3.weight'.format(idx)].data -= model.state_dict()['model.{}.m.3.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.0.weight'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.1.weight'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.2.weight'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.3.weight'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data -= model.state_dict()['model.{}.m.0.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data -= model.state_dict()['model.{}.m.1.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data -= model.state_dict()['model.{}.m.2.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data -= model.state_dict()['model.{}.m.3.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.bias'.format(idx2)].data\n",
"\n",
"# reparametrized YOLOR\n",
"for i in range(255):\n",
" model.state_dict()['model.{}.m.0.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.0.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.1.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.1.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.2.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.2.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.3.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.3.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].mul(state_dict['model.{}.ia.0.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].mul(state_dict['model.{}.ia.1.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].mul(state_dict['model.{}.ia.2.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].mul(state_dict['model.{}.ia.3.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data *= state_dict['model.{}.im.0.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data *= state_dict['model.{}.im.1.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data *= state_dict['model.{}.im.2.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data *= state_dict['model.{}.im.3.implicit'.format(idx2)].data.squeeze()\n",
"\n",
"# model to be saved\n",
"ckpt = {'model': deepcopy(model.module if is_parallel(model) else model).half(),\n",
" 'optimizer': None,\n",
" 'training_results': None,\n",
" 'epoch': -1}\n",
"\n",
"# save reparameterized model\n",
"torch.save(ckpt, 'cfg/deploy/yolov7-w6.pt')\n"
]
},
{
"cell_type": "markdown",
"id": "5f093d43",
"metadata": {},
"source": [
"## YOLOv7-E6 reparameterization"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "aa2b2142",
"metadata": {},
"outputs": [],
"source": [
"# import\n",
"from copy import deepcopy\n",
"from models.yolo import Model\n",
"import torch\n",
"from utils.torch_utils import select_device, is_parallel\n",
"\n",
"device = select_device('0', batch_size=1)\n",
"# model trained by cfg/training/*.yaml\n",
"ckpt = torch.load('cfg/training/yolov7-e6.pt', map_location=device)\n",
"# reparameterized model in cfg/deploy/*.yaml\n",
"model = Model('cfg/deploy/yolov7-e6.yaml', ch=3, nc=80).to(device)\n",
"\n",
"# copy intersect weights\n",
"state_dict = ckpt['model'].float().state_dict()\n",
"exclude = []\n",
"intersect_state_dict = {k: v for k, v in state_dict.items() if k in model.state_dict() and not any(x in k for x in exclude) and v.shape == model.state_dict()[k].shape}\n",
"model.load_state_dict(intersect_state_dict, strict=False)\n",
"model.names = ckpt['model'].names\n",
"model.nc = ckpt['model'].nc\n",
"\n",
"idx = 140\n",
"idx2 = 144\n",
"\n",
"# copy weights of lead head\n",
"model.state_dict()['model.{}.m.0.weight'.format(idx)].data -= model.state_dict()['model.{}.m.0.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.1.weight'.format(idx)].data -= model.state_dict()['model.{}.m.1.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.2.weight'.format(idx)].data -= model.state_dict()['model.{}.m.2.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.3.weight'.format(idx)].data -= model.state_dict()['model.{}.m.3.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.0.weight'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.1.weight'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.2.weight'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.3.weight'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data -= model.state_dict()['model.{}.m.0.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data -= model.state_dict()['model.{}.m.1.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data -= model.state_dict()['model.{}.m.2.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data -= model.state_dict()['model.{}.m.3.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.bias'.format(idx2)].data\n",
"\n",
"# reparametrized YOLOR\n",
"for i in range(255):\n",
" model.state_dict()['model.{}.m.0.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.0.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.1.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.1.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.2.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.2.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.3.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.3.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].mul(state_dict['model.{}.ia.0.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].mul(state_dict['model.{}.ia.1.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].mul(state_dict['model.{}.ia.2.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].mul(state_dict['model.{}.ia.3.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data *= state_dict['model.{}.im.0.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data *= state_dict['model.{}.im.1.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data *= state_dict['model.{}.im.2.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data *= state_dict['model.{}.im.3.implicit'.format(idx2)].data.squeeze()\n",
"\n",
"# model to be saved\n",
"ckpt = {'model': deepcopy(model.module if is_parallel(model) else model).half(),\n",
" 'optimizer': None,\n",
" 'training_results': None,\n",
" 'epoch': -1}\n",
"\n",
"# save reparameterized model\n",
"torch.save(ckpt, 'cfg/deploy/yolov7-e6.pt')\n"
]
},
{
"cell_type": "markdown",
"id": "a3bccf89",
"metadata": {},
"source": [
"## YOLOv7-D6 reparameterization"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e5216b70",
"metadata": {},
"outputs": [],
"source": [
"# import\n",
"from copy import deepcopy\n",
"from models.yolo import Model\n",
"import torch\n",
"from utils.torch_utils import select_device, is_parallel\n",
"\n",
"device = select_device('0', batch_size=1)\n",
"# model trained by cfg/training/*.yaml\n",
"ckpt = torch.load('cfg/training/yolov7-d6.pt', map_location=device)\n",
"# reparameterized model in cfg/deploy/*.yaml\n",
"model = Model('cfg/deploy/yolov7-d6.yaml', ch=3, nc=80).to(device)\n",
"\n",
"# copy intersect weights\n",
"state_dict = ckpt['model'].float().state_dict()\n",
"exclude = []\n",
"intersect_state_dict = {k: v for k, v in state_dict.items() if k in model.state_dict() and not any(x in k for x in exclude) and v.shape == model.state_dict()[k].shape}\n",
"model.load_state_dict(intersect_state_dict, strict=False)\n",
"model.names = ckpt['model'].names\n",
"model.nc = ckpt['model'].nc\n",
"\n",
"idx = 162\n",
"idx2 = 166\n",
"\n",
"# copy weights of lead head\n",
"model.state_dict()['model.{}.m.0.weight'.format(idx)].data -= model.state_dict()['model.{}.m.0.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.1.weight'.format(idx)].data -= model.state_dict()['model.{}.m.1.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.2.weight'.format(idx)].data -= model.state_dict()['model.{}.m.2.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.3.weight'.format(idx)].data -= model.state_dict()['model.{}.m.3.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.0.weight'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.1.weight'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.2.weight'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.3.weight'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data -= model.state_dict()['model.{}.m.0.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data -= model.state_dict()['model.{}.m.1.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data -= model.state_dict()['model.{}.m.2.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data -= model.state_dict()['model.{}.m.3.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.bias'.format(idx2)].data\n",
"\n",
"# reparametrized YOLOR\n",
"for i in range(255):\n",
" model.state_dict()['model.{}.m.0.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.0.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.1.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.1.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.2.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.2.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.3.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.3.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].mul(state_dict['model.{}.ia.0.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].mul(state_dict['model.{}.ia.1.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].mul(state_dict['model.{}.ia.2.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].mul(state_dict['model.{}.ia.3.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data *= state_dict['model.{}.im.0.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data *= state_dict['model.{}.im.1.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data *= state_dict['model.{}.im.2.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data *= state_dict['model.{}.im.3.implicit'.format(idx2)].data.squeeze()\n",
"\n",
"# model to be saved\n",
"ckpt = {'model': deepcopy(model.module if is_parallel(model) else model).half(),\n",
" 'optimizer': None,\n",
" 'training_results': None,\n",
" 'epoch': -1}\n",
"\n",
"# save reparameterized model\n",
"torch.save(ckpt, 'cfg/deploy/yolov7-d6.pt')\n"
]
},
{
"cell_type": "markdown",
"id": "334c273b",
"metadata": {},
"source": [
"## YOLOv7-E6E reparameterization"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "635fd8d2",
"metadata": {},
"outputs": [],
"source": [
"# import\n",
"from copy import deepcopy\n",
"from models.yolo import Model\n",
"import torch\n",
"from utils.torch_utils import select_device, is_parallel\n",
"\n",
"device = select_device('0', batch_size=1)\n",
"# model trained by cfg/training/*.yaml\n",
"ckpt = torch.load('cfg/training/yolov7-e6e.pt', map_location=device)\n",
"# reparameterized model in cfg/deploy/*.yaml\n",
"model = Model('cfg/deploy/yolov7-e6e.yaml', ch=3, nc=80).to(device)\n",
"\n",
"# copy intersect weights\n",
"state_dict = ckpt['model'].float().state_dict()\n",
"exclude = []\n",
"intersect_state_dict = {k: v for k, v in state_dict.items() if k in model.state_dict() and not any(x in k for x in exclude) and v.shape == model.state_dict()[k].shape}\n",
"model.load_state_dict(intersect_state_dict, strict=False)\n",
"model.names = ckpt['model'].names\n",
"model.nc = ckpt['model'].nc\n",
"\n",
"idx = 261\n",
"idx2 = 265\n",
"\n",
"# copy weights of lead head\n",
"model.state_dict()['model.{}.m.0.weight'.format(idx)].data -= model.state_dict()['model.{}.m.0.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.1.weight'.format(idx)].data -= model.state_dict()['model.{}.m.1.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.2.weight'.format(idx)].data -= model.state_dict()['model.{}.m.2.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.3.weight'.format(idx)].data -= model.state_dict()['model.{}.m.3.weight'.format(idx)].data\n",
"model.state_dict()['model.{}.m.0.weight'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.1.weight'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.2.weight'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.3.weight'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data -= model.state_dict()['model.{}.m.0.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data -= model.state_dict()['model.{}.m.1.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data -= model.state_dict()['model.{}.m.2.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data -= model.state_dict()['model.{}.m.3.bias'.format(idx)].data\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.bias'.format(idx2)].data\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.bias'.format(idx2)].data\n",
"\n",
"# reparametrized YOLOR\n",
"for i in range(255):\n",
" model.state_dict()['model.{}.m.0.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.0.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.1.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.1.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.2.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.2.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
" model.state_dict()['model.{}.m.3.weight'.format(idx)].data[i, :, :, :] *= state_dict['model.{}.im.3.implicit'.format(idx2)].data[:, i, : :].squeeze()\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data += state_dict['model.{}.m.0.weight'.format(idx2)].mul(state_dict['model.{}.ia.0.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data += state_dict['model.{}.m.1.weight'.format(idx2)].mul(state_dict['model.{}.ia.1.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data += state_dict['model.{}.m.2.weight'.format(idx2)].mul(state_dict['model.{}.ia.2.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data += state_dict['model.{}.m.3.weight'.format(idx2)].mul(state_dict['model.{}.ia.3.implicit'.format(idx2)]).sum(1).squeeze()\n",
"model.state_dict()['model.{}.m.0.bias'.format(idx)].data *= state_dict['model.{}.im.0.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.1.bias'.format(idx)].data *= state_dict['model.{}.im.1.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.2.bias'.format(idx)].data *= state_dict['model.{}.im.2.implicit'.format(idx2)].data.squeeze()\n",
"model.state_dict()['model.{}.m.3.bias'.format(idx)].data *= state_dict['model.{}.im.3.implicit'.format(idx2)].data.squeeze()\n",
"\n",
"# model to be saved\n",
"ckpt = {'model': deepcopy(model.module if is_parallel(model) else model).half(),\n",
" 'optimizer': None,\n",
" 'training_results': None,\n",
" 'epoch': -1}\n",
"\n",
"# save reparameterized model\n",
"torch.save(ckpt, 'cfg/deploy/yolov7-e6e.pt')\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "63a62625",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}