| import os
|
| import time
|
| import math
|
| import argparse
|
| import numpy as np
|
| from PIL import Image, ImageDraw
|
|
|
| import mlx.core as mx
|
| from mlxDeepDanBooru import *
|
|
|
| parser = argparse.ArgumentParser(description='Thanks for choosing mlxDeepDanBooruModel.')
|
| parser.add_argument("path", type=str, help='image path or image-dir path')
|
| args = parser.parse_args()
|
|
|
| IMAGEDIR = f'example'
|
|
|
| in_path = args.path
|
| user_file = ""
|
| if os.path.isdir(in_path):
|
| IMAGEDIR = in_path
|
| user_file = ""
|
|
|
| if os.path.isfile(in_path):
|
| user_file = in_path
|
| IMAGEDIR = ""
|
|
|
| print(f"user_file: {user_file}")
|
| print(f"IMAGEDIR: {IMAGEDIR}")
|
|
|
| mlx_dan = mlxDeepDanBooruModel()
|
|
|
| mlx_dan.load_weights(f"models/model-resnet_custom_v3_fp32_mlx.safetensors")
|
|
|
|
|
|
|
| mx.eval(mlx_dan.parameters())
|
|
|
|
|
| model_tags = np.load(f'models/tags-resnet_custom_v3_mlx.npy')
|
|
|
|
|
|
|
| def danbooru_tags(fpath):
|
| results = {}
|
| tags = []
|
|
|
| try:
|
| pic = img_resize(fpath)
|
| a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
|
|
|
|
|
|
|
|
|
| x = mx.array(a)
|
| y = mlx_dan(x)[0]
|
| ylst = y.tolist()
|
| ylen = len(ylst)
|
| for i in range(ylen):
|
| if ylst[i] >= 0.72:
|
| tags.append(model_tags[i].item())
|
|
|
| except Exception as err:
|
| print(err)
|
| tags = []
|
|
|
| results[fpath] = tags
|
| return results
|
|
|
| def image_infer(fpath):
|
| tags = danbooru_tags(fpath)
|
| return tags
|
|
|
|
|
| if __name__ == "__main__":
|
| image_count = 0
|
| t1 = time.time()
|
|
|
| if user_file != "":
|
| tags = image_infer(user_file)
|
| image_count += 1
|
| print(tags)
|
|
|
| if IMAGEDIR != "":
|
| for root, _, files in os.walk(IMAGEDIR,True):
|
| for f in files:
|
| if not f[-4:].lower() in (".png",".jpg","jpeg",".bmp"):
|
| continue
|
| fpath = os.path.join(root,f).replace("\\","/")
|
| tags = image_infer(fpath)
|
| image_count += 1
|
| print(tags)
|
|
|
| t2 = time.time()
|
|
|
| if image_count > 0:
|
| print(f'\n\n{image_count} images: infer speed(with mlx): {(t2 - t1)/image_count} seconds per image')
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
