94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
import os
|
|
import sys
|
|
import urllib.request
|
|
import zipfile
|
|
import subprocess
|
|
from tqdm import tqdm
|
|
|
|
# URL to download the checkpoints (colorization_checkpoint.zip)
|
|
CHECKPOINT_URL = (
|
|
"https://github.com/zhangmozhe/Deep-Exemplar-based-Video-Colorization/"
|
|
"releases/download/v1.0/colorization_checkpoint.zip"
|
|
)
|
|
ZIP_FILENAME = "colorization_checkpoint.zip"
|
|
|
|
def ensure_opencv_contrib():
|
|
"""
|
|
Checks if cv2.ximgproc is available. If not, installs opencv-contrib-python.
|
|
"""
|
|
try:
|
|
import cv2
|
|
_ = cv2.ximgproc # just a test import
|
|
print("[install.py] opencv-contrib-python is already installed.")
|
|
except (ImportError, AttributeError):
|
|
print("[install.py] Installing opencv-contrib-python for WLS filter support...")
|
|
subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "opencv-contrib-python"])
|
|
print("[install.py] opencv-contrib-python installed successfully.")
|
|
|
|
def download_file(url: str, dest_path: str) -> None:
|
|
"""
|
|
Downloads a file from a URL to a local destination with a progress bar.
|
|
|
|
Args:
|
|
url (str): The URL from which to download the file.
|
|
dest_path (str): The file path to save the downloaded file.
|
|
"""
|
|
with tqdm(unit='B', unit_scale=True, miniters=1, desc=os.path.basename(dest_path)) as t:
|
|
def reporthook(block_num, block_size, total_size):
|
|
if total_size > 0:
|
|
t.total = total_size
|
|
t.update(block_size)
|
|
urllib.request.urlretrieve(url, dest_path, reporthook=reporthook)
|
|
|
|
def extract_zip(zip_path: str, extract_to: str) -> None:
|
|
"""
|
|
Extracts a zip file into the specified directory.
|
|
|
|
Args:
|
|
zip_path (str): The path to the zip file.
|
|
extract_to (str): The directory where files will be extracted.
|
|
"""
|
|
with zipfile.ZipFile(zip_path, "r") as zf:
|
|
zf.extractall(extract_to)
|
|
|
|
def main() -> None:
|
|
"""
|
|
Main function that downloads and extracts the checkpoint files.
|
|
This will create the necessary `checkpoints/` and `data/` directories.
|
|
"""
|
|
# 1) Ensure opencv-contrib-python is present (for cv2.ximgproc)
|
|
ensure_opencv_contrib()
|
|
|
|
# 2) Determine the root directory (where this script resides)
|
|
root_dir = os.path.dirname(os.path.abspath(__file__))
|
|
zip_path = os.path.join(root_dir, ZIP_FILENAME)
|
|
|
|
# 3) Download the checkpoint zip if it's missing
|
|
if not os.path.exists(zip_path):
|
|
print(f"Downloading checkpoints from {CHECKPOINT_URL} ...")
|
|
try:
|
|
download_file(CHECKPOINT_URL, zip_path)
|
|
print("Download complete.")
|
|
except Exception as e:
|
|
print(f"Error during download: {e}")
|
|
return
|
|
else:
|
|
print(f"Zip file already exists: {zip_path}")
|
|
|
|
# 4) Extract the downloaded zip file
|
|
print(f"Extracting {zip_path} into {root_dir} ...")
|
|
try:
|
|
extract_zip(zip_path, root_dir)
|
|
print("Extraction complete.")
|
|
except Exception as e:
|
|
print(f"Error during extraction: {e}")
|
|
return
|
|
|
|
# Optionally remove the zip file after successful extraction
|
|
# os.remove(zip_path)
|
|
|
|
print("Checkpoint files are installed in the `checkpoints/` and `data/` directories.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|