-

README.md

@@ -57,31 +57,27 @@ Image Interrogate:
 
 ```python
 import numpy as np
-from PIL import Image, ImageDraw
+from PIL import Image
 
 # using apple silicon's MLX 
 # not Pytorch
 import mlx.core as mx
 from mlxDeepDanBooru.mlx_deep_danbooru_model import mlxDeepDanBooruModel
 
-
 model_path = "models/model-resnet_custom_v3_fp32_mlx.safetensors"
-tags_path = 'models/tags-resnet_custom_v3_mlx.npy'
+model_tags = np.load('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):
     tags = []
-    pic = Image.open(fpath).convert("RGB").resize((512, 512))
-    a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
 
     try:
+        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]
         for i, p in enumerate(y):
@@ -91,31 +87,23 @@ def danbooru_tags(fpath):
                 tags.append(model_tags[i].item())
     except Exception as err:
         print(err)
+        tags = []
 
     return tags
 
-image_count = 0
 def image_infer(fpath):
-    global image_count
     tags = danbooru_tags(fpath)
-    image_count += 1
     return tags
 
-
-t1 = time.time()
 tags_1 = image_infer("example/1.png")
 tags_2 = image_infer("example/2.png")
 
-t2 = time.time()
-
 print(tags_1)
 # will show tags: ['1girl', 'beach', 'black_hair', 'blurry', 'blurry_background', 'blurry_foreground', 'building', 'bush', 'christmas_tree', 'day', 'depth_of_field', 'field', 'grass', 'lake', 'looking_at_viewer', 'mountain', 'nature', 'outdoors', 'palm_leaf', 'palm_tree', 'park', 'park_bench', 'path', 'photo_background', 'plant', 'river', 'road', 'skirt', 'sky', 'smile', 'striped', 'striped_dress', 'striped_shirt', 'tree', 'vertical-striped_shirt', 'vertical_stripes', 'rating:safe']
 
 print(tags_2)
 # will show tags: ['1girl', '3d', 'blurry', 'blurry_background', 'blurry_foreground', 'brown_eyes', 'brown_hair', 'bush', 'christmas_tree', 'cosplay_photo', 'day', 'depth_of_field', 'field', 'floral_print', 'foliage', 'forest', 'garden', 'grass', 'jungle', 'lips', 'long_hair', 'long_sleeves', 'looking_at_viewer', 'nature', 'on_grass', 'outdoors', 'palm_tree', 'park', 'path', 'plant', 'potted_plant', 'realistic', 'smile', 'solo', 'tree', 'upper_body', 'white_dress', 'rating:safe']
 
-print("-----------")
-print(f'infer speed(with mlx): {(t2 - t1)/image_count} seconds per image')
 ```