|
- import glob
- import os
- import click
-
- import deepdanbooru as dd
- import numpy as np
- import PIL.Image
- import tensorflow as tf
-
- from tqdm import tqdm
- from typing import List
-
-
- def load_model(path='/pretrainmodel/model-resnet_custom_v3.h5') -> tf.keras.Model:
- model = tf.keras.models.load_model(path)
- return model
-
-
- @click.command()
- @click.option('--input', default='/dataset/danbooru-2020-512', help='Input dir')
- @click.option('--output', default='/model/danbooru-2020.npy', help='Output file')
- @click.option('--ckpt_url', default='/pretrainmodel/model-resnet_custom_v3.h5', help='Output file')
- def main(input, output, ckpt_url):
- if input.endswith('/'):
- input = input[:-1]
-
- score_threshold = 0.5
- model = load_model(ckpt_url)
-
- files = glob.glob(os.path.join(input, '**', '*.jpg'))
- files += glob.glob(os.path.join(input, '**', '*.png'))
- files += glob.glob(os.path.join(input, '**', '*.bmp'))
- files = sorted(files)
-
- save_dict = dict()
- for path in tqdm(files):
- basename = '.'.join(os.path.basename(path).split('.')[:-1])
- image = PIL.Image.open(path)
- _, height, width, _ = model.input_shape
- image = np.asarray(image)
- image = tf.image.resize(image,
- size=(height, width),
- method=tf.image.ResizeMethod.AREA,
- preserve_aspect_ratio=True)
- image = image.numpy()
- image = dd.image.transform_and_pad_image(image, width, height)
- image = image / 255.
- probs = model.predict(image[None, ...])[0]
- probs = probs.astype(float)
- ids = []
- for idx, prob in enumerate(probs.tolist()):
- if prob < score_threshold:
- continue
- ids.append(idx)
- ids = np.asarray(ids, dtype=np.uint16)
-
- srckey = path.replace('\\', '/').replace('//', '/')
- key = srckey.split(input)[1]
- save_dict.update({ key: ids })
-
- np.save(output, save_dict, allow_pickle=True)
- # np.load(output, allow_pickle=True).item()
-
- if __name__ == '__main__':
- main()
|