|
import spaces |
|
import gradio as gr |
|
import os |
|
import shutil |
|
from pathlib import Path |
|
from toon3d.scripts.viser_vis import main as viser_vis_main |
|
import viser |
|
import time |
|
import threading |
|
|
|
viewer_thread_instance = None |
|
stop_event = threading.Event() |
|
shared_url = None |
|
|
|
_HEADER_ = ''' |
|
<h2>Toon3D: Seeing Cartoons from a New Perspective</h2> |
|
Toon3D lifts hand-drawn images to 3D with a piecewise-rigid deformation optimization at hand-labeled keypoints and using monocular depth as a prior. The project page is at <a href='https://toon3d.studio/' target='_blank'>https://toon3d.studio/</a> and the Toon3D Labeler is at <a href='https://labeler.toon3d.studio/' target='_blank'>https://labeler.toon3d.studio/</a>. Follow the steps below to run Toon3D! |
|
|
|
<div style="margin-top: 20px; font-size: 16px; line-height: 1.6;"> |
|
<div style="display: flex; justify-content: space-between;"> |
|
<div style="width: 49%;"> |
|
<ol> |
|
<li><strong>Prepare and Process Data</strong> |
|
<ul> |
|
<li>Upload images and click on "Process Data" to generate processed data.</li> |
|
<li>Download the processed data.</li> |
|
</ul> |
|
</li> |
|
<li><strong>Label Data</strong> |
|
<ul> |
|
<li>Upload the processed data and label points using the labeler ("Upload ZIP").</li> |
|
<li>Click export and upload the points.json to the "Labeled Points" section.</li> |
|
</ul> |
|
</li> |
|
</ol> |
|
</div> |
|
<div style="width: 49%;"> |
|
<ol start="3"> |
|
<li><strong>Generate 3D Output</strong> |
|
<ul> |
|
<li>Click on "Run Toon3D" to run the structure from motion pipeline.</li> |
|
<li>Download the output and inspect locally (point cloud, mesh, Nerfstudio dataset).</li> |
|
</ul> |
|
</li> |
|
<li><strong>View in Web!</strong> |
|
<ul> |
|
<li>Click on "Open Viewer" to view the output in an interactive viewer powered by <a href="https://viser.studio/">Viser</a>.</li> |
|
</ul> |
|
<ul> |
|
<li>Reach out if you have any questions!</li> |
|
</ul> |
|
</li> |
|
</ol> |
|
</div> |
|
</div> |
|
</div> |
|
''' |
|
|
|
def check_input_images(input_images): |
|
if input_images is None: |
|
raise gr.Error("No images uploaded!") |
|
|
|
@spaces.GPU(duration=120) |
|
def process_images(input_images, compute_segment_anything): |
|
|
|
images_path = "/tmp/gradio/images" |
|
processed_path = "/tmp/gradio/processed" |
|
|
|
|
|
os.system(f"rm -rf {images_path}") |
|
os.system(f"rm -rf {processed_path}") |
|
|
|
|
|
os.system(f"mkdir -p {images_path}") |
|
os.system(f"mkdir -p {processed_path}") |
|
|
|
for fileobj in input_images: |
|
shutil.copyfile(fileobj.name, images_path + "/" + os.path.basename(fileobj.name)) |
|
|
|
|
|
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 = " --no-compute-segment-anything" |
|
|
|
|
|
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" |
|
shutil.make_archive(zip_folder, 'zip', zip_folder) |
|
|
|
return zip_folder + ".zip" |
|
|
|
def toggle_labeler_visibility(visible): |
|
if visible: |
|
return '<iframe src="https://labeler.toon3d.studio/" style="display: block; margin: auto; width: 100%; height: 100vh;" frameborder="0"></iframe>' |
|
else: |
|
return "" |
|
|
|
def check_input_toon3d(processed_data_zip, labeled_data): |
|
if processed_data_zip is None: |
|
raise gr.Error("No images uploaded!") |
|
|
|
@spaces.GPU(duration=120) |
|
def run_toon3d(processed_data_zip, labeled_data): |
|
|
|
data_prefix = "/tmp/gradio/inputs" |
|
processed_path = f"{data_prefix}/toon3d-dataset" |
|
output_prefix = "/tmp/gradio/outputs" |
|
nerfstudio_folder = "/tmp/gradio/nerfstudio" |
|
|
|
os.system(f"rm -rf {processed_path}") |
|
os.system(f"rm -rf {output_prefix}") |
|
os.system(f"rm -rf {nerfstudio_folder}") |
|
|
|
shutil.unpack_archive(processed_data_zip.name, processed_path) |
|
shutil.copyfile(labeled_data.name, f"{processed_path}/points.json") |
|
|
|
|
|
toon3d_cmd = f"tnd-run --dataset toon3d-dataset --data_prefix {data_prefix} --output_prefix {output_prefix} --nerfstudio_folder {nerfstudio_folder} --no-view-point-cloud" |
|
os.system(toon3d_cmd) |
|
|
|
|
|
|
|
output_dirs = Path(output_prefix) / "toon3d-dataset" / "run" |
|
output_dir = Path(output_dirs / sorted(os.listdir(output_dirs))[-1]) |
|
|
|
zip_folder = str(output_dir) |
|
shutil.make_archive(zip_folder, 'zip', zip_folder) |
|
|
|
return zip_folder + ".zip" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def viewer_thread(processed_data_zip, labeled_data, toon3d_output_zip): |
|
global shared_url |
|
data_prefix = Path("/tmp/gradio/inputs") |
|
processed_path = f"{data_prefix}/toon3d-dataset" |
|
|
|
viewer_folder = "/tmp/gradio/viewer/toon3d-dataset/run/temp" |
|
os.system(f"rm -rf {viewer_folder}") |
|
shutil.unpack_archive(toon3d_output_zip.name, viewer_folder) |
|
shutil.unpack_archive(processed_data_zip.name, processed_path) |
|
shutil.copyfile(labeled_data.name, f"{processed_path}/points.json") |
|
|
|
viser_server = viser.ViserServer() |
|
url = viser_server.request_share_url() |
|
shared_url = url |
|
print(url) |
|
|
|
viser_vis_main( |
|
data_prefix=data_prefix, |
|
dataset="toon3d-dataset", |
|
output_prefix=Path("/tmp/gradio/viewer"), |
|
output_method=Path("run"), |
|
server=viser_server, |
|
visible=True, |
|
return_early=True |
|
) |
|
while not stop_event.is_set(): |
|
time.sleep(1) |
|
|
|
viser_server.stop() |
|
|
|
def kill_viewer(): |
|
global viewer_thread_instance, stop_event |
|
if viewer_thread_instance and viewer_thread_instance.is_alive(): |
|
stop_event.set() |
|
viewer_thread_instance.join() |
|
viewer_thread_instance = None |
|
print("Viewer has been stopped.") |
|
else: |
|
print("No viewer is running.") |
|
|
|
def get_html_for_shared_url(url): |
|
return f'<h1>Open <a href="{url}" target="_blank">{url}</a>!</h1>' |
|
|
|
def check_input_open_viewer(processed_data_zip, labeled_data, toon3d_output_zip): |
|
if processed_data_zip is None: |
|
raise gr.Error("No processed data uploaded!") |
|
if labeled_data is None: |
|
raise gr.Error("No labeled points uploaded!") |
|
if toon3d_output_zip is None: |
|
raise gr.Error("No Toon3D output uploaded!") |
|
|
|
def start_viewer(processed_data_zip, labeled_data, toon3d_output_zip): |
|
kill_viewer() |
|
|
|
global viewer_thread_instance, stop_event, shared_url |
|
stop_event.clear() |
|
shared_url = None |
|
if viewer_thread_instance is None or not viewer_thread_instance.is_alive(): |
|
viewer_thread_instance = threading.Thread(target=viewer_thread, args=(processed_data_zip, labeled_data, toon3d_output_zip)) |
|
viewer_thread_instance.start() |
|
while not shared_url: |
|
|
|
time.sleep(0.1) |
|
return get_html_for_shared_url(shared_url) |
|
else: |
|
print("Viewer is already running.") |
|
return get_html_for_shared_url(shared_url) |
|
|
|
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"]) |
|
with gr.Column(): |
|
compute_segment_anything = gr.Checkbox(label="Compute Segment Anything? (slow)", value=False) |
|
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) |
|
with gr.Row(variant="panel"): |
|
labeler_frame = gr.HTML() |
|
labeler_visible.change(toggle_labeler_visibility, inputs=[labeler_visible], outputs=[labeler_frame]) |
|
with gr.Row(variant="panel"): |
|
labeled_data = gr.File(label="Labeled Points", file_count="single", file_types=[".json"]) |
|
run_toon3d_button = gr.Button("Run Toon3D", elem_id="run_toon3d_button", variant="primary") |
|
toon3d_output_zip = gr.File(label="Toon3D Output", file_count="single", file_types=[".zip"], interactive=True) |
|
with gr.Row(variant="panel"): |
|
open_viewer_button = gr.Button("Open Viewer", elem_id="open_viser_button", variant="primary") |
|
with gr.Row(variant="panel"): |
|
viser_link = gr.HTML() |
|
|
|
process_data_button.click(fn=check_input_images, inputs=[input_images]).success( |
|
fn=process_images, |
|
inputs=[input_images, compute_segment_anything], |
|
outputs=[processed_data_zip], |
|
) |
|
|
|
run_toon3d_button.click(fn=check_input_toon3d, inputs=[processed_data_zip, labeled_data]).success( |
|
fn=run_toon3d, |
|
inputs=[processed_data_zip, labeled_data], |
|
outputs=[toon3d_output_zip], |
|
) |
|
|
|
open_viewer_button.click(fn=check_input_open_viewer, inputs=[processed_data_zip, labeled_data, toon3d_output_zip]).success( |
|
fn=start_viewer, |
|
inputs=[processed_data_zip, labeled_data, toon3d_output_zip], |
|
outputs=[viser_link], |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.queue(max_size=10) |
|
demo.launch() |
|
|