nv-tlabs / GSCNN

Gated-Shape CNN for Semantic Segmentation (ICCV 2019)
https://nv-tlabs.github.io/GSCNN/
Other
915 stars 200 forks source link

Segment single image using saved model #84

Closed as24-asingla closed 2 years ago

as24-asingla commented 3 years ago

Do we have an inference routine for single images available using the default saved model?

Tiga002 commented 3 years ago

Have you managed to solve it? Since I am trying to deploy the trained model to a mobile robot for testing in the woods nearby.

as24-asingla commented 3 years ago

Unfortunately, not :(

isledge commented 2 years ago

It's not that hard.

# Make sure to import the remaining packages needed by GSCNN.  Here, I'm
# just loading additional packages that are used below.
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

from PIL import Image

from torchvision import transforms

# Insert the argument-parsing code here.  Point to the snapshot that is
# to be used.

# Construct the network and load the weights from the snapshot specified by
# the parsed arguments.  Note that I have modified get_optimizer() to also
# return the network with the loaded weights.
assert_and_infer_cfg(args)
criterion, criterion_val = loss.get_loss(args)
net = network.get_net(args, criterion)
optim, scheduler, net = optimizer.get_optimizer(args, net)

# Release unoccupied cached memory to reduce fragmentation.
torch.cuda.empty_cache()

# Specify the computing device.  Put the network in evaluation mode.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device).eval()

# Load a test image.
test_img = numpy.array(Image.open("./test-image.jpg"), numpy.float32)

# Set up whatever pre-processes and transformations you wish to apply to 
# an image.  Here, I'm just converting an image to a PyTorch Tensor.
convert_tensor = transforms.ToTensor()

# Since we're not updating the network weights, disable the autogradient
# engine to speed up inference.
with torch.no_grad():
    # Convert the input image to a PyTorch Tensor.  Unsqueeze the first 
    # dimension so that the input is of the proper shape for the GSCNN.
    # Store both the segmentation and edge responses of the GSCNN.
    seg, edge = net(torch.unsqueeze(convert_tensor(test_img),0))

# The segmentation response is a tuple, as the GSCNN specifies class
# probabilities for nineteen classes.  To display the results, we will
# need to individually convert all of the PyTorch Tensors in the tuple
# into numpy arrays.
seg_tuple = tuple(t.cpu().data.numpy() for t in seg)

# Specify which one of the nineteen classes, class_index, are to be
# displayed.  Remove the redundant dimension from the array and
# plot the result.
class_index = 0
seg_imgplot = plt.imshow(numpy.squeeze(seg_tuple)[class_index])

# Convert the edge response from a PyTorch Tensor to a numpy array,
# remove the redundant dimension, and plot the result.
edge_imgplot = plt.imshow(numpy.squeeze(edge.cpu().data.numpy()))

If you want to display the result as a colormap, then you'll want to use some of the functions from https://github.com/NVIDIA/semantic-segmentation particularly those from the ImageDumper class.

as24-asingla commented 2 years ago

Amazing, Thank you. I will give it a try.