Refactor depth image saving logic

app.py

@@ -105,8 +105,17 @@ with gr.Blocks(css=css) as demo:
             # pred_depth = pred_depth.cpu().numpy().squeeze()
             # output_image = plt.imsave('depth.png', pred_depth, cmap='jet')
 
-            raw_depth = Image.fromarray(
-                pred_depth.cpu().numpy().astype('uint16'))
+            # 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_16bit)
+
+            # Continue with your file saving operations
             tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
             raw_depth.save(tmp.name)