Refactor image processing and saving logic

app.py

@@ -82,17 +82,13 @@ with gr.Blocks(css=css) as demo:
     def on_submit(image):
         original_image = image.copy()
 
+        # This is for resizing the image to 518x518
         h, w = image.shape[:2]
 
-        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
-        # image = transform({'image': image})['image']
-        # image = torch.from_numpy(image).unsqueeze(0).to(DEVICE)
-
         image = np.asarray(image, dtype=np.float32) / 255.0
         image = torch.from_numpy(image.transpose((2, 0, 1)))
         image = Normalize(mean=[0.485, 0.456, 0.406], std=[
             0.229, 0.224, 0.225])(image)
-        # image = torch.from_numpy(image).unsqueeze(0)
         with torch.no_grad():
             image = torch.autograd.Variable(image.unsqueeze(0))
             print("== Processing image")
@@ -106,19 +102,12 @@ with gr.Blocks(css=css) as demo:
             # Convert the PyTorch tensor to a NumPy array and squeeze
             pred_depth = pred_depth.cpu().numpy().squeeze()
 
-            # Convert to uint8 if necessary for the colormap
-            pred_output_depth = pred_depth.astype(np.uint8)
-
-            # Apply color map
-            output_image = cv2.applyColorMap(
-                pred_output_depth, cv2.COLORMAP_INFERNO)[:, :, ::-1]
-
             # Continue with your file saving operations
             tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
             # cv2.imwrite(tmp.name, output_image)
             plt.imsave(tmp.name, pred_depth, cmap='jet')
 
-            return [(original_image, output_image), tmp.name]
+            return [(original_image, tmp.name), tmp.name]
 
     submit.click(on_submit, inputs=[input_image], outputs=[
                  depth_image_slider, raw_file])