Refactor depth image saving logic

app.py

@@ -102,10 +102,15 @@ with gr.Blocks(css=css) as demo:
             pred_depth = post_process_depth(
                 pred_depths_r_list[-1], pred_depths_r_list_flipped[-1])
             print("== Finished processing image")
-            pred_depth = pred_depth.cpu().numpy().squeeze()
-            output_image = plt.imsave('depth.png', pred_depth, cmap='jet')
+            # pred_depth = pred_depth.cpu().numpy().squeeze()
+            # output_image = plt.imsave('depth.png', pred_depth, cmap='jet')
 
-            return [(original_image, original_image), output_image]
+            raw_depth = Image.fromarray(
+                pred_depth.cpu().numpy().astype('uint16'))
+            tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
+            raw_depth.save(tmp.name)
+
+            return [(original_image, original_image), tmp.name]
 
     submit.click(on_submit, inputs=[input_image], outputs=[
                  depth_image_slider, raw_file])