Spaces:
Running
Running
File size: 1,533 Bytes
59d7fdc 58d222b 2fb5d7b 3476774 59d7fdc 952a281 59d7fdc 952a281 59d7fdc 13908e1 952a281 59d7fdc 952a281 59d7fdc 433a10b | 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 | import torch
torch.set_grad_enabled(False)
import gradio as gr
from datetime import datetime
from inference_tagger_standalone import *
# from huggingface_hub import hf_hub_download
# hf_hub_download(repo_id="lodestones/tagger-experiment", filename="tagger_proto.safetensors", local_dir=".")
import os
os.system("wget -nv https://huggingface.co/lodestones/tagger-experiment/resolve/main/tagger_proto.safetensors")
model = Tagger(checkpoint_path="./tagger_proto.safetensors", vocab_path="./tagger_vocab.json", max_size=1024)
def get_tags(image, top_k):
current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
print(f"{current_datetime}: started.")
results = model.predict(image, topk=top_k,)#threshold=threshold)
temp = []
return_dict = dict()
for rank, (tag, score) in enumerate(results, 1):
return_dict[tag] = score
temp.append(tag.replace(" ", "_"))
return_str = " ".join(temp)
current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
print(f"{current_datetime}: finished.\n")
return return_str, return_dict
demo = gr.Interface(
get_tags,
inputs=[
gr.Image(label="Source", sources=['upload',], type='filepath'),
# gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.30, label="Threshold"),
gr.Slider(minimum=0, maximum=500, step=1, value=85, label="Top K")
],
outputs=[
gr.Textbox(label="Tag String"),
gr.Label(label="Tag Predictions", num_top_classes=500),
],
)
demo.launch(ssr_mode=False)
|