File size: 5,565 Bytes
b1d950f 6ac1aaf d9fca73 6ac1aaf 37bff96 17dc0c3 37bff96 788fe60 31e248b b1f936e 1421531 b1f936e 788fe60 37bff96 d9fca73 23fe6a2 aaa673b 69fa5d2 aaa673b 17dc0c3 69fa5d2 aaa673b 37bff96 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
import os
import subprocess
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as T
import gradio as gr
from pytorch_lightning import seed_everything
import os
import requests
import csv
import spaces
from triton.fb import build_paths
def plot_feats(image, lr, hr):
from featup.util import pca, remove_axes
assert len(image.shape) == len(lr.shape) == len(hr.shape) == 3
seed_everything(0)
[lr_feats_pca, hr_feats_pca], _ = pca(
[lr.unsqueeze(0), hr.unsqueeze(0)], dim=9)
fig, ax = plt.subplots(3, 3, figsize=(15, 15))
ax[0, 0].imshow(image.permute(1, 2, 0).detach().cpu())
ax[1, 0].imshow(image.permute(1, 2, 0).detach().cpu())
ax[2, 0].imshow(image.permute(1, 2, 0).detach().cpu())
ax[0, 0].set_title("Image", fontsize=22)
ax[0, 1].set_title("Original", fontsize=22)
ax[0, 2].set_title("Upsampled Features", fontsize=22)
ax[0, 1].imshow(lr_feats_pca[0, :3].permute(1, 2, 0).detach().cpu())
ax[0, 0].set_ylabel("PCA Components 1-3", fontsize=22)
ax[0, 2].imshow(hr_feats_pca[0, :3].permute(1, 2, 0).detach().cpu())
ax[1, 1].imshow(lr_feats_pca[0, 3:6].permute(1, 2, 0).detach().cpu())
ax[1, 0].set_ylabel("PCA Components 4-6", fontsize=22)
ax[1, 2].imshow(hr_feats_pca[0, 3:6].permute(1, 2, 0).detach().cpu())
ax[2, 1].imshow(lr_feats_pca[0, 6:9].permute(1, 2, 0).detach().cpu())
ax[2, 0].set_ylabel("PCA Components 7-9", fontsize=22)
ax[2, 2].imshow(hr_feats_pca[0, 6:9].permute(1, 2, 0).detach().cpu())
remove_axes(ax)
plt.tight_layout()
plt.close(fig) # Close plt to avoid additional empty plots
return fig
def download_image(url, save_path):
response = requests.get(url)
with open(save_path, 'wb') as file:
file.write(response.content)
base_url = "https://marhamilresearch4.blob.core.windows.net/feature-upsampling-public/sample_images/"
sample_images_urls = {
"skate.jpg": base_url + "skate.jpg",
"car.jpg": base_url + "car.jpg",
"plant.png": base_url + "plant.png",
}
sample_images_dir = "/tmp/sample_images"
# Ensure the directory for sample images exists
os.makedirs(sample_images_dir, exist_ok=True)
# Download each sample image
for filename, url in sample_images_urls.items():
save_path = os.path.join(sample_images_dir, filename)
# Download the image if it doesn't already exist
if not os.path.exists(save_path):
print(f"Downloading {filename}...")
download_image(url, save_path)
else:
print(f"{filename} already exists. Skipping download.")
os.environ['TORCH_HOME'] = '/tmp/.cache'
os.environ['GRADIO_EXAMPLES_CACHE'] = '/tmp/gradio_cache'
csv.field_size_limit(100000000)
options = ['dino16', 'vit', 'dinov2', 'clip', 'resnet50']
image_input = gr.Image(label="Choose an image to featurize",
height=480,
type="pil",
image_mode='RGB',
sources=['upload', 'webcam', 'clipboard']
)
model_option = gr.Radio(options, value="dino16",
label='Choose a backbone to upsample')
def find_cuda_home():
try:
# Define the search string and the directory
search_string = "CUDA"
search_directory = "/usr"
# Use subprocess to run the grep command
command = ['grep', '-r', search_string, search_directory]
output = subprocess.check_output(command).decode()
print(output)
for line in output.split('\n'):
if 'Cuda compilation tools' in line:
version = line.split()[-1]
return f"/usr/local/cuda-{version.split('.')[0]}.{version.split('.')[1]}"
except Exception as e:
print(f"Error finding CUDA_HOME: {e}")
return None
@spaces.GPU
def upsample_features(image, model_option):
os.environ["CUDA_HOME"] = build_paths.cuda()
print(os.environ["CUDA_HOME"])
os.environ['PATH'] = '/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'
os.environ['LD_LIBRARY_PATH'] = '/usr/local/nvidia/lib:/usr/local/nvidia/lib64'
# Install the required package from GitHub
subprocess.check_call(
["pip", "install", "git+https://github.com/mhamilton723/FeatUp"])
from featup.util import norm, unnorm
models = {o: torch.hub.load("mhamilton723/FeatUp", o) for o in options}
# Image preprocessing
input_size = 224
transform = T.Compose([
T.Resize(input_size),
T.CenterCrop((input_size, input_size)),
T.ToTensor(),
norm
])
image_tensor = transform(image).unsqueeze(0).cuda()
# Load the selected model
upsampler = models[model_option].cuda()
hr_feats = upsampler(image_tensor)
lr_feats = upsampler.model(image_tensor)
upsampler.cpu()
return plot_feats(unnorm(image_tensor)[0], lr_feats[0], hr_feats[0])
demo = gr.Interface(fn=upsample_features,
inputs=[image_input, model_option],
outputs="plot",
title="Feature Upsampling Demo",
description="This demo allows you to upsample features of an image using selected models.",
examples=[
["/tmp/sample_images/skate.jpg", "dino16"],
["/tmp/sample_images/car.jpg", "dinov2"],
["/tmp/sample_images/plant.png", "dino16"],
]
)
demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)
|