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

[MuhammadBilal848]

Search before asking

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

YOLOv5 Component

Other

Bug

I trained a YOLOv5 custom segmentation model and tried inferencing it without CLI using the following code.

import torch
import numpy as np
import cv2

# model = torch.hub.load('/home/user/Downloads/YOLOv5/yolov5', 'custom', path='best_seg.pt', source='local') # works
model = torch.hub.load('.', 'custom', path='best_seg.pt', source='local') # works
results = model('res.jpg')
print(results)

Error:

YOLOv5 🚀 v7.0-345-g8003649c Python-3.8.0 torch-2.3.1+cu121 CPU

Fusing layers... 
Model summary: 165 layers, 7428089 parameters, 0 gradients, 25.8 GFLOPs
WARNING ⚠️ YOLOv5 SegmentationModel is not yet AutoShape compatible. You will not be able to run inference with this model.
Traceback (most recent call last):
  File "inference_segmentation.py", line 10, in <module>
    results = model('res.jpg')
  File "/home/bilal/Downloads/YOLOv5/v5yolo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/bilal/Downloads/YOLOv5/v5yolo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "./models/common.py", line 659, in forward
    b, ch, h, w = im.shape  # batch, channel, height, width
AttributeError: 'str' object has no attribute 'shape'

Environment

No response

Minimal Reproducible Example

No response

Additional

No response

Are you willing to submit a PR?

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

[UltralyticsAssistant]

👋 Hello @MuhammadBilal848, thank you for reaching out and using YOLOv5 🚀! This is an automated response, and an Ultralytics engineer will assist you soon.

It seems you're encountering an AttributeError because the input to the model should be a NumPy array or a preprocessed image rather than a string. Let's update your code to fix this issue.

Suggested Solution

You should read the image using a library like OpenCV to convert it into an array before passing it to the model. Here's how you can update your code:

import torch
import cv2

# Load the model
model = torch.hub.load('.', 'custom', path='best_seg.pt', source='local')  

# Read image
img = cv2.imread('res.jpg')

# Ensure the image is in the correct format
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# Run inference
results = model(img)

# Print results
print(results)

Important Notes

• Ensure that the image file 'res.jpg' is in the correct directory or provide the full path.
• The cv2.cvtColor(img, cv2.COLOR_BGR2RGB) part ensures that the image is in the RGB format, which is standard for most pre-trained models.

Best Practices for Issue Submission

To help us provide you with the best assistance, consider including:

• A detailed description of your environment, including hardware and installed libraries.
• A minimal, reproducible example if possible.
• Any additional context or screenshots that might be helpful.

Further Assistance

For more details on using YOLOv5, you can check out our Tutorials. If you're interested in our latest updates, consider trying out YOLOv8!

We’re excited to help you resolve this issue, and appreciate your patience while our team gets back to you. 😊

[MuhammadBilal848]

The error is resolved! The following code is used to resolve this.

import torch
import matplotlib.pyplot as plt
import numpy as np
import cv2
from torchvision.transforms import functional as F

# model = torch.hub.load('/home/user/Downloads/YOLOv5/yolov5', 'custom', path='best_seg.pt', source='local') # this also works
model = torch.hub.load('.', 'custom', path='best_seg.pt', source='local')
image = cv2.imread('res.jpg')  # Read image with OpenCV
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB

image = F.to_tensor(image)  # Convert to PyTorch tensor
image = image.unsqueeze(0)  # Add batch dimension

results = model(image)
print(results)

[glenn-jocher]

@MuhammadBilal848 great to hear the issue is resolved! If you have any more questions or need further assistance, feel free to ask.