Update app.py

app.py

@@ -1,42 +1,203 @@
 import os
-os.system("pip install gradio==2.4.6")
-import gradio as gr
-os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'")
-os.system("git clone https://github.com/ShuhongChen/bizarre-pose-estimator.git")
-os.chdir("bizarre-pose-estimator")
+import subprocess # Using subprocess for better command execution
+import time # For potential retries or delays if needed
 
-os.system("wget https://i.imgur.com/IkJzlaE.jpeg")
+# --- Dependency Installation ---
+# Note: Ideally, manage dependencies via requirements.txt for HF Spaces build process.
+# Running pip here might be slow on startup or cause unexpected version conflicts.
+print("Ensuring dependencies are installed...")
+# os.system("pip install gradio==2.4.6") # Prefer version from requirements.txt
+# os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'") # Prefer from requirements.txt
 
-os.system("gdown https://drive.google.com/uc?id=1HyOKbl2iLYNFVN3FuL3ENSe_L47Cxk69")
+import gradio as gr # Import after potential install
 
+# --- 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}")
+        # Consider exiting if clone fails
+        exit()
+else:
+    print(f"Repository '{repo_dir}' already exists.")
 
-os.system("unzip bizarre_pose_models.zip")
-os.system("cp -a ./bizarre_pose_models/. .")
+# --- Change Directory ---
+try:
+    os.chdir(repo_dir)
+    print(f"Changed directory to: {os.getcwd()}")
+except FileNotFoundError:
+    print(f"ERROR: Failed to change directory to '{repo_dir}'.")
+    exit() # Cannot continue without the repo code
 
+# --- 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}")
 
-os.system("ls")
+# --- Model Download ---
+model_zip_file = "bizarre_pose_models.zip"
+model_extract_dir = "bizarre_pose_models" # Expected folder name inside zip
+model_id = "17N5PutpYJTlKuNB6bdDaiQsPSIkYtiPm" # The correct FILE ID
+download_successful = False
+
+# Check if model files already exist (e.g., from previous run/unzip)
+# Check for a key file/folder that should exist after copy
+# Using the checkpoint file as an indicator
+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.")
+    download_successful = True # Treat as success for logic flow
+    models_ready = True
+else:
+    print(f"Model checkpoint not found at '{model_checkpoint_rel_path}'. Attempting download...")
+    models_ready = False
+    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:
+            result = subprocess.run(gdown_command, shell=True, check=True, capture_output=True, text=True)
+            print("Gdown download successful.")
+            # print("stdout:", result.stdout) # Often too verbose
+            # print("stderr:", result.stderr)
+            download_successful = True
+        except FileNotFoundError:
+             print("ERROR: 'gdown' command not found. Install via requirements.txt.")
+        except subprocess.CalledProcessError as e:
+            print(f"ERROR: gdown download failed with exit code {e.returncode}")
+            print("Stderr:", e.stderr)
+            print("Stdout:", e.stdout) # May contain useful info from gdown
+            print("Download might fail due to quotas or permissions on Google Drive.")
+            print("Consider manually uploading 'bizarre_pose_models.zip' to the Space root if errors persist.")
+
+    # --- 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) # -o overwrite, -q quiet
+            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
+
+                # --- Delete Zip File ---
+                print(f"Attempting to remove zip file to save space: {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}")
+                # --- End Zip Deletion ---
+
+            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}")
+    else:
+        if 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):
+    # Gradio passes a PIL Image or Numpy array depending on config, need filepath
+    # Save the uploaded image temporarily if it's not already a path
+    if isinstance(img_input, str):
+        img_path = img_input # Already a path (e.g., from examples)
+    else:
+        # Save PIL/Numpy input to a temporary file
+        from PIL import Image
+        import numpy as np
+        if isinstance(img_input, np.ndarray):
+            img_pil = Image.fromarray(img_input)
+        else: # Assume PIL Image
+            img_pil = img_input
+        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"
+
+    # Check paths relative to current directory ('bizarre-pose-estimator')
+    if not os.path.exists(img_path):
+      print(f"ERROR: Input image path does not exist: {img_path}"); return None
+    if not os.path.exists(model_checkpoint):
+       print(f"ERROR: Model checkpoint not found at: {model_checkpoint}"); return None
+
+    command = f"python3 -m _scripts.pose_estimator '{img_path}' '{model_checkpoint}'" # Add quotes for paths
+    print(f"Executing command: {command}")
+    try:
+        result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True, timeout=120) # Add timeout
+        print("Inference stdout:", result.stdout)
+        print("Inference stderr:", result.stderr)
+        print("Inference completed.")
+        if os.path.exists(output_sample_path):
+            return output_sample_path # Return path to the output image
+        else:
+            print(f"ERROR: Inference ran but 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:
+        # Clean up temporary input file if created
+        if not isinstance(img_input, str) 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}")
 
-def inference(img):
-  os.system("python3 -m _scripts.pose_estimator "+img+" ./_train/character_pose_estim/runs/feat_concat+data.ckpt")
 
-  return "./_samples/character_pose_estim.png"
-  
-  
 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>"
 
-examples=[["IkJzlaE.jpeg"]]
+# Ensure example file exists before adding it
+examples_list = [[example_img_file]] if os.path.exists(example_img_file) else []
+
 gr.Interface(
-    inference, 
-    gr.inputs.Image(type="filepath", label="Input"), 
-    gr.outputs.Image(type="file", label="Output"),
+    fn=inference, # Use fn= alias
+    # Input needs to accept various types, filepath ensures it works with examples
+    inputs=gr.Image(type="pil", label="Input"), # Use PIL input type, save internally
+    outputs=gr.Image(type="filepath", label="Output"), # Output is path to generated file
     title=title,
     description=description,
     article=article,
-    allow_flagging="never",
-    examples=examples,
-    enable_queue=True
-    ).launch()
-
+    allow_flagging="never", # Original setting
+    examples=examples_list,
+    enable_queue=True # Original setting
+    # Consider adding cache_examples=True if examples are static
+).launch() # share=True is needed only if not running on HF platform