How to convert PyTorch tensor to image and send it with flask?

Question

I have a neural network in PyTorch which gives an image as a Tensor. I have converted it to a numpy array and then followed the explanation here to send that image to the html. The problem is that it's always black.

This is my code in the PyTorch:

def getLinearImage():
  with torch.no_grad():
      test_z = torch.randn(1, 100).to(device)
      generated = G(test_z)
      generated = generated.cpu()
      numpy = generated[0].view(64, 64).numpy()
      numpu = (numpy + 1) * 128
      return numpy

This is the code in the flask where arr is the returned value from getLinearImage()

def getImage(arr):
    img = Image.fromarray(arr.astype("uint8"))
    file_object = io.BytesIO()

    img.save(file_object, "PNG")  
    file_object.seek(0)

    return send_file(file_object, mimetype="image/PNG")

If I open a static image and I send it to getImage() it works but won't work with the generated one. In the html I call it like: <img src="/getLinearImage" alt="User Image" height="100px" width="100px">

Answer 1

Logically speaking, since the static image works, the error is somewhere in your getLinearImage code. I would suggest running things through using PDB (or a debugger of your choice) to figure out why it's not generated correctly.

That said, I create a variable in your code:

numpu = (numpy + 1) * 128

which you don't seem to use, since you return the other variable afterwards:

return numpy

Could that be your problem?

Also: I presume that when you created this, you saved the original image locally to ensure something gets generated in the first place?