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 = "

Transfer Learning for Pose Estimation of Illustrated Characters | Github Repo

" 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.")