Spaces:
				
			
			
	
			
			
		Paused
		
	
	
	
			
			
	
	
	
	
		
		
		Paused
		
	Update app.py
Browse files
    	
        app.py
    CHANGED
    
    | @@ -15,35 +15,6 @@ os.environ["PIXEL3DMM_CODE_BASE"] = f"{os.getcwd()}" | |
| 15 | 
             
            os.environ["PIXEL3DMM_PREPROCESSED_DATA"] = f"{os.getcwd()}/proprocess_results"
         | 
| 16 | 
             
            os.environ["PIXEL3DMM_TRACKING_OUTPUT"] = f"{os.getcwd()}/tracking_results"
         | 
| 17 |  | 
| 18 | 
            -
             | 
| 19 | 
            -
            def sh(cmd): subprocess.check_call(cmd, shell=True)
         | 
| 20 | 
            -
             | 
| 21 | 
            -
            # only do this once per VM restart
         | 
| 22 | 
            -
            sh("pip install -e .")
         | 
| 23 | 
            -
            sh("cd src/pixel3dmm/preprocessing/facer && pip install -e .")
         | 
| 24 | 
            -
            sh("cd src/pixel3dmm/preprocessing/PIPNet/FaceBoxesV2/utils && sh make.sh")
         | 
| 25 | 
            -
             | 
| 26 | 
            -
             | 
| 27 | 
            -
            def install_cuda_toolkit():
         | 
| 28 | 
            -
                CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
         | 
| 29 | 
            -
                CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
         | 
| 30 | 
            -
                subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
         | 
| 31 | 
            -
                subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
         | 
| 32 | 
            -
                subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
         | 
| 33 | 
            -
             | 
| 34 | 
            -
                os.environ["CUDA_HOME"] = "/usr/local/cuda"
         | 
| 35 | 
            -
                os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
         | 
| 36 | 
            -
                os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
         | 
| 37 | 
            -
                    os.environ["CUDA_HOME"],
         | 
| 38 | 
            -
                    "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
         | 
| 39 | 
            -
                )
         | 
| 40 | 
            -
                # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
         | 
| 41 | 
            -
                os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0"
         | 
| 42 | 
            -
                print("==> finished installation")
         | 
| 43 | 
            -
                
         | 
| 44 | 
            -
             | 
| 45 | 
            -
            install_cuda_toolkit()
         | 
| 46 | 
            -
             | 
| 47 | 
             
            # Utility to select first image from a folder
         | 
| 48 | 
             
            def first_image_from_dir(directory):
         | 
| 49 | 
             
                patterns = ["*.jpg", "*.png", "*.jpeg"]
         | 
| @@ -73,7 +44,7 @@ def reset_all(): | |
| 73 | 
             
            @spaces.GPU()
         | 
| 74 | 
             
            def preprocess_image(image_array, state):
         | 
| 75 | 
             
                if image_array is None:
         | 
| 76 | 
            -
                    return "β Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive= | 
| 77 |  | 
| 78 | 
             
                session_id = str(uuid.uuid4())
         | 
| 79 | 
             
                base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id)
         | 
| @@ -96,14 +67,14 @@ def preprocess_image(image_array, state): | |
| 96 |  | 
| 97 | 
             
                crop_dir = os.path.join(base_dir, "cropped")
         | 
| 98 | 
             
                image = first_image_from_dir(crop_dir)
         | 
| 99 | 
            -
                return "β
 Step 1 complete. Ready for Normals.", image, state, gr.update(interactive= | 
| 100 |  | 
| 101 | 
             
            # Step 2: Normals inference β normals image
         | 
| 102 | 
             
            @spaces.GPU()
         | 
| 103 | 
             
            def step2_normals(state):
         | 
| 104 | 
             
                session_id = state.get("session_id")
         | 
| 105 | 
             
                if not session_id:
         | 
| 106 | 
            -
                    return "β State lost. Please start from Step 1.", None, state, gr.update(interactive= | 
| 107 |  | 
| 108 | 
             
                try:
         | 
| 109 | 
             
                    p = subprocess.run([
         | 
| @@ -115,14 +86,14 @@ def step2_normals(state): | |
| 115 |  | 
| 116 | 
             
                normals_dir = os.path.join(state["base_dir"], "p3dmm", "normals")
         | 
| 117 | 
             
                image = first_image_from_dir(normals_dir)
         | 
| 118 | 
            -
                return "β
 Step 2 complete. Ready for UV Map.", image, state, gr.update(interactive= | 
| 119 |  | 
| 120 | 
             
            # Step 3: UV map inference β uv map image
         | 
| 121 | 
             
            @spaces.GPU()
         | 
| 122 | 
             
            def step3_uv_map(state):
         | 
| 123 | 
             
                session_id = state.get("session_id")
         | 
| 124 | 
             
                if not session_id:
         | 
| 125 | 
            -
                    return "β State lost. Please start from Step 1.", None, state, gr.update(interactive= | 
| 126 |  | 
| 127 | 
             
                try:
         | 
| 128 | 
             
                    p = subprocess.run([
         | 
| @@ -134,14 +105,14 @@ def step3_uv_map(state): | |
| 134 |  | 
| 135 | 
             
                uv_dir = os.path.join(state["base_dir"], "p3dmm", "uv_map")
         | 
| 136 | 
             
                image = first_image_from_dir(uv_dir)
         | 
| 137 | 
            -
                return "β
 Step 3 complete. Ready for Tracking.", image, state, gr.update(interactive= | 
| 138 |  | 
| 139 | 
             
            # Step 4: Tracking β final tracking image
         | 
| 140 | 
             
            @spaces.GPU()
         | 
| 141 | 
             
            def step4_track(state):
         | 
| 142 | 
             
                session_id = state.get("session_id")
         | 
| 143 | 
             
                if not session_id:
         | 
| 144 | 
            -
                    return "β State lost. Please start from Step 1.", None, state, gr.update(interactive= | 
| 145 |  | 
| 146 | 
             
                script = os.path.join(os.environ["PIXEL3DMM_CODE_BASE"], "scripts", "track.py")
         | 
| 147 | 
             
                try:
         | 
| @@ -154,7 +125,7 @@ def step4_track(state): | |
| 154 |  | 
| 155 | 
             
                tracking_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "frames")
         | 
| 156 | 
             
                image = first_image_from_dir(tracking_dir)
         | 
| 157 | 
            -
                return "β
 Pipeline complete!", image, state, gr.update(interactive= | 
| 158 |  | 
| 159 | 
             
            # Build Gradio UI
         | 
| 160 | 
             
            demo = gr.Blocks()
         | 
|  | |
| 15 | 
             
            os.environ["PIXEL3DMM_PREPROCESSED_DATA"] = f"{os.getcwd()}/proprocess_results"
         | 
| 16 | 
             
            os.environ["PIXEL3DMM_TRACKING_OUTPUT"] = f"{os.getcwd()}/tracking_results"
         | 
| 17 |  | 
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
| 18 | 
             
            # Utility to select first image from a folder
         | 
| 19 | 
             
            def first_image_from_dir(directory):
         | 
| 20 | 
             
                patterns = ["*.jpg", "*.png", "*.jpeg"]
         | 
|  | |
| 44 | 
             
            @spaces.GPU()
         | 
| 45 | 
             
            def preprocess_image(image_array, state):
         | 
| 46 | 
             
                if image_array is None:
         | 
| 47 | 
            +
                    return "β Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive=False)
         | 
| 48 |  | 
| 49 | 
             
                session_id = str(uuid.uuid4())
         | 
| 50 | 
             
                base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id)
         | 
|  | |
| 67 |  | 
| 68 | 
             
                crop_dir = os.path.join(base_dir, "cropped")
         | 
| 69 | 
             
                image = first_image_from_dir(crop_dir)
         | 
| 70 | 
            +
                return "β
 Step 1 complete. Ready for Normals.", image, state, gr.update(interactive=False), gr.update(interactive=True)
         | 
| 71 |  | 
| 72 | 
             
            # Step 2: Normals inference β normals image
         | 
| 73 | 
             
            @spaces.GPU()
         | 
| 74 | 
             
            def step2_normals(state):
         | 
| 75 | 
             
                session_id = state.get("session_id")
         | 
| 76 | 
             
                if not session_id:
         | 
| 77 | 
            +
                    return "β State lost. Please start from Step 1.", None, state, gr.update(interactive=False), gr.update(interactive=False)
         | 
| 78 |  | 
| 79 | 
             
                try:
         | 
| 80 | 
             
                    p = subprocess.run([
         | 
|  | |
| 86 |  | 
| 87 | 
             
                normals_dir = os.path.join(state["base_dir"], "p3dmm", "normals")
         | 
| 88 | 
             
                image = first_image_from_dir(normals_dir)
         | 
| 89 | 
            +
                return "β
 Step 2 complete. Ready for UV Map.", image, state, gr.update(interactive=False), gr.update(interactive=True)
         | 
| 90 |  | 
| 91 | 
             
            # Step 3: UV map inference β uv map image
         | 
| 92 | 
             
            @spaces.GPU()
         | 
| 93 | 
             
            def step3_uv_map(state):
         | 
| 94 | 
             
                session_id = state.get("session_id")
         | 
| 95 | 
             
                if not session_id:
         | 
| 96 | 
            +
                    return "β State lost. Please start from Step 1.", None, state, gr.update(interactive=False), gr.update(interactive=False)
         | 
| 97 |  | 
| 98 | 
             
                try:
         | 
| 99 | 
             
                    p = subprocess.run([
         | 
|  | |
| 105 |  | 
| 106 | 
             
                uv_dir = os.path.join(state["base_dir"], "p3dmm", "uv_map")
         | 
| 107 | 
             
                image = first_image_from_dir(uv_dir)
         | 
| 108 | 
            +
                return "β
 Step 3 complete. Ready for Tracking.", image, state, gr.update(interactive=False), gr.update(interactive=True)
         | 
| 109 |  | 
| 110 | 
             
            # Step 4: Tracking β final tracking image
         | 
| 111 | 
             
            @spaces.GPU()
         | 
| 112 | 
             
            def step4_track(state):
         | 
| 113 | 
             
                session_id = state.get("session_id")
         | 
| 114 | 
             
                if not session_id:
         | 
| 115 | 
            +
                    return "β State lost. Please start from Step 1.", None, state, gr.update(interactive=False)
         | 
| 116 |  | 
| 117 | 
             
                script = os.path.join(os.environ["PIXEL3DMM_CODE_BASE"], "scripts", "track.py")
         | 
| 118 | 
             
                try:
         | 
|  | |
| 125 |  | 
| 126 | 
             
                tracking_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "frames")
         | 
| 127 | 
             
                image = first_image_from_dir(tracking_dir)
         | 
| 128 | 
            +
                return "β
 Pipeline complete!", image, state, gr.update(interactive=False)
         | 
| 129 |  | 
| 130 | 
             
            # Build Gradio UI
         | 
| 131 | 
             
            demo = gr.Blocks()
         | 
