Refactor depth image processing and saving

app.py

@@ -102,22 +102,13 @@ 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')
-
-            # Remove batch and channel dimensions
-            pred_depth_squeezed = pred_depth.squeeze()
-
-            # If pred_depth is in an acceptable range for 16-bit representation
-            # No need for further normalization
-            # pred_depth_16bit = pred_depth_squeezed.cpu().numpy().astype('uint16')
-
-            # Convert to PIL image
-            raw_depth = Image.fromarray(pred_depth_squeezed)
+            pred_depth = pred_depth.cpu().numpy().squeeze()
+            output_image = cv2.applyColorMap(
+                pred_depth, cv2.COLORMAP_INFERNO)[:, :, ::-1]
 
             # Continue with your file saving operations
             tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
-            raw_depth.save(tmp.name)
+            output_image.save(tmp.name)
 
             return [(original_image, original_image), tmp.name]