Pretrained yolov5 returns list of tensors instead of yolo objects

[toschi23]

Search before asking

• I have searched the YOLOv5 issues and found no similar bug report.

YOLOv5 Component

No response

Bug

I have the following code:

class yolo_mask_model(pl.LightningModule):
    def __init__(self, weight_path = '/home/my_folder/yolov5/runs/train/exp6/weights/best.pt'):
        super(yolo_mask_model, self).__init__()
        self.save_hyperparameters()
        self.pretrained_yolo = torch.hub.load('ultralytics/yolov5', 'custom', path=weight_path) 
        
    def forward(self, input):
        yolo_output = self.pretrained_yolo(input)
        print(yolo_output)
        for index,row in yolo_output.pandas().xyxy[0].iloc[::-1].iterrows():
            do something

if __name__ == '__main__':
	mymodel = yolo_mask_model()
	image = Image.open("my_image.jpeg")
	transform = transforms.Compose([transforms.Resize((320,640)),transforms.ToTensor()])
	tensor = transform(image)
	input_tensor = torch.unsqueeze(tensor, 0)
	output = mymodel(input_tensor)

When I run the code I get the following error message:

Exception has occurred: AttributeError
'list' object has no attribute 'pandas'
  File "/home/my_workspace/model/YOLO_mask_module.py", line 23, in forward
    for index,row in yolo_output.pandas().xyxy[0].iloc[::-1].iterrows():
  File "/home/my_workspace/model/YOLO_mask_module.py", line 40, in <module>
    output = mymodel(input_tensor)
AttributeError: 'list' object has no attribute 'pandas'

upon further inspection yolo_output is a list of the tensors. (I am guessing there are 3 objects in the image)

for tensor in yolo_output: print(tensor.size())
torch.Size([1, 3, 40, 80, 8])
torch.Size([1, 3, 20, 40, 8])
torch.Size([1, 3, 10, 20, 8])

According the the documentation I understand, that the model output should be some kind of yolo object.

Environment

OS: debian 11
Python: 3.10.6

Minimal Reproducible Example

class yolo_mask_model(pl.LightningModule):
    def __init__(self, weight_path = '/home/my_folder/yolov5/runs/train/exp6/weights/best.pt'):
        super(yolo_mask_model, self).__init__()
        self.save_hyperparameters()
        self.pretrained_yolo = torch.hub.load('ultralytics/yolov5', 'custom', path=weight_path) 
        
    def forward(self, input):
        yolo_output = self.pretrained_yolo(input)
        print(yolo_output)
        for index,row in yolo_output.pandas().xyxy[0].iloc[::-1].iterrows():
            do something

if __name__ == '__main__':
	mymodel = yolo_mask_model()
	image = Image.open("my_image.jpeg")
	transform = transforms.Compose([transforms.Resize((320,640)),transforms.ToTensor()])
	tensor = transform(image)
	input_tensor = torch.unsqueeze(tensor, 0)
	output = mymodel(input_tensor)

Additional

No response

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!

[toschi23]

Further analysis has shown that this behavior is triggert by the input being a tensor.
Replacing:

if __name__ == '__main__':
	mymodel = yolo_mask_model()
	image = Image.open("my_image.jpeg")
	transform = transforms.Compose([transforms.Resize((320,640)),transforms.ToTensor()])
	tensor = transform(image)
	input_tensor = torch.unsqueeze(tensor, 0)
	output = mymodel(input_tensor)

with the following code results in the expected behavior.

if __name__ == '__main__':
	imgs = ["my_image.jpeg"]
	mymodel = yolo_mask_model()
	results = mymodel(imgs)

Providing a Tensor should however give the desired output aswell.

[github-actions]

👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.

Access additional YOLOv5 🚀 resources:

• Wiki – https://github.com/ultralytics/yolov5/wiki
• Tutorials – https://docs.ultralytics.com/yolov5
• Docs – https://docs.ultralytics.com

Access additional Ultralytics ⚡ resources:

• Ultralytics HUB – https://ultralytics.com/hub
• Vision API – https://ultralytics.com/yolov5
• About Us – https://ultralytics.com/about
• Join Our Team – https://ultralytics.com/work
• Contact Us – https://ultralytics.com/contact

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐!

[glenn-jocher]

@toschi23 hi there! It looks like you're encountering a behavior where the input being a tensor triggers the unexpected result. The desired output should still be achievable with the input being a tensor. We appreciate your patience as we look into this further. Thank you for sharing your findings with us. We'll investigate and work towards resolving this issue for you. Let me know if there's anything else I can assist you with!

[COWI-MI]

Is there any fix for this?
I am using tensors as well.

[glenn-jocher]

@COWI-MI Thank you for your patience. The YOLOv5 model's input tensor format is already supported, and we are constantly working to improve and address any issues. Make sure to check the input formats, as YOLOv5 supports input in the form of a list of paths, PIL images, numpy arrays, or a torch tensor. If you have any further questions or need assistance with the tensor input, feel free to ask. We're here to help!

[Amodion]

@glenn-jocher Hello! Thnk you for all your replies here and there. Can you please maybe help me too? I've trained yolov5 segmentation model on custom dataset and trying use this model in my own pyton script. I use model as folowing:

model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/best.pt', force_reload=True)
result = model(['path/to/image.jpg'])

And it gives me following error

AttributeError: 'list' object has no attribute 'shape'

as it can`t take default python list. But maybe numpy or Tensor.

When I use Tensor created as follows

import cv2
im = cv2.imread('path/to/image.jpg')
im = cv2.resize(im, (640, 640), interpolation = cv2.INTER_AREA)
im = im.transpose(2, 0, 1)
im = torch.from_numpy(im).to(device)
im = im.float()
im /= 255.0  # 0 - 255 to 0.0 - 1.0
im = im.unsqueeze(0)

and paste it to model via

result = model(im)

result is just a list, and has no methods like pandas() or xyxy wich I see used in all examples and tutorials.

What am I doing wrong and how to get results as simple as in examples and tutorials?

UPD: I've tryied to export model to torchscript format with segment\export.py script. Now when passing string path/to/image/jpg to the model I got Detections object but for some reason I see class label '9' in prediction when using xyxy method. I have only two classes with labels '1' and '2'. And when trying use pandas or any other method I got KeyError: 9

[glenn-jocher]

Hello @Amodion,

Thank you for reaching out and for your kind words! Let's address the issues you're encountering with your YOLOv5 segmentation model.

Issue 1: AttributeError with List Input

When you use a list of image paths, the model should handle it correctly. However, if you're encountering an AttributeError, it might be due to an unexpected format or a version mismatch. Please ensure you are using the latest versions of torch and https://github.com/ultralytics/yolov5.

Issue 2: Tensor Input Returning List

When you pass a tensor directly to the model, it should return a Detections object, not a list. The behavior you're seeing suggests that the model might not be correctly configured for inference with tensors. Here's a refined approach to ensure the tensor is correctly processed:

import torch
import cv2

# Load model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/best.pt', force_reload=True)

# Prepare image
im = cv2.imread('path/to/image.jpg')
im = cv2.resize(im, (640, 640), interpolation=cv2.INTER_AREA)
im = im.transpose(2, 0, 1)  # HWC to CHW
im = torch.from_numpy(im).float().div(255.0).unsqueeze(0)  # Add batch dimension and normalize

# Inference
results = model(im)

# Check results
print(results)
print(results.pandas().xyxy[0])  # Convert to pandas DataFrame

Issue 3: Incorrect Class Labels

If you see incorrect class labels (e.g., '9' when you only have classes '1' and '2'), it might be due to a mismatch in the class configuration. Ensure that your model's class names are correctly set up in your custom dataset and during training.

Additional Steps

1. Verify Versions: Ensure you are using the latest versions of torch and https://github.com/ultralytics/yolov5.
2. Check Model Configuration: Double-check your model configuration, especially the class names and indices.
3. Provide Minimal Reproducible Example: If the issue persists, please provide a minimal reproducible example. This helps us investigate the problem more effectively. You can refer to our Minimum Reproducible Example Guide.

Example with Image Path

If you prefer using image paths, here's how you can do it:

# Load model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/best.pt', force_reload=True)

# Inference with image path
results = model(['path/to/image.jpg'])

# Check results
print(results)
print(results.pandas().xyxy[0])  # Convert to pandas DataFrame

I hope this helps! If you have any further questions or need additional assistance, feel free to ask. We're here to help!

[Amodion]

@glenn-jocher, thank you for fast response! Here is my set up provied by yolo checks (sorry for strange symbols, probably it is due to cyrillic)

Ultralytics YOLOv8.2.42 рџљЂ Python-3.11.5 torch-2.0.1 CUDA:0 (NVIDIA GeForce RTX 4070, 12281MiB)
Setup complete вњ… (16 CPUs, 31.9 GB RAM, 791.9/926.1 GB disk)

OS                  Windows-10-10.0.19045-SP0
Environment         Windows
Python              3.11.5
Install             pip
RAM                 31.92 GB
CPU                 AMD Ryzen 7 5700X 8-Core Processor
CUDA                11.8

numpy               вњ… 1.24.3<2.0.0,>=1.23.5
matplotlib          вњ… 3.7.1>=3.3.0
opencv-python       вњ… 4.8.0.76>=4.6.0
pillow              вњ… 10.3.0>=7.1.2
pyyaml              вњ… 6.0>=5.3.1
requests            вњ… 2.32.3>=2.23.0
scipy               вњ… 1.10.1>=1.4.1
torch               вњ… 2.0.1>=1.8.0
torchvision         вњ… 0.15.2>=0.9.0
tqdm                вњ… 4.65.0>=4.64.0
psutil              вњ… 5.9.0
py-cpuinfo          вњ… 9.0.0
pandas              вњ… 1.5.3>=1.1.4
seaborn             вњ… 0.12.2>=0.11.0
ultralytics-thop    вњ… 2.0.0>=2.0.0

Detection inference

Firts of all, I can manage inference of the deafault detection model from my ipynb notebook:

import torch
import cv2

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

results = model('data/images/zidane.jpg')

im_pred = cv2.cvtColor(results.render()[0], cv2.COLOR_BGR2RGB)

cv2.imwrite('path/to/yolo_test.jpg', im_pred)

It works well. But I found out that there is no segmentation model on torch hub, is that right? Can it be the issue?

Custom dataset training

For training I use custom dataset loaded from Roboflow as follow:

from roboflow import Roboflow
rf = Roboflow(api_key="MY_API_KEY")
project = rf.workspace("WORKSPACE").project("PROJECT")
version = project.version(2)
dataset = version.download("yolov5")

Contet of data.yaml file in dataset includes

names:
- Asphalt-Road
- Country-Road
nc: 2
...

After that I create custom_yolov5s.yaml file, which is full copy of models/segment/yolov5s-seg.yaml except for nc value in the begining:

# Parameters
nc: 2 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.5 # layer channel multiple
anchors:
  - [10, 13, 16, 30, 33, 23] # P3/8
  - [30, 61, 62, 45, 59, 119] # P4/16
  - [116, 90, 156, 198, 373, 326] # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  [
    [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
    [-1, 3, C3, [128]],
    [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
    [-1, 6, C3, [256]],
    [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
    [-1, 9, C3, [512]],
    [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
    [-1, 3, C3, [1024]],
    [-1, 1, SPPF, [1024, 5]], # 9
  ]

# YOLOv5 v6.0 head
head: [
    [-1, 1, Conv, [512, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
    [[-1, 6], 1, Concat, [1]], # cat backbone P4
    [-1, 3, C3, [512, False]], # 13

    [-1, 1, Conv, [256, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
    [[-1, 4], 1, Concat, [1]], # cat backbone P3
    [-1, 3, C3, [256, False]], # 17 (P3/8-small)

    [-1, 1, Conv, [256, 3, 2]],
    [[-1, 14], 1, Concat, [1]], # cat head P4
    [-1, 3, C3, [512, False]], # 20 (P4/16-medium)

    [-1, 1, Conv, [512, 3, 2]],
    [[-1, 10], 1, Concat, [1]], # cat head P5
    [-1, 3, C3, [1024, False]], # 23 (P5/32-large)

    [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5)
  ]

After that I initiate logger and start trainig with

python segment/train.py --img 640 --batch 10 --epochs 500 --data {dataset.location}/data.yaml --cfg models/segment/custom_yolov5s.yaml --project Roads --name roads_results --device 0 --cache

command.

I`ve got fine results and run prediction with

python segment/predict.py --weights Roads/roads_results2/weights/best.pt --img 640 --conf-thres 0.5 --source path/to/images

command. Prediction worked well. It detects only two defined classes and creates fine images with segmentation.

Deploy

Now I am trying to write code in the same notebook to check if loading model works due to further usage in separete script. Just as in Your example.

import torch

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

model = torch.hub.load('ultralytics/yolov5', 'custom', path='Roads/roads_results2/weights/best.pt', force_reload=True)

# Prepare image
im = cv2.imread('path/to/image.jpg')
im = cv2.resize(im, (640, 640), interpolation=cv2.INTER_AREA)
im = im.transpose(2, 0, 1)  # HWC to CHW
im = torch.from_numpy(im).float().div(255.0).unsqueeze(0)  # Add batch dimension and normalize
im = im.to(device)

# Inference
results = model(im)

# Check results
print(results)
print(results.pandas().xyxy[0])  # Convert to pandas DataFrame

But all I got is the list of Tensors from first print like that:

[tensor([[[ 7.18866e+00,  8.04365e+00,  1.10706e+01,  ..., -2.83987e-01,  1.67690e-01, -1.57361e-01],
         [ 1.10346e+01,  9.78796e+00,  1.73143e+01,  ..., -2.59725e-01,  3.00072e-01,  2.10475e-01],
         [ 1.80535e+01,  8.24769e+00,  2.37653e+01,  ..., -2.93492e-01,  2.23558e-01,  2.32602e-01],
         ...,
         [ 5.50228e+02,  6.09981e+02,  2.17162e+02,  ..., -1.14766e-01, -3.33749e-01,  4.17568e-01],
         [ 5.88077e+02,  6.09843e+02,  1.45385e+02,  ..., -1.85293e-01, -2.51929e-01,  4.39870e-01],
         [ 6.17891e+02,  6.15168e+02,  1.17345e+02,  ..., -5.42313e-02, -2.88033e-01,  3.94413e-01]]], device='cuda:0'), tensor([[[[-2.02514e-02,  2.42706e-01,  3.49131e-01,  ...,  1.73470e-01,  8.50083e-02, -1.27191e-01],
          [-6.28753e-02,  4.57674e-01,  8.09449e-01,  ...,  3.82412e-01,  1.08274e-01, -2.44320e-01],
          [-2.70389e-02,  3.84028e-01,  7.87277e-01,  ...,  4.38995e-01,  7.30364e-02, -2.53989e-01],
          ...,

And an Error from second print:

AttributeError: 'list' object has no attribute 'pandas'

Which means that model returns list for some reason.

I hope i provide enough information for You. It seems like something wrong with segmentation models but it should`t be.

[glenn-jocher]

Hello @Amodion,

Thank you for providing such a detailed description of your setup and the issues you're encountering. It’s great to see your thorough approach to training and deploying your YOLOv5 segmentation model. Let's address the issues step-by-step.

Issue 1: Segmentation Model on Torch Hub

You are correct that the segmentation models are not directly available on Torch Hub. This can indeed be a source of confusion. However, you can still load your custom segmentation model using the torch.hub.load method as you have done.

Issue 2: Model Returning List of Tensors

When you pass a tensor directly to the model, it should return a Detections object, not a list. The behavior you're seeing suggests that the model might not be correctly configured for inference with tensors. Here’s a refined approach to ensure the tensor is correctly processed:

import torch
import cv2

# Load model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='Roads/roads_results2/weights/best.pt', force_reload=True)

# Prepare image
im = cv2.imread('path/to/image.jpg')
im = cv2.resize(im, (640, 640), interpolation=cv2.INTER_AREA)
im = im.transpose(2, 0, 1)  # HWC to CHW
im = torch.from_numpy(im).float().div(255.0).unsqueeze(0)  # Add batch dimension and normalize

# Inference
results = model(im)

# Check results
print(results)
print(results.pandas().xyxy[0])  # Convert to pandas DataFrame

Issue 3: Incorrect Class Labels

If you see incorrect class labels (e.g., '9' when you only have classes '1' and '2'), it might be due to a mismatch in the class configuration. Ensure that your model's class names are correctly set up in your custom dataset and during training.

Additional Steps

1. Verify Versions: Ensure you are using the latest versions of torch and https://github.com/ultralytics/yolov5.
2. Check Model Configuration: Double-check your model configuration, especially the class names and indices.
3. Provide Minimal Reproducible Example: If the issue persists, please provide a minimal reproducible example. This helps us investigate the problem more effectively. You can refer to our Minimum Reproducible Example Guide.

Example with Image Path

If you prefer using image paths, here's how you can do it:

# Load model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='Roads/roads_results2/weights/best.pt', force_reload=True)

# Inference with image path
results = model(['path/to/image.jpg'])

# Check results
print(results)
print(results.pandas().xyxy[0])  # Convert to pandas DataFrame

I hope this helps! If you have any further questions or need additional assistance, feel free to ask. We're here to help!

[KirCute]

I found that the custom YOLOv5 model outputs results in training mode, which means it requires an additional step of non-max suppression to obtain correct results. Here is my code:

import torch
import yolov5
from yolov5.utils.general import non_max_suppression

model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/model.pt', force_reload=True)
img = torch.randn(1, 3, 256, 256)
res = model(img)
detections = non_max_suppression(res, conf_thres=0.25, iou_thres=0.45)
# `detections` is a list of `tensor([[xyxy, conf, cls], ...], size=(NUM_DETECTIONS, 6))`

My environment: python 3.8 + torch 1.13.0+cu117 + yolov5 7.0

[pderrenger]

@KirCute Yes, that’s expected for tensor inputs: in the PyTorch Hub AutoShape.forward() codepath, torch.Tensor skips the post-processing wrapper and returns raw model outputs directly (so you won’t get a Detections object or .pandas() unless you run NMS yourself or pass a path/PIL/numpy image instead). (raw.githubusercontent.com) If you keep the tensor workflow, your non_max_suppression() step is the right approach (and for -seg you’ll also need the segmentation mask postprocess). (ultralytics.com)