Create app.py

app.py

@@ -0,0 +1,64 @@
+import json
+import random
+random.seed(999)
+import torch
+from torchvision.transforms import transforms
+import gradio as gr
+
+model = torch.load('model.pth', map_location=torch.device('cpu'))
+model.eval()
+transform = transforms.Compose([
+    transforms.Resize((384, 384)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+    mean=[
+        0.5,
+        0.5,
+        0.5,
+    ], std=[
+        0.5,
+        0.5,
+        0.5,
+    ])
+])
+
+with open("tags_9940.json", "r") as file:
+    allowed_tags = json.load(file)
+
+allowed_tags = sorted(allowed_tags)
+allowed_tags.append("explicit")
+allowed_tags.append("questionable")
+allowed_tags.append("safe")
+
+def create_tags(image, threshold):
+    img = image.convert('RGB')
+    tensor = transform(img).unsqueeze(0)
+
+    with torch.no_grad():
+        logits = model(tensor)
+        probabilities = torch.nn.functional.sigmoid(logits[0])
+        indices = torch.where(probabilities > threshold)[0]
+        values = probabilities[indices]
+
+    temp = []
+    tag_score = dict()
+    for i in range(indices.size(0)):
+        temp.append([allowed_tags[indices[i]], values[i].item()])
+        tag_score[allowed_tags[indices[i]]] = values[i].item()
+    # temp = sorted(temp, key=lambda x: x[1], reverse=True)
+    # print("Before adding implicated tags, there are " + str(len(temp)) + " tags")
+    temp = [t[0] for t in temp]
+    text_no_impl = " ".join(temp)
+    return text_no_impl, tag_score
+
+demo = gr.Interface(
+    create_tags,
+    inputs=[gr.Image(label="Source", sources=['upload', 'webcam'], type='pil'), gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.30, label="Threshold")],
+    outputs=[
+        gr.Textbox(label="Tag String"),
+        gr.Label(label="Tag Predictions", num_top_classes=200),
+    ],
+    allow_flagging="never",
+)
+
+demo.launch()