import os
import subprocess

# Clone the repository
if not os.path.exists("wdv3-jax"):
    repo_url = "https://github.com/SmilingWolf/wdv3-jax.git"
    subprocess.run(["git", "clone", repo_url])

# Change directory to the cloned repository
os.chdir("wdv3-jax")

subprocess.run(["bash", "setup.sh"])

model_type = "swinv2"  # or "convnext" or "vit"
image_path = "../test-img.png"
with open("output.txt", "w") as output_file:
    subprocess.run(["python", "wdv3_jax.py", model_type, image_path], stdout=output_file)