import spaces import os import subprocess import tempfile import uuid import glob import shutil import time import gradio as gr import sys from PIL import Image # Set environment variables os.environ["PIXEL3DMM_CODE_BASE"] = f"{os.getcwd()}" os.environ["PIXEL3DMM_PREPROCESSED_DATA"] = f"{os.getcwd()}/proprocess_results" os.environ["PIXEL3DMM_TRACKING_OUTPUT"] = f"{os.getcwd()}/tracking_results" def sh(cmd): subprocess.check_call(cmd, shell=True) # only do this once per VM restart sh("pip install -e .") sh("cd src/pixel3dmm/preprocessing/facer && pip install -e .") sh("cd src/pixel3dmm/preprocessing/PIPNet/FaceBoxesV2/utils && sh make.sh") def install_cuda_toolkit(): CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run" CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL) subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE]) subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE]) subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"]) os.environ["CUDA_HOME"] = "/usr/local/cuda" os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"]) os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % ( os.environ["CUDA_HOME"], "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"], ) # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0" print("==> finished installation") install_cuda_toolkit() # Utility to select first image from a folder def first_image_from_dir(directory): patterns = ["*.jpg", "*.png", "*.jpeg"] files = [] for p in patterns: files.extend(glob.glob(os.path.join(directory, p))) if not files: return None return sorted(files)[0] # Step 1: Preprocess the input image (Save and Crop) @spaces.GPU() def preprocess_image(image_array, state): # Check if an image was uploaded if image_array is None: return "❌ Please upload an image first.", None, state # Step 1a: Save the uploaded image session_id = str(uuid.uuid4()) base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id) os.makedirs(base_dir, exist_ok=True) state.update({"session_id": session_id, "base_dir": base_dir}) img = Image.fromarray(image_array) saved_image_path = os.path.join(base_dir, f"{session_id}.png") img.save(saved_image_path) state["image_path"] = saved_image_path # Step 1b: Run the preprocessing script try: p = subprocess.run([ "python", "scripts/run_preprocessing.py", "--video_or_images_path", saved_image_path ], check=True, capture_output=True, text=True) except subprocess.CalledProcessError as e: err = f"❌ Preprocess failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}" # Clean up created directory on failure shutil.rmtree(base_dir) return err, None, state crop_dir = os.path.join(base_dir, "cropped") image = first_image_from_dir(crop_dir) return "✅ Preprocessing complete", image, state # Step 2: Normals inference → normals image @spaces.GPU() def step2_normals(state): session_id = state.get("session_id") if not session_id: return "❌ Please preprocess an image first.", None, state try: # Execute the network inference for normals p = subprocess.run([ "python", "scripts/network_inference.py", "model.prediction_type=normals", f"video_name={session_id}" ], check=True, capture_output=True, text=True) except subprocess.CalledProcessError as e: err = f"❌ Normal map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}" return err, None, state normals_dir = os.path.join(state["base_dir"], "p3dmm", "normals") image = first_image_from_dir(normals_dir) return "✅ Step 2: Normals inference complete", image, state # Step 3: UV map inference → uv map image @spaces.GPU() def step3_uv_map(state): session_id = state.get("session_id") if not session_id: return "❌ Please preprocess an image first.", None, state try: # Execute the network inference for UV map p = subprocess.run([ "python", "scripts/network_inference.py", "model.prediction_type=uv_map", f"video_name={session_id}" ], check=True, capture_output=True, text=True) except subprocess.CalledProcessError as e: err = f"❌ UV map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}" return err, None, state uv_dir = os.path.join(state["base_dir"], "p3dmm", "uv_map") image = first_image_from_dir(uv_dir) return "✅ Step 3: UV map inference complete", image, state # Step 4: Tracking → final tracking image @spaces.GPU() def step4_track(state): session_id = state.get("session_id") if not session_id: return "❌ Please preprocess an image first.", None, state script = os.path.join(os.environ["PIXEL3DMM_CODE_BASE"], "scripts", "track.py") try: # Execute the tracking script p = subprocess.run([ "python", script, f"video_name={session_id}" ], check=True, capture_output=True, text=True) except subprocess.CalledProcessError as e: err = f"❌ Tracking failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}" return err, None, state tracking_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "frames") image = first_image_from_dir(tracking_dir) return "✅ Step 4: Tracking complete", image, state # Build Gradio UI demo = gr.Blocks() with demo: gr.Markdown("## Image Processing Pipeline") with gr.Row(): with gr.Column(): image_in = gr.Image(label="Upload Image", type="numpy", height=512) status = gr.Textbox(label="Status", lines=2, interactive=False) state = gr.State({}) with gr.Column(): with gr.Row(): crop_img = gr.Image(label="Preprocessed", height=256) normals_img = gr.Image(label="Normals", height=256) with gr.Row(): uv_img = gr.Image(label="UV Map", height=256) track_img = gr.Image(label="Tracking", height=256) with gr.Row(): preprocess_btn = gr.Button("Step 1: Preprocess") normals_btn = gr.Button("Step 2: Normals") uv_map_btn = gr.Button("Step 3: UV Map") track_btn = gr.Button("Step 4: Track") # Pipeline execution preprocess_btn.click(fn=preprocess_image, inputs=[image_in, state], outputs=[status, crop_img, state]) normals_btn.click(fn=step2_normals, inputs=[state], outputs=[status, normals_img, state]) uv_map_btn.click(fn=step3_uv_map, inputs=[state], outputs=[status, uv_img, state]) track_btn.click(fn=step4_track, inputs=[state], outputs=[status, track_img, state]) # ------------------------------------------------------------------ # START THE GRADIO SERVER # ------------------------------------------------------------------ demo.queue() demo.launch(share=True, ssr_mode=False)