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)
