Refactor image processing code to convert PyTorch tensor to NumPy array and apply color map

app.py

@@ -102,16 +102,22 @@ 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")
+
+            # Convert the PyTorch tensor to a NumPy array and squeeze
             pred_depth = pred_depth.cpu().numpy().squeeze()
-            pred_outpu_depth = pred_depth.cpu().numpy().astype(np.uint8)
+
+            # Convert to uint8 if necessary for the colormap
+            pred_output_depth = pred_depth.astype(np.uint8)
+
+            # Apply color map
             output_image = cv2.applyColorMap(
-                pred_outpu_depth, cv2.COLORMAP_INFERNO)[:, :, ::-1]
+                pred_output_depth, cv2.COLORMAP_INFERNO)[:, :, ::-1]
 
             # Continue with your file saving operations
             tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
-            output_image.save(tmp.name)
+            cv2.imwrite(tmp.name, output_image)
 
-            return [(original_image, original_image), tmp.name]
+            return [(original_image, output_image), tmp.name]
 
     submit.click(on_submit, inputs=[input_image], outputs=[
                  depth_image_slider, raw_file])