gradio update

app.py

@@ -11,6 +11,7 @@ import threading
 viewer_thread_instance = None
 stop_event = threading.Event()
 shared_url = None
+compute_segment_anything = False
 
 _HEADER_ = '''
 <h2>Toon3D: Seeing Cartoons from a New Perspective</h2>
@@ -61,7 +62,7 @@ def check_input_images(input_images):
         raise gr.Error("No images uploaded!")
 
 @spaces.GPU(duration=120)
-def process_images(input_images):
+def process_images(input_images, compute_segment_anything):
 
     images_path = "/tmp/gradio/images"
     processed_path = "/tmp/gradio/processed"
@@ -81,8 +82,15 @@ def process_images(input_images):
     download_cmd = "tnd-download-data sam --save-dir /tmp/gradio"
     os.system(download_cmd)
 
+    if compute_segment_anything:
+        sam_cmd = " --compute-segment-anything"
+    else:
+        sam_cmd = ""
+
     # process the data
     process_data_cmd = f"tnd-process-data initialize --dataset toon3d-dataset --input_path {images_path} --data_prefix {processed_path} --sam_checkpoint_prefix /tmp/gradio/sam-checkpoints"
+    process_data_cmd += sam_cmd
+    print(process_data_cmd)
     os.system(process_data_cmd)
 
     zip_folder = "/tmp/gradio/processed/toon3d-dataset"
@@ -95,6 +103,10 @@ def toggle_labeler_visibility(visible):
         return '<iframe src="https://labeler.toon3d.studio/" style="display: block; margin: auto; width: 100%; height: 100vh;" frameborder="0"></iframe>'
     else:
         return ""
+    
+def set_compute_segment_anything(value):
+    global compute_segment_anything
+    compute_segment_anything = value
 
 def check_input_toon3d(processed_data_zip, labeled_data):
     if processed_data_zip is None:
@@ -233,7 +245,10 @@ with gr.Blocks(title="Toon3D") as demo:
     gr.Markdown(_HEADER_)
     with gr.Row(variant="panel"):
             input_images = gr.File(label="Upload Images", file_count="multiple", file_types=[".jpg", "jpeg", "png"])
-            process_data_button = gr.Button("Process Data", elem_id="process_data_button", variant="primary")
+            with gr.Column():
+                compute_segment_anything = gr.Checkbox(label="Compute Segment Anything? (slow)", value=False)
+                compute_segment_anything.change(set_compute_segment_anything, inputs=[compute_segment_anything])
+                process_data_button = gr.Button("Process Data", elem_id="process_data_button", variant="primary")
             processed_data_zip = gr.File(label="Processed Data", file_count="single", file_types=[".zip"], interactive=True)
     with gr.Row(variant="panel"):
         labeler_visible = gr.Checkbox(label="Show Labeler", value=False)
@@ -251,7 +266,7 @@ with gr.Blocks(title="Toon3D") as demo:
 
     process_data_button.click(fn=check_input_images, inputs=[input_images]).success(
         fn=process_images,
-        inputs=[input_images],
+        inputs=[input_images, compute_segment_anything],
         outputs=[processed_data_zip],
     )
 

requirements.txt

@@ -1,2 +1,2 @@
-toon3d==0.0.2
+toon3d==0.0.3
 segment_anything @ git+https://github.com/facebookresearch/segment-anything.git