valhalla commited on
Commit
683367b
1 Parent(s): ce425e1

Update swin_ir_2.py

Browse files
Files changed (1) hide show
  1. swin_ir_2.py +2 -1
swin_ir_2.py CHANGED
@@ -1322,4 +1322,5 @@ def preprocesss_image(image: PIL.Image.Image) -> torch.FloatTensor:
1322
  def postprocess_image(output: torch.Tensor) -> PIL.Image.Image:
1323
  output = output.data.float().cpu().clamp_(0, 1).numpy()
1324
  output = (output * 255).round().astype("uint8")
1325
- output = output.transpose(1, 2, 0)
 
 
1322
  def postprocess_image(output: torch.Tensor) -> PIL.Image.Image:
1323
  output = output.data.float().cpu().clamp_(0, 1).numpy()
1324
  output = (output * 255).round().astype("uint8")
1325
+ output = output.transpose(1, 2, 0)
1326
+ return PIL.Image.fromarray(output)