Update app.py

app.py

@@ -36,7 +36,7 @@ css = """
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 dtype = torch.float32
 variant = None
-checkpoint_path = "GonzaloMG/marigold-e2e-ft-depth"
+checkpoint_path = "GonzaloMG/marigold-e2e-ft-normals"
 unet         = UNet2DConditionModel.from_pretrained(checkpoint_path, subfolder="unet")   
 vae          = AutoencoderKL.from_pretrained(checkpoint_path, subfolder="vae")  
 text_encoder = CLIPTextModel.from_pretrained(checkpoint_path, subfolder="text_encoder")  
@@ -55,28 +55,28 @@ pipe = pipe.to(DEVICE)
 pipe.unet.eval()
 
 
-title = "# End-to-End Fine-Tuned Marigold for Depth Estimation"
+title = "# End-to-End Fine-Tuned Marigold for Normals Estimation"
 description = """ Please refer to our [paper](https://arxiv.org/abs/2409.11355) and [GitHub](https://vision.rwth-aachen.de/diffusion-e2e-ft) for more details."""
     
 @spaces.GPU
-def predict_depth(image, processing_res_choice):
+def predict_normals(image, processing_res_choice):
     with torch.no_grad():
-        pipe_out = pipe(image, denoising_steps=1, ensemble_size=1, noise="zeros", normals=False, processing_res=processing_res_choice, match_input_res=True)
-    pred = pipe_out.depth_np
-    pred_colored = pipe_out.depth_colored
+        pipe_out = pipe(image, denoising_steps=1, ensemble_size=1, noise="zeros", normals=True, processing_res=processing_res_choice, match_input_res=True)
+    pred = pipe_out.normal_np
+    pred_colored = pipe_out.normal_colored
     return pred, pred_colored
 
 with gr.Blocks(css=css) as demo:
     gr.Markdown(title)
     gr.Markdown(description)
-    gr.Markdown("### Depth Prediction demo")
+    gr.Markdown("### Normals Prediction demo")
 
     with gr.Row():
         input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input')
-        depth_image_slider = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5)
+        normals_image_slider = ImageSlider(label="Surface Normals with Slider View", elem_id='img-display-output', position=0.5)
 
     with gr.Row():
-        submit = gr.Button(value="Compute Depth")
+        submit = gr.Button(value="Compute Normals")
         processing_res_choice = gr.Radio(
                 [
                     ("Recommended (768)", 768),
@@ -86,8 +86,7 @@ with gr.Blocks(css=css) as demo:
                 value=768,
             )
 
-    gray_depth_file = gr.File(label="Grayscale depth map", elem_id="download",)
-    raw_file = gr.File(label="Raw Depth Data (.npy)", elem_id="download")
+    raw_file = gr.File(label="Raw Normals Data (.npy)", elem_id="download")
 
     cmap = matplotlib.colormaps.get_cmap('Spectral_r')
 
@@ -98,30 +97,30 @@ with gr.Blocks(css=css) as demo:
             return None
     
         pil_image = Image.fromarray(image.astype('uint8'))
-        depth_npy, depth_colored = predict_depth(pil_image, processing_res_choice)
+        normal_npy, normal_colored = predict_normals(pil_image, processing_res_choice)
     
-        # Save the npy data (raw depth map)
-        tmp_npy_depth = tempfile.NamedTemporaryFile(suffix='.npy', delete=False)
-        np.save(tmp_npy_depth.name, depth_npy)
+        # Save the npy data (raw normals)
+        tmp_npy_normal = tempfile.NamedTemporaryFile(suffix='.npy', delete=False)
+        np.save(tmp_npy_normal.name, normal_npy)
     
         # Save the grayscale depth map
-        depth_gray = (depth_npy * 65535.0).astype(np.uint16)
-        tmp_gray_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
-        Image.fromarray(depth_gray).save(tmp_gray_depth.name, mode="I;16")
+        # depth_gray = (depth_npy * 65535.0).astype(np.uint16)
+        # tmp_gray_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
+        # Image.fromarray(depth_gray).save(tmp_gray_depth.name, mode="I;16")
     
-        # Save the colored depth map
-        tmp_colored_depth = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
-        depth_colored.save(tmp_colored_depth.name)
+        # Save the colored normals map
+        tmp_colored_normal = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
+        normal_colored.save(tmp_colored_normal.name)
    
-        return [(image, depth_colored),  tmp_gray_depth.name, tmp_npy_depth.name]
+        return [(image, normal_colored), tmp_npy_normal.name]
 
-    submit.click(on_submit, inputs=[input_image, processing_res_choice], outputs=[depth_image_slider, gray_depth_file, raw_file])
+    submit.click(on_submit, inputs=[input_image, processing_res_choice], outputs=[normals_image_slider, raw_file])
 
     example_files = os.listdir('assets/examples')
     example_files.sort()
     example_files = [os.path.join('assets/examples', filename) for filename in example_files]
     example_files = [[image, 768] for image in example_files]
-    examples = gr.Examples(examples=example_files, inputs=[input_image, processing_res_choice], outputs=[depth_image_slider, gray_depth_file, raw_file], fn=on_submit)
+    examples = gr.Examples(examples=example_files, inputs=[input_image, processing_res_choice], outputs=[normals_image_slider, raw_file], fn=on_submit)
 
 
 if __name__ == '__main__':