Skip to content

Instantly share code, notes, and snippets.

@l4rz
Created January 22, 2021 11:28
Show Gist options
  • Save l4rz/32315d435a662cb8e15febf888139eb7 to your computer and use it in GitHub Desktop.
Save l4rz/32315d435a662cb8e15febf888139eb7 to your computer and use it in GitHub Desktop.
import torch
import clip
from PIL import Image
from pathlib import Path
import random, sys, string
import shutil
import os
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)
filelist = sorted(Path(sys.argv[1]).rglob('*.jpg'))
for file in filelist:
file = str(file)
image = preprocess(Image.open(file)).unsqueeze(0).to(device)
text = clip.tokenize(["a girl in lingerie", "a naked girl"]).to(device)
with torch.no_grad():
image_features = model.encode_image(image)
text_features = model.encode_text(text)
logits_per_image, logits_per_text = model(image, text3)
probs = logits_per_image.softmax(dim=-1).cpu().numpy()
print(file, "lingerie/naked probs:", probs)
if probs[0][0] > 0.9:
shutil.copyfile(file, "lingerie/"+str(probs3[0][1])+"-"+os.path.basename(file))
if probs[0][1] > 0.9:
shutil.copyfile(file, "naked/"+str(probs3[0][1])+"-"+os.path.basename(file))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment