Testing on a Single Image

[Prabhav55221]

Hi @qubvel!

I have been using SMP for image segmentation on MRI images of the human brain. I used the UnetPlusPlus model to train and that worked perfectly. Further, I was also able to test the same on the validation dataset. But what I want to do, is to test on a single image and to further visualize the image. So basically, I need the prediction of the model to visualize the segmentation. Is there any way to do this in the SMP docs? For further context, I am currently testing like this:

_To conclude,

1. I need to access the predictions of the MODEL.
2. I need to be able to run the model on single images._

Any help would be appreciated!

[Nachimak28]

Hi @Prabhav55221 You could simply use the way the CamVid notebook example does using the model.predict() function.

Load your model and put it in eval mode:

model = torch.load('your_model_path.pth')
model.eval()

Import torchvision transforms if not done already

from PIL import Image
from torchvision import transforms

Open the single image you want to get predictions from using PIL (pillow)

img = Image.open('path/to/your/image.jpg')

Create the transform pipeline as per you might have done for your model validation transforms Mine was as follows

tfms = transforms.Compose([
        transforms.Resize(img_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

Put the image to device

# unsqueeze provides the batch dimension
img_tensor = tfms(img).to('cuda').unsqueeze(0)

Run prediction over single image

output = model.predict(img_tensor)
# to binarize the output since I had only 1 class
output = (output.squeeze().cpu().numpy().round())
# the output pixel intensities will be 0 for background and 1 for foreground class predicted

# you can get it in binary mask format as such incase you need it
pred_mask = np.asarray(output, dtype=np.uint8)*255

plt.imshow(pred_mask)
plt.show()

Hopefully this explanation is relevant to your problem statement

[thalia27]

Hi @Prabhav55221 You could simply use the way the CamVid notebook example does using the model.predict() function.

Load your model and put it in eval mode:

model = torch.load('your_model_path.pth')
model.eval()

Import torchvision transforms if not done already

from PIL import Image
from torchvision import transforms

Open the single image you want to get predictions from using PIL (pillow)

img = Image.open('path/to/your/image.jpg')

Create the transform pipeline as per you might have done for your model validation transforms Mine was as follows

tfms = transforms.Compose([
        transforms.Resize(img_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

Put the image to device

# unsqueeze provides the batch dimension
img_tensor = tfms(img).to('cuda').unsqueeze(0)

Run prediction over single image

output = model.predict(img_tensor)
# to binarize the output since I had only 1 class
output = (output.squeeze().cpu().numpy().round())
# the output pixel intensities will be 0 for background and 1 for foreground class predicted

# you can get it in binary mask format as such incase you need it
pred_mask = np.asarray(output, dtype=np.uint8)*255

plt.imshow(pred_mask)
plt.show()

Hopefully this explanation is relevant to your problem statement

@Nachimak28 Thanks Man!!!

[github-actions[bot]]

This issue is stale because it has been open 60 days with no activity. Remove stale label or comment or this will be closed in 7 days.

[github-actions[bot]]

This issue was closed because it has been stalled for 7 days with no activity.