File size: 1,647 Bytes
aebb6ba 370ff13 aebb6ba 370ff13 aebb6ba 370ff13 aebb6ba 370ff13 aebb6ba 370ff13 aebb6ba 370ff13 aebb6ba 370ff13 aebb6ba 370ff13 aebb6ba 370ff13 aebb6ba 370ff13 aebb6ba 370ff13 |
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 |
import os
import time
import glob
import numpy as np
from PIL import Image, ImageDraw
import mlx.core as mx
from mlxDeepDanBooru.mlx_deep_danbooru_model import mlxDeepDanBooruModel
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed, wait, FIRST_COMPLETED
from copy import deepcopy
ROOTDIR = os.path.dirname(os.path.abspath(__file__))
IMAGEDIR = f'{ROOTDIR}/example'
model_path = f"{ROOTDIR}/models/model-resnet_custom_v3_mlx.npz"
tags_path = f'{ROOTDIR}/models/tags-resnet_custom_v3_mlx.npy'
mlx_dan = mlxDeepDanBooruModel()
mlx_dan.load_weights(model_path)
mx.eval(mlx_dan.parameters())
model_tags = np.load(tags_path)
#print(f'total tags: {len(model_tags)}')
def danbooru_tags(fpath):
results = {}
tags = []
pic = Image.open(fpath).convert("RGB").resize((512, 512))
a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
x = mx.array(a)
y = mlx_dan(x)[0]
try:
for n in range(10):
mlx_dan(x)
for i, p in enumerate(y):
if p >= 0.55:
#print(model_tags[i].item(), p)
tags.append(model_tags[i].item())
except Exception as err:
print(err)
results[fpath] = tags
return results
def image_infer(fpath):
tags = danbooru_tags(fpath)
return tags
t1 = time.time()
tags_1 = image_infer(f'{IMAGEDIR}/1.png')
tags_2 = image_infer(f'{IMAGEDIR}/2.png')
t2 = time.time()
print(tags_1)
print(tags_2)
print(f'2 images: infer speed(with mlx): {(t2 - t1)/2} seconds per image')
|