MOT数据集制作,yolov5生成MOT数据集中det.txt文件

import torch
import os
from PIL import Image
import pandas as pd
# Model
#model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5l, yolov5x, etc. 'ultralytics/yolov5', 'custom',
model = torch.hub.load( 'ultralytics/yolov5','custom','/home/se8502/Desktop/gaogao/yolov5-master/runs/train/exp4/weights/best.pt')  # custom trained model

# Images
#im = './MOTFish/images/train/Black-porgy-03-01/img1/1.jpg'  # or file, Path, URL, PIL, OpenCV, numpy, list



directory = './MOTFish/images/train/Black-porgy-06/img1/'

frames = [] # batch of images
for filename in os.listdir(directory):
    frame = Image.open('./MOTFish/images/train/Black-porgy-06/img1/'+filename)
    frames.append(frame)
# Inference
#results = model(im)
results = model(frames)
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
results.save()

# print(results.xyxy[0])  # im predictions (tensor)
mypath = "./test.txt"
txt_outfile = open(mypath, "a")
cn = 0

for i in range(len(results.pandas().xyxy)):
    df = results.pandas().xyxy[i].to_json()
    df = eval(df)
    if df['class']:
        #write('frame class xmin ymin xmax ymax confidence')
        write = str(i+1) + ' ' + str(df['class']['0']) + ' ' + str(df['xmin']['0']) + ' ' + str(df['ymin']['0']) + ' ' + str(df['xmax']['0']) + ' ' + str(df['ymax']['0']) + ' ' + str(df['confidence']['0']) + '\n'
        txt_outfile.write(write)


版权声明:本文为qwazp3526cn原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接和本声明。