Update app.py

app.py

@@ -101,8 +101,8 @@ 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-output')
-        # dis_image_slider = ImageSlider(label="Pedict View", type="pil", elem_id='img-display-output')
+        dis_image = 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',upload_count=2)
     submit = gr.Button(value="Compute")
 
     def on_submit(image):
@@ -110,15 +110,16 @@ with gr.Blocks(css=css) as demo:
 
         DIS_map = predict(np.array(image))
         DIS_map = (DIS_map - DIS_map.min()) / (DIS_map.max() - DIS_map.min()) * 255.0
-        matting = (DIS_map[...,None] / 255.0 * original_image) + (255-DIS_map[...,None])
-        return matting.astype('uint8')
+        # matting = (DIS_map[...,None] / 255.0 * original_image) + (255-DIS_map[...,None])
+        alpha_img = np.concatenate([np.array(original_image),DIS_map[...,None]],axis=-1).astype(np.uint16)
+        return alpha_img
 
-    submit.click(on_submit, inputs=[input_image], outputs=dis_image_slider)
+    submit.click(on_submit, inputs=[input_image], outputs=dis_image)
 
     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, fn=on_submit)
+    examples = gr.Examples(examples=example_files, inputs=[input_image], outputs=dis_image, fn=on_submit)
 
 if __name__ == '__main__':
     demo.queue().launch(share=True)