File size: 3,156 Bytes
032e687 | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | import argparse
import concurrent.futures
import json
import os
from pathlib import Path
from tqdm import tqdm
from detectron2.data.detection_utils import read_image
def parse_args():
"""
Parse input arguments
"""
parser = argparse.ArgumentParser(description="Get image information")
parser.add_argument("--image_dir", help="path to image", type=str)
parser.add_argument("--json_path", help="path to json", type=str)
parser.add_argument("--output_path", type=str)
parser.add_argument("--max_workers", type=int, default=100)
args = parser.parse_args()
return args
args = parse_args()
def get_image_info(line, image_dir):
result = {
"status": "",
"id": None,
"file_name": None,
"height": None,
"width": None,
"channel": None,
}
file_name = line["file_name"]
file_path = os.path.join(image_dir, file_name)
if not os.path.isfile(file_path):
result["status"] = "NOFOUND"
print(line)
return result
try:
image = read_image(file_path, format="BGR")
except Exception as e:
result["status"] = "TRUNCATED"
print(e, line)
return result
result["status"] = "SUCCESS"
result["id"] = line["id"]
result["file_name"] = line["file_name"]
result["height"] = image.shape[0]
result["width"] = image.shape[1]
result["channel"] = image.shape[2]
return result
def get_images_info(data, image_dir, record_file):
with tqdm(total=len(data)) as pbar:
with concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor:
# Submit up to `chunk_size` tasks at a time to avoid too many pending tasks.
chunk_size = min(50000, args.max_workers * 500)
for i in range(0, len(data), chunk_size):
futures = [
executor.submit(get_image_info, line, image_dir)
for line in data[i : i + chunk_size]
]
for future in concurrent.futures.as_completed(futures):
r = future.result()
status, image_id, file_name, height, width, channel = (
r["status"],
r["id"],
r["file_name"],
r["height"],
r["width"],
r["channel"],
)
if status == "SUCCESS":
record_file.write(f"{image_id} {file_name} {height} {width} {channel}\n")
elif status == "NOFOUND":
pass
elif status == "TRUNCATED":
pass
else:
assert False
pbar.update(1)
def main():
print("loading", args.json_path)
json_data = json.load(open(args.json_path, "r"))
images = json_data["images"]
record_file = open(args.output_path, "w", encoding="utf8")
get_images_info(images, args.image_dir, record_file)
record_file.close()
if __name__ == "__main__":
main()
|