Update app.py

app.py

@@ -92,7 +92,6 @@ def predict(image):
     image = transforms(image).unsqueeze(0)
     DIS_map = model.inference(image.to(device),depth.to(device))[0][0][0].cpu()
     DIS_map = cv2.resize(np.array(DIS_map), (W,H))
-    # return cv2.resize(np.array(depth[0][0]), (W,H))
     return DIS_map
 
 with gr.Blocks(css=css) as demo:
@@ -102,29 +101,24 @@ with gr.Blocks(css=css) as demo:
 
     with gr.Row():
         input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input')
-        dis_image_slider = gr.Image(label="Pedict View", type='numpy', elem_id='img-display-input')
-        # dis_image_slider = ImageSlider(label="Pedict View", elem_id='img-display-output', position=0.5)
+        dis_image_slider = gr.Image(label="Pedict View",type='numpy', elem_id='img-display-output')
+        # dis_image_slider = ImageSlider(label="Pedict View", type="pil", elem_id='img-display-output')
     submit = gr.Button(value="Compute")
-    raw_file = gr.File(label="16-bit raw output", elem_id="download",)
 
     def on_submit(image):
         original_image = image.copy()
 
-        DIS_map = predict(image)
+        DIS_map = predict(np.array(image))
         DIS_map = (DIS_map - DIS_map.min()) / (DIS_map.max() - DIS_map.min()) * 255.0
-        raw_DIS_map = Image.fromarray(DIS_map.astype('uint16'))
-        tmp_raw_DIS_map = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
-        raw_DIS_map.save(tmp_raw_DIS_map.name)
+        matting = (DIS_map[...,None] / 255.0 * original_image)
+        return matting.astype('uint8')
 
-        # return [[original_image, DIS_map.astype(np.uint8)], tmp_raw_DIS_map.name]
-        return [DIS_map.astype(np.uint8), tmp_raw_DIS_map.name]
-
-    submit.click(on_submit, inputs=[input_image], outputs=[dis_image_slider, raw_file])
+    submit.click(on_submit, inputs=[input_image], outputs=dis_image_slider)
 
     example_files = os.listdir('assets/examples')
     example_files.sort()
     example_files = [os.path.join('assets/examples', filename) for filename in example_files]
-    examples = gr.Examples(examples=example_files, inputs=[input_image], outputs=[dis_image_slider, raw_file], fn=on_submit)
+    examples = gr.Examples(examples=example_files, inputs=[input_image], outputs=dis_image_slider, fn=on_submit)
 
 if __name__ == '__main__':
     demo.queue().launch(share=True)