Update app.py

app.py

@@ -48,13 +48,13 @@ def get_canny_filter(image):
     image = np.concatenate([image, image, image], axis=2)
     canny_image = Image.fromarray(image)
     return canny_image
-def generate_images(canvas):        
+def generate_images(canvas,output_type):        
     base64_img = canvas
     image_data = base64.b64decode(base64_img.split(',')[1])
     input_img = Image.open(BytesIO(image_data)).convert(
         'RGB').resize((512, 512))
     input_img = input_img.filter(ImageFilter.GaussianBlur(radius=2))
-    input_img = get_canny_filter(input_img)
+    input_img = get_canny_filter(input_img) if output_type == "Canny" else pass
     return input_img
 def placeholder_fn(axis):
     pass
@@ -64,6 +64,7 @@ with gr.Blocks() as b:
         with gr.Column():
             rotation_axis = gr.Radio(choices=["x", "y", "z"], value="x", label="Joint rotation axis")
             pose_template = gr.Radio(choices=["regular", "ballet", "handstand", "split", "kick", "chilling"], value="regular", label="Pose template")
+            output_type = gr.Radio(choices=["Depth", "Canny"], value="Depth", label="Image Output")
             run_button = gr.Button("Generate")
             out_im = gr.Image()
     
@@ -78,7 +79,7 @@ with gr.Blocks() as b:
                             queue=False,
                             _js=js_pose_template)
     run_button.click(fn=generate_images,
-                     inputs=[canvas],
+                     inputs=[canvas,output_type],
                      outputs=[out_im],
                      _js=get_js_image)
     b.load(None,None,None,_js=load_js)