CODEWITHSUNDEEP
pythonnumpypytorchstreamlittorchvision
afrologicinsect
afrologicinsect

Reputation: 171

ValueError: not enough values to unpack (expected 4, got 3) on torchvision

Trying to classify images using custom Pytorch model on Streamlit.

Now before prediction, the served image is first converted with torchvision, however, the following error pops up when I try to predict:

ValueError: not enough values to unpack (expected 4, got 3)

Traceback: File

"c:\users\pc\anaconda3\envs\tf-gpu2.5\lib\site-packages\streamlit\scriptrunner\script_runner.py", line 557, in _run_script exec(code, module.dict) File "app_v2.py", line 48, in predict(image, model) File "app_v2.py", line 25, in predict pred = model(x) File "c:\users\pc\anaconda3\envs\tf-gpu2.5\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl result = self.forward(*input, **kwargs) File "c:\users\pc\anaconda3\envs\tf-gpu2.5\lib\site-packages\torch\autograd\grad_mode.py", line 27, in decorate_context return func(*args, **kwargs) File "C:\Users\PC/.cache\torch\hub\ultralytics_yolov5_master\models\common.py", line 573, in forward return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference File "c:\users\pc\anaconda3\envs\tf-gpu2.5\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl result = self.forward(*input, **kwargs) File "C:\Users\PC/.cache\torch\hub\ultralytics_yolov5_master\models\common.py", line 443, in forward b, ch, h, w = im.shape # batch, channel, height, width

code basically looks like:

## Serve Image
uploaded_file = st.file_uploader('Upload image...', type=['jpeg', 'jpg', 'png'])

if uploaded_file is not None:
    ## convert
    image = Image.open(uploaded_file).convert('L')
    predict(image, model)

def predict(image, model):
    ## convert
    x = torchvision.transforms.ToTensor()(image)
    pred = model(x)
    pred = pred.detach().numpy()

Would appreciate the assistance. Thanks

Upvotes: 0

Views: 6065

Answers (3)

Wok
Wok

Reputation: 5313

You are probably using an RGB image (dimension 3), which is missing the dimension for the batch size, hence an expected dimension 4. To prepend this dimension to your image, use the following.

With NumPy:

import numpy as np

image = np.expand_dims(image, axis=0)

With PyTorch:

import torch

image = torch.unsqueeze(image, dim=0)

Alternatively:

image = image[None, :]

Or, solely because you preprend this dimension, you can omit the colon ::

image = image[None]

Upvotes: 0

afrologicinsect
afrologicinsect

Reputation: 171

So here's how I solved this problem, for anyone else, this was ultimately an Image Classification challenge using Pytorch's custom model on streamlit. So rigged the whole thing.

The following works.

import streamlit as st
from PIL import Image
import numpy as np
import torch
import io

## LOAD MODEL
run_model_path = './model_file.pt'
model = torch.hub.load('ultralytics/yolov5', 'custom', path=run_model_path)
model.eval()

## Collect
img_file_buffer = st.sidebar.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
if img_file_buffer is not None:
    image = np.array(Image.open(img_file_buffer))
    ## Predict Image
    result = model(image)
    ## Convert Image format
    output = io.BytesIO()
    out_image = np.squeeze(result.render()) ## to numpy array
    o_img = Image.fromarray(out_image) ## to PIL
    o_img.save(output, format='JPEG')
    result_img = output.getvalue() ## Back to JPEG
    ## Show Result
    st.image(result_img)

So train (hard stuff) save in the same directory you run your app.

Phew.

Upvotes: 0

Shai
Shai

Reputation: 114796

Your input image x has only three dimensions: channels (RGB), height and width. PyTorch expects input images to have four dimensions: batch, channels, height and width.

You need to add a singleton "batch" dimension:

pred = model(x[None, ...])

Upvotes: 0

Related Questions