yolov5 outputs' sizes are wrong

Question

I am trying to use yolov5n6 pretrained model on local but it seems there are some errors.

I expect the detection size to be torch.Size(3,4) because I sent same image to the model three time in batch. I printed the result and copy-paste next to the each print function.

What is the possible reason to get unexpected outputs?

import torch
import cv2

X = cv2.imread("/path/to/image/normalized_breast.PNG")
X = cv2.resize(X,(1024,1024))
 
print(X.shape) # (1024, 1024, 3)

model = torch.hub.load('ultralytics/yolov5', 'yolov5n6', pretrained=True, classes=1)
checkpoint_ = torch.load('/path/to/pretrained/model/yolov5_breast.pt')['model']
model.load_state_dict(checkpoint_.state_dict())

img = torch.from_numpy(X/255).permute(2,0,1).float().unsqueeze(0)
batch = torch.concatenate([img,img,img]).to("cuda")

print(batch.size()) # torch.Size([3, 3, 1024, 1024])

with torch.no_grad():
    model.eval()
    detections = model(batch)

print(len(detections),         # 2
      detections[0].size(),      # torch.Size([3, 65280, 6])
      detections[1][0].size()) # torch.Size([3, 3, 128, 128, 6])

I also tried following method but it didn't solve the problem too.

import torch
import cv2
from models.experimental import attempt_load
from models.yolo import Model
from utils.general import intersect_dicts,LOGGER

X = cv2.imread("image/path/yoloV5/normalized_breast.PNG")
X = cv2.resize(X,(1024,1024))
print(X.shape)

weights = "/pretrained/model/path/yolov5_breast.pt"
device = "cuda"
cfg = "/yaml/file/path/yolov5/models/hub/yolov5n6.yaml"
nc =1
resume = False

ckpt = torch.load(weights, map_location='cpu')  # load checkpoint to CPU to avoid CUDA memory leak
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=3).to(device)  # create
exclude = ['anchor'] if cfg and not resume else []  # exclude keys
csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # intersect
model.load_state_dict(csd, strict=False)  # load

img = torch.from_numpy(X/255).permute(2,0,1).float().unsqueeze(0)
batch = torch.concatenate([img,img,img]).to("cuda")

print(batch.size())

with torch.no_grad():
    model.eval()
    detections = model(batch)

print(len(detections),
          detections[0].size(),
          detections[1][0].size())

yolov5 outputs' sizes are wrong

0 Answers0