Update app.py

app.py

@@ -1,15 +1,17 @@
 import os
-import subprocess # Using subprocess for better command execution
-import time # For potential retries or delays if needed
+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
 
-# --- 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
+# 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()
 
-import gradio as gr # Import after potential install
 
 # --- Repository Setup ---
 repo_dir = "bizarre-pose-estimator"
@@ -21,18 +23,21 @@ if not os.path.exists(repo_dir):
         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.")
 
 # --- Change Directory ---
 try:
-    os.chdir(repo_dir)
-    print(f"Changed directory to: {os.getcwd()}")
+    # 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() # Cannot continue without the repo code
+    exit()
 
 # --- Download Example Image ---
 example_img_url = "https://i.imgur.com/IkJzlaE.jpeg"
@@ -47,21 +52,17 @@ if not os.path.exists(example_img_file):
 
 # --- 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
+model_extract_dir = "bizarre_pose_models"
+model_id = "17N5PutpYJTlKuNB6bdDaiQsPSIkYtiPm"
 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
+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.")
-    download_successful = True # Treat as success for logic flow
-    models_ready = True
+    models_ready = True # Models are already copied and ready
 else:
-    print(f"Model checkpoint not found at '{model_checkpoint_rel_path}'. Attempting download...")
-    models_ready = False
+    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
@@ -70,26 +71,21 @@ else:
         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)
+            # 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.")
-            # 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 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 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.")
+            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) # -o overwrite, -q quiet
+            subprocess.run(["unzip", "-oq", model_zip_file], check=True)
             print("Unzip successful.")
 
             print("Copying model files...")
@@ -100,23 +96,16 @@ else:
                 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}")
+                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}")
-                # --- End Zip Deletion ---
-
-            else:
-                print(f"ERROR: Directory '{source_dir_unzip}' not found after unzip.")
-
+                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}")
-    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.")
+    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:")
@@ -126,78 +115,60 @@ if not models_ready:
 
 # --- 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)
+    if isinstance(img_input, str): img_path = img_input
     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
+        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"
 
-    # 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
+    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}'" # Add quotes for paths
+    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) # 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
+        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:
-        # 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}")
+        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>"
-
-# Ensure example file exists before adding it
 examples_list = [[example_img_file]] if os.path.exists(example_img_file) else []
 
-gr.Interface(
-    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", # 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
+print("Setting up Gradio Interface...")
+# Check if models are ready before launching
+if models_ready:
+    # REMOVE enable_queue=True
+    # Also update inputs/outputs to newer syntax if possible, but focus on fixing the error first
+    # Using older gr.inputs/gr.outputs syntax for now as it might be what gradio 2.4.6 expects if it somehow overrides
+    ui = gr.Interface(
+        fn=inference,
+        inputs=gr.inputs.Image(type="filepath", label="Input"), # Using filepath to simplify function logic
+        outputs=gr.outputs.Image(type="file", label="Output"), # Function returns filepath
+        title=title,
+        description=description,
+        article=article,
+        allow_flagging="never",
+        examples=examples_list
+        # enable_queue=True # REMOVED
+    )
+    print("Launching Gradio...")
+    ui.launch() # queue() method is used in newer Gradio if needed, launch() handles basics
+else:
+    print("Gradio launch aborted because model files are not ready.")