1# /// script
2# requires-python = ">=3.11"
3# dependencies = [
4# "onnxruntime",
5# "pillow",
6# "tqdm",
7# ]
8# ///
9import argparse
10from PIL import Image
11import onnxruntime as rt
12import numpy as np
13import json
14from tqdm import tqdm
15
16def resize_and_pad(image_path, size: int, typ=np.float32):
17 # Open the image
18 img = Image.open(image_path)
19 # Resize the image while maintaining aspect ratio
20 img.thumbnail((size, size), Image.LANCZOS)
21 # Create a new image with a black background (zero-filled)
22 new_image = Image.new("RGB", (size, size), (0, 0, 0))
23 # Calculate the position to paste the resized image
24 paste_x = (size - img.width) // 2
25 paste_y = (size - img.height) // 2
26 # Paste the resized image onto the new image
27 new_image.paste(img, (paste_x, paste_y))
28 # Convert the image to a NumPy array
29 image_array = np.array(new_image)
30 image = image_array.transpose(2,0,1).astype(typ)
31 image /= 255
32 return image
33
34def sort_and_yield_indices(arr, thresh=0.1):
35 # Sort the array in descending order and get the sorted indices
36 sorted_indices = np.argsort(arr)[::-1]
37 sorted_arr = arr[sorted_indices]
38
39 # Yield indices of elements that are >= 0.0001
40 for index in sorted_indices:
41 if arr[index] < thresh:
42 break
43 yield (index,arr[index])
44
45def window(array, size):
46 for i in tqdm(range(0, len(array), size)):
47 yield array[i : i + size]
48
49# import line_profiler
50# @line_profiler.profile
51def main():
52 parser = argparse.ArgumentParser(description='Process some images.')
53
54 parser.add_argument('-m', '--model-path', required=True, help='Path to the model.')
55 parser.add_argument('-q', '--quant', type=int, required=True, choices=[16, 32], help='Quantization size (16 or 32).')
56 parser.add_argument('-b', '--batch-size', type=int, default=16, help='Batch size')
57 parser.add_argument('-tags', '--tag-path', required=True, help='Path to the tags.')
58 parser.add_argument('-thresh', type=float, help='Tag threshold', default=0.1)
59 parser.add_argument('-d', '--data-path', type=str, help='Data dump file')
60 parser.add_argument('--replace-path', type=str, help='replace path (i.e. `from:to`)')
61
62 parser.add_argument('images', nargs='+', type=str, help='Image to be processed')
63
64 args = parser.parse_args()
65
66 replace_path_from, replace_path_to = None, None
67 if args.replace_path:
68 replace_path_from, replace_path_to = args.replace_path.split(':')
69
70 import csv
71 excluded = set()
72 if args.data_path:
73 with open(args.data_path, 'r') as f:
74 csvreader = csv.reader(f, escapechar='\\', quoting=csv.QUOTE_NONE)
75 for row in csvreader:
76 if replace_path_from:
77 excluded.add(row[0].replace(replace_path_from, replace_path_to))
78 else:
79 excluded.add(row[0])
80 args.images = list(filter(lambda x: x not in excluded, args.images))
81 print(f"{len(args.images)} images to be processed")
82 if not args.images:
83 return
84
85 vocab = json.load(open(args.tag_path,"r"))
86 if args.quant == 16:
87 typ = np.float16
88 elif args.quant == 32:
89 typ = np.float32
90 img_size = 224
91
92 # initialize onnx runtime inference session
93 sess_opt = rt.SessionOptions()
94 sess = rt.InferenceSession(args.model_path, sess_opt)
95 input_name = sess.get_inputs()[0].name
96 output_name = sess.get_outputs()[0].name
97
98 import concurrent.futures
99
100 if args.data_path:
101 outf = open(args.data_path, 'a')
102 else:
103 outf = None
104 for images in window(args.images, args.batch_size):
105 with concurrent.futures.ThreadPoolExecutor() as executor:
106 inp_tensors = np.array(list(executor.map(lambda img: resize_and_pad(img, img_size,typ=typ), images)))
107 results = sess.run([output_name], {input_name: inp_tensors})[0]
108
109 if outf:
110 for (K, path) in zip(results, images):
111 if replace_path_from:
112 path=path.replace(replace_path_to,replace_path_from)
113 outf.write(path.replace(",", "\\,"))
114 for index, v in sort_and_yield_indices(K, args.thresh):
115 outf.write(",")
116 outf.write(vocab[index])
117 outf.write("\n")
118 else:
119 for (K, path) in zip(results, images):
120 print(path)
121 for index, v in sort_and_yield_indices(K, args.thresh):
122 print(f" {vocab[index]}:{v}")
123
124if __name__ == '__main__':
125 main()
126
127
