|
import json |
|
import os |
|
import zipfile |
|
from io import BytesIO |
|
from tempfile import NamedTemporaryFile |
|
import tempfile |
|
|
|
import gradio as gr |
|
import pandas as pd |
|
from PIL import Image |
|
import safetensors.torch |
|
import spaces |
|
import timm |
|
from timm.models import VisionTransformer |
|
import torch |
|
from torchvision.transforms import transforms |
|
from torchvision.transforms import InterpolationMode |
|
import torchvision.transforms.functional as TF |
|
from torch.utils.data import Dataset, DataLoader |
|
|
|
|
|
torch.set_grad_enabled(False) |
|
|
|
class Fit(torch.nn.Module): |
|
def __init__( |
|
self, |
|
bounds: tuple[int, int] | int, |
|
interpolation = InterpolationMode.LANCZOS, |
|
grow: bool = True, |
|
pad: float | None = None |
|
): |
|
super().__init__() |
|
|
|
self.bounds = (bounds, bounds) if isinstance(bounds, int) else bounds |
|
self.interpolation = interpolation |
|
self.grow = grow |
|
self.pad = pad |
|
|
|
def forward(self, img: Image) -> Image: |
|
wimg, himg = img.size |
|
hbound, wbound = self.bounds |
|
|
|
hscale = hbound / himg |
|
wscale = wbound / wimg |
|
|
|
if not self.grow: |
|
hscale = min(hscale, 1.0) |
|
wscale = min(wscale, 1.0) |
|
|
|
scale = min(hscale, wscale) |
|
if scale == 1.0: |
|
return img |
|
|
|
hnew = min(round(himg * scale), hbound) |
|
wnew = min(round(wimg * scale), wbound) |
|
|
|
img = TF.resize(img, (hnew, wnew), self.interpolation) |
|
|
|
if self.pad is None: |
|
return img |
|
|
|
hpad = hbound - hnew |
|
wpad = wbound - wnew |
|
|
|
tpad = hpad // 2 |
|
bpad = hpad - tpad |
|
|
|
lpad = wpad // 2 |
|
rpad = wpad - lpad |
|
|
|
return TF.pad(img, (lpad, tpad, rpad, bpad), self.pad) |
|
|
|
def __repr__(self) -> str: |
|
return ( |
|
f"{self.__class__.__name__}(" + |
|
f"bounds={self.bounds}, " + |
|
f"interpolation={self.interpolation.value}, " + |
|
f"grow={self.grow}, " + |
|
f"pad={self.pad})" |
|
) |
|
|
|
class CompositeAlpha(torch.nn.Module): |
|
def __init__( |
|
self, |
|
background: tuple[float, float, float] | float, |
|
): |
|
super().__init__() |
|
|
|
self.background = (background, background, background) if isinstance(background, float) else background |
|
self.background = torch.tensor(self.background).unsqueeze(1).unsqueeze(2) |
|
|
|
def forward(self, img: torch.Tensor) -> torch.Tensor: |
|
if img.shape[-3] == 3: |
|
return img |
|
|
|
alpha = img[..., 3, None, :, :] |
|
|
|
img[..., :3, :, :] *= alpha |
|
|
|
background = self.background.expand(-1, img.shape[-2], img.shape[-1]) |
|
if background.ndim == 1: |
|
background = background[:, None, None] |
|
elif background.ndim == 2: |
|
background = background[None, :, :] |
|
|
|
img[..., :3, :, :] += (1.0 - alpha) * background |
|
return img[..., :3, :, :] |
|
|
|
def __repr__(self) -> str: |
|
return ( |
|
f"{self.__class__.__name__}(" + |
|
f"background={self.background})" |
|
) |
|
|
|
transform = transforms.Compose([ |
|
Fit((384, 384)), |
|
transforms.ToTensor(), |
|
CompositeAlpha(0.5), |
|
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), |
|
transforms.CenterCrop((384, 384)), |
|
]) |
|
|
|
model = timm.create_model( |
|
"vit_so400m_patch14_siglip_384.webli", |
|
pretrained=False, |
|
num_classes=9083, |
|
) |
|
|
|
safetensors.torch.load_model(model, "JTP_PILOT-e4-vit_so400m_patch14_siglip_384.safetensors") |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
model.to(device) |
|
model.eval() |
|
|
|
with open("tagger_tags.json", "r") as file: |
|
tags = json.load(file) |
|
allowed_tags = list(tags.keys()) |
|
|
|
for idx, tag in enumerate(allowed_tags): |
|
allowed_tags[idx] = tag.replace("_", " ") |
|
|
|
sorted_tag_score = {} |
|
|
|
@spaces.GPU(duration=9) |
|
def run_classifier(image, threshold): |
|
global sorted_tag_score |
|
img = image.convert('RGB') |
|
tensor = transform(img).unsqueeze(0).to(device) |
|
with torch.no_grad(): |
|
logits = model(tensor) |
|
probabilities = torch.nn.functional.sigmoid(logits[0]) |
|
indices = torch.topk(probabilities, 250).indices |
|
values = probabilities[indices] |
|
|
|
tag_score = dict() |
|
for i in range(indices.size(0)): |
|
tag_score[allowed_tags[indices[i]]] = values[i].item() |
|
sorted_tag_score = dict(sorted(tag_score.items(), key=lambda item: item[1], reverse=True)) |
|
|
|
return create_tags(threshold) |
|
|
|
def create_tags(threshold): |
|
global sorted_tag_score |
|
filtered_tag_score = {key: value for key, value in sorted_tag_score.items() if value > threshold} |
|
text_no_impl = ", ".join(filtered_tag_score.keys()) |
|
return text_no_impl, filtered_tag_score |
|
|
|
|
|
class ImageDataset(Dataset): |
|
def __init__(self, image_files, transform): |
|
self.image_files = image_files |
|
self.transform = transform |
|
|
|
def __len__(self): |
|
return len(self.image_files) |
|
|
|
def __getitem__(self, idx): |
|
img_path = self.image_files[idx] |
|
img = Image.open(img_path).convert('RGB') |
|
return self.transform(img), os.path.basename(img_path) |
|
|
|
@spaces.GPU(duration=299) |
|
def process_images(images, threshold): |
|
dataset = ImageDataset(images, transform) |
|
|
|
dataloader = DataLoader(dataset, batch_size=64, num_workers=0, pin_memory=True, drop_last=False) |
|
|
|
all_results = [] |
|
|
|
with torch.no_grad(): |
|
for batch, filenames in dataloader: |
|
|
|
batch = batch.to(device) |
|
with torch.no_grad(): |
|
logits = model(batch) |
|
probabilities = torch.nn.functional.sigmoid(logits) |
|
|
|
for i, prob in enumerate(probabilities): |
|
indices = torch.where(prob > threshold)[0] |
|
values = prob[indices] |
|
|
|
temp = [] |
|
tag_score = dict() |
|
for j in range(indices.size(0)): |
|
temp.append([allowed_tags[indices[j]], values[j].item()]) |
|
tag_score[allowed_tags[indices[j]]] = values[j].item() |
|
|
|
tags = ", ".join([t[0] for t in temp]) |
|
all_results.append((filenames[i], tags, tag_score)) |
|
|
|
return all_results |
|
|
|
def is_valid_image(file_path): |
|
try: |
|
with Image.open(file_path) as img: |
|
img.verify() |
|
return True |
|
except: |
|
return False |
|
|
|
def process_zip(zip_file, threshold): |
|
if zip_file is None: |
|
return None, None |
|
|
|
with tempfile.TemporaryDirectory() as temp_dir: |
|
with zipfile.ZipFile(zip_file.name, 'r') as zip_ref: |
|
zip_ref.extractall(temp_dir) |
|
|
|
all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)] |
|
image_files = [f for f in all_files if is_valid_image(f)] |
|
results = process_images(image_files, threshold) |
|
|
|
temp_file = NamedTemporaryFile(delete=False, suffix=".zip") |
|
with zipfile.ZipFile(temp_file, "w") as zip_ref: |
|
for image_name, text_no_impl, _ in results: |
|
with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file: |
|
file.write(text_no_impl.encode()) |
|
temp_file.seek(0) |
|
df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags']) |
|
|
|
return temp_file.name, df |
|
|
|
@spaces.GPU(duration=120) |
|
def process_images_light(images, threshold): |
|
dataset = ImageDataset(images, transform) |
|
|
|
dataloader = DataLoader(dataset, batch_size=32, num_workers=0, pin_memory=True, drop_last=False) |
|
|
|
all_results = [] |
|
|
|
with torch.no_grad(): |
|
for batch, filenames in dataloader: |
|
|
|
batch = batch.to(device) |
|
with torch.no_grad(): |
|
logits = model(batch) |
|
probabilities = torch.nn.functional.sigmoid(logits) |
|
|
|
for i, prob in enumerate(probabilities): |
|
indices = torch.where(prob > threshold)[0] |
|
values = prob[indices] |
|
|
|
temp = [] |
|
tag_score = dict() |
|
for j in range(indices.size(0)): |
|
temp.append([allowed_tags[indices[j]], values[j].item()]) |
|
tag_score[allowed_tags[indices[j]]] = values[j].item() |
|
|
|
tags = ", ".join([t[0] for t in temp]) |
|
all_results.append((filenames[i], tags, tag_score)) |
|
|
|
return all_results |
|
|
|
def process_zip_light(zip_file, threshold): |
|
if zip_file is None: |
|
return None, None |
|
|
|
with tempfile.TemporaryDirectory() as temp_dir: |
|
with zipfile.ZipFile(zip_file.name, 'r') as zip_ref: |
|
zip_ref.extractall(temp_dir) |
|
|
|
all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)] |
|
image_files = [f for f in all_files if is_valid_image(f)] |
|
results = process_images_light(image_files, threshold) |
|
|
|
temp_file = NamedTemporaryFile(delete=False, suffix=".zip") |
|
with zipfile.ZipFile(temp_file, "w") as zip_ref: |
|
for image_name, text_no_impl, _ in results: |
|
with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file: |
|
file.write(text_no_impl.encode()) |
|
temp_file.seek(0) |
|
df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags']) |
|
|
|
return temp_file.name, df |
|
|
|
with gr.Blocks(css=".output-class { display: none; }") as demo: |
|
gr.Markdown(""" |
|
## Joint Tagger Project: PILOT Demo |
|
This tagger is designed for use on furry images (though may very well work on out-of-distribution images, potentially with funny results). A threshold of 0.2 is recommended. Lower thresholds often turn up more valid tags, but can also result in some amount of hallucinated tags. |
|
|
|
This tagger is the result of joint efforts between members of the RedRocket team. Special thanks to Minotoro at frosting.ai for providing the compute power for this project. |
|
|
|
Usage Note for batch tagging: |
|
|
|
the normal version is limited to 300s and uses batch size 64 |
|
|
|
the light version is limited to 120s with batch size 32 |
|
|
|
if your image count is low use the light version for lower gpu wait time (most of the time you instantly get a gpu anyway) |
|
""") |
|
|
|
with gr.Tabs(): |
|
with gr.TabItem("Single Image"): |
|
with gr.Row(): |
|
with gr.Column(): |
|
image_input = gr.Image(label="Source", sources=['upload'], type='pil', height=512, show_label=False) |
|
threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold") |
|
with gr.Column(): |
|
tag_string = gr.Textbox(label="Tag String") |
|
label_box = gr.Label(label="Tag Predictions", num_top_classes=250, show_label=False) |
|
|
|
image_input.upload( |
|
fn=run_classifier, |
|
inputs=[image_input, threshold_slider], |
|
outputs=[tag_string, label_box] |
|
) |
|
|
|
threshold_slider.input( |
|
fn=create_tags, |
|
inputs=[threshold_slider], |
|
outputs=[tag_string, label_box] |
|
) |
|
|
|
with gr.TabItem("Multiple Images"): |
|
with gr.Row(): |
|
with gr.Column(): |
|
zip_input = gr.File(label="Upload ZIP file", file_types=['.zip']) |
|
multi_threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold") |
|
process_button = gr.Button("Process Images") |
|
with gr.Column(): |
|
zip_output = gr.File(label="Download Tagged Text Files (ZIP)") |
|
dataframe_output = gr.Dataframe(label="Image Tags Summary") |
|
|
|
process_button.click( |
|
fn=process_zip, |
|
inputs=[zip_input, multi_threshold_slider], |
|
outputs=[zip_output, dataframe_output] |
|
) |
|
with gr.TabItem("Multiple Images (Light)"): |
|
with gr.Row(): |
|
with gr.Column(): |
|
zip_input_light = gr.File(label="Upload ZIP file", file_types=['.zip']) |
|
multi_threshold_slider_light = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold") |
|
process_button_light = gr.Button("Process Images (Light)") |
|
with gr.Column(): |
|
zip_output_light = gr.File(label="Download Tagged Text Files (ZIP)") |
|
dataframe_output_light = gr.Dataframe(label="Image Tags Summary") |
|
|
|
process_button_light.click( |
|
fn=process_zip_light, |
|
inputs=[zip_input_light, multi_threshold_slider_light], |
|
outputs=[zip_output_light, dataframe_output_light] |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.queue().launch() |
