Spaces:
Sleeping
Sleeping
| import os | |
| import subprocess | |
| import time | |
| # REMOVE pip installs from here - manage via requirements.txt | |
| # os.system("pip install gradio==2.4.6") # REMOVE | |
| # os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'") # REMOVE | |
| # Import gradio AFTER potentially installing it (though it should already be there) | |
| try: | |
| import gradio as gr | |
| except ImportError: | |
| print("ERROR: Gradio not installed. Ensure it's in requirements.txt") | |
| exit() | |
| # --- Repository Setup --- | |
| repo_dir = "bizarre-pose-estimator" | |
| if not os.path.exists(repo_dir): | |
| print(f"Cloning repository '{repo_dir}'...") | |
| clone_command = "git clone https://github.com/ShuhongChen/bizarre-pose-estimator.git" | |
| try: | |
| subprocess.run(clone_command, shell=True, check=True, capture_output=True, text=True) | |
| print("Repository cloned successfully.") | |
| except subprocess.CalledProcessError as e: | |
| print(f"ERROR: Failed to clone repository: {e.stderr}") | |
| exit() | |
| else: | |
| print(f"Repository '{repo_dir}' already exists.") | |
| # --- Change Directory --- | |
| try: | |
| # Check if already in the directory (important for restarts) | |
| if os.path.basename(os.getcwd()) != repo_dir: | |
| os.chdir(repo_dir) | |
| print(f"Changed directory to: {os.getcwd()}") | |
| else: | |
| print(f"Already in directory: {os.getcwd()}") | |
| except FileNotFoundError: | |
| print(f"ERROR: Failed to change directory to '{repo_dir}'.") | |
| exit() | |
| # --- Download Example Image --- | |
| example_img_url = "https://i.imgur.com/IkJzlaE.jpeg" | |
| example_img_file = "IkJzlaE.jpeg" | |
| if not os.path.exists(example_img_file): | |
| print(f"Downloading example image: {example_img_file}") | |
| try: | |
| subprocess.run(["curl", "-fL", example_img_url, "-o", example_img_file], check=True) | |
| print("Example image downloaded.") | |
| except (subprocess.CalledProcessError, FileNotFoundError) as e: | |
| print(f"Warning: Failed to download example image {example_img_url}. Error: {e}") | |
| # --- Model Download --- | |
| model_zip_file = "bizarre_pose_models.zip" | |
| model_extract_dir = "bizarre_pose_models" | |
| model_id = "17N5PutpYJTlKuNB6bdDaiQsPSIkYtiPm" | |
| download_successful = False | |
| models_ready = False | |
| model_checkpoint_rel_path = "_train/character_pose_estim/runs/feat_concat+data.ckpt" | |
| if os.path.exists(model_checkpoint_rel_path): | |
| print("Model checkpoint already exists. Skipping download and unzip.") | |
| models_ready = True # Models are already copied and ready | |
| else: | |
| print(f"Model checkpoint not found at '{model_checkpoint_rel_path}'. Checking for zip...") | |
| if os.path.exists(model_zip_file): | |
| print(f"Zip file '{model_zip_file}' already exists. Skipping download.") | |
| download_successful = True | |
| else: | |
| print(f"Attempting to download model weights using gdown (ID: {model_id})...") | |
| gdown_command = f"gdown --id {model_id} -O {model_zip_file}" | |
| print(f"Executing: {gdown_command}") | |
| try: | |
| # Set timeout for gdown (e.g., 5 minutes) | |
| result = subprocess.run(gdown_command, shell=True, check=True, capture_output=True, text=True, timeout=300) | |
| print("Gdown download successful.") | |
| download_successful = True | |
| except FileNotFoundError: print("ERROR: 'gdown' command not found.") | |
| except subprocess.TimeoutExpired: print("ERROR: gdown download timed out.") | |
| except subprocess.CalledProcessError as e: | |
| print(f"ERROR: gdown download failed: {e.returncode}\nStderr:{e.stderr}\nStdout:{e.stdout}") | |
| # --- Unzip and Copy --- | |
| if download_successful and os.path.exists(model_zip_file): | |
| print(f"Unzipping '{model_zip_file}'...") | |
| try: | |
| if os.path.exists(model_extract_dir): print(f"Note: '{model_extract_dir}' already exists. Overwriting.") | |
| subprocess.run(["unzip", "-oq", model_zip_file], check=True) | |
| print("Unzip successful.") | |
| print("Copying model files...") | |
| source_dir_unzip = f"./{model_extract_dir}/" | |
| dest_dir = "." | |
| if os.path.exists(source_dir_unzip): | |
| os.system(f"cp -a {source_dir_unzip}. {dest_dir}") # Use os.system like original | |
| print("Model files copied.") | |
| models_ready = True # Models should now be ready | |
| print(f"Attempting to remove zip file: {model_zip_file}") | |
| try: | |
| os.remove(model_zip_file) | |
| print(f"Removed zip file: {model_zip_file}") | |
| except OSError as e: print(f"Warning: Error removing zip file {model_zip_file}: {e}") | |
| else: print(f"ERROR: Directory '{source_dir_unzip}' not found after unzip.") | |
| except FileNotFoundError: print(f"ERROR: 'unzip' command not available.") | |
| except subprocess.CalledProcessError as e: print(f"ERROR: Unzip failed: {e}") | |
| elif not download_successful: print("Download failed previously, cannot unzip.") | |
| elif not os.path.exists(model_zip_file): print(f"Zip file '{model_zip_file}' not found, cannot unzip.") | |
| # --- Final Check and LS --- | |
| print("\nCurrent directory contents:") | |
| os.system("ls -la") | |
| if not models_ready: | |
| print("\n\nERROR: MODEL FILES ARE NOT SET UP CORRECTLY. APP MAY NOT WORK.\n\n") | |
| # --- Gradio Interface --- | |
| def inference(img_input): | |
| if isinstance(img_input, str): img_path = img_input | |
| else: | |
| from PIL import Image | |
| import numpy as np | |
| if isinstance(img_input, np.ndarray): img_pil = Image.fromarray(img_input) | |
| else: img_pil = img_input # Assume PIL | |
| temp_img_name = "temp_gradio_input.png"; img_pil.save(temp_img_name); img_path = temp_img_name | |
| print(f"Saved Gradio input to temporary file: {img_path}") | |
| print(f"Running inference on: {img_path}") | |
| output_sample_path = "./_samples/character_pose_estim.png" | |
| model_checkpoint = "./_train/character_pose_estim/runs/feat_concat+data.ckpt" | |
| if not os.path.exists(img_path): print(f"ERROR: Input path does not exist: {img_path}"); return None | |
| if not os.path.exists(model_checkpoint): print(f"ERROR: Model checkpoint not found: {model_checkpoint}"); return None | |
| command = f"python3 -m _scripts.pose_estimator '{img_path}' '{model_checkpoint}'" | |
| print(f"Executing command: {command}") | |
| try: | |
| result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True, timeout=120) | |
| print("Inference stdout:", result.stdout); print("Inference stderr:", result.stderr); print("Inference completed.") | |
| if os.path.exists(output_sample_path): return output_sample_path | |
| else: print(f"ERROR: Output file not found: {output_sample_path}"); return None | |
| except subprocess.TimeoutExpired: print("ERROR: Inference script timed out."); return None | |
| except subprocess.CalledProcessError as e: print(f"ERROR: Inference script failed: {e.returncode}\nStderr:{e.stderr}\nStdout:{e.stdout}"); return None | |
| finally: | |
| if not isinstance(img_input, str) and 'temp_img_name' in locals() and os.path.exists(temp_img_name): | |
| try: os.remove(temp_img_name); print(f"Removed temporary file: {temp_img_name}") | |
| except OSError as e: print(f"Warning: Failed to remove temp file {temp_img_name}: {e}") | |
| title = "bizarre-pose-estimator" | |
| description = "Gradio demo for Transfer Learning for Pose Estimation of Illustrated Characters. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below." | |
| article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.01819' target='_blank'>Transfer Learning for Pose Estimation of Illustrated Characters</a> | <a href='https://github.com/ShuhongChen/bizarre-pose-estimator' target='_blank'>Github Repo</a></p>" | |
| example_img_file = "IkJzlaE.jpeg" | |
| examples_list = [[example_img_file]] if os.path.exists(example_img_file) else [] | |
| print("Setting up Gradio Interface...") | |
| if models_ready: | |
| # --- USE NEWER GRADIO SYNTAX for inputs/outputs --- | |
| ui = gr.Interface( | |
| fn=inference, | |
| # Define inputs using component classes directly | |
| inputs=gr.Image(type="filepath", label="Input"), # Keep type="filepath" | |
| # Define outputs using component classes directly | |
| outputs=gr.Image(type="filepath", label="Output"), # Keep type="filepath" | |
| # --- End of syntax change --- | |
| title=title, | |
| description=description, | |
| article=article, | |
| allow_flagging="never", | |
| examples=examples_list | |
| ) | |
| print("Launching Gradio...") | |
| ui.launch() | |
| else: | |
| print("Gradio launch aborted because model files are not ready.") |