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

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

CodePudding user response：

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, ...])

CodePudding user response：

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.