76 lines
2.7 KiB
Python
Executable File
76 lines
2.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import random
|
|
import argparse
|
|
from functools import partial
|
|
from common import *
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--data-dir', required=True)
|
|
parser.add_argument('--scan-id', help='Process only this specific scan ID')
|
|
parser.add_argument('--scan-id-range', help='Process scan IDs in range (e.g., "358626-358630" or "358626-358630,359000-359010")')
|
|
parser.add_argument('--force', action='store_true', help='Overwrite existing sbs.jpg files')
|
|
return parser.parse_args()
|
|
|
|
def process_scan(scan_dir, force=False):
|
|
if not os.path.isdir(scan_dir):
|
|
return "scan_dir not found"
|
|
frame_file = os.path.join(scan_dir, 'frame.jpg')
|
|
std_file = os.path.join(scan_dir, 'std.jpg')
|
|
if not os.path.exists(frame_file) or not os.path.exists(std_file):
|
|
return "frame.jpg or std.jpg not found"
|
|
sbs_file = os.path.join(scan_dir, 'sbs.jpg')
|
|
try:
|
|
if force or not os.path.exists(sbs_file):
|
|
frame_img = Image.open(frame_file)
|
|
std_img = Image.open(std_file)
|
|
sbs_img = make_side_by_side_img_with_margins(frame_img, std_img)
|
|
if sbs_img:
|
|
sbs_img.save(sbs_file)
|
|
else:
|
|
return "make_side_by_side_img_with_margins failed"
|
|
return "ok"
|
|
except Exception as e:
|
|
return f"error: {e}"
|
|
|
|
def main():
|
|
args = parse_args()
|
|
data_dir = args.data_dir
|
|
scans_dir = os.path.join(data_dir, 'scans')
|
|
scan_ids = os.listdir(scans_dir)
|
|
if args.scan_id:
|
|
if args.scan_id not in scan_ids:
|
|
print(f"Error: scan-id '{args.scan_id}' not found")
|
|
sys.exit(1)
|
|
scan_ids = [args.scan_id]
|
|
elif args.scan_id_range:
|
|
ranges = parse_ranges(args.scan_id_range)
|
|
filtered_scan_ids = []
|
|
for scan_id in scan_ids:
|
|
try:
|
|
scan_id_int = int(scan_id)
|
|
for val_range in ranges:
|
|
if in_range(scan_id_int, val_range):
|
|
filtered_scan_ids.append(scan_id)
|
|
break
|
|
except ValueError:
|
|
# Skip non-numeric scan IDs
|
|
continue
|
|
scan_ids = filtered_scan_ids
|
|
if not scan_ids:
|
|
print(f"Error: No scan IDs found in range '{args.scan_id_range}'")
|
|
sys.exit(1)
|
|
pool = Pool(cpu_count())
|
|
counts = defaultdict(int)
|
|
process_scan_with_force = partial(process_scan, force=args.force)
|
|
for result in tqdm(pool.imap(process_scan_with_force, [os.path.join(scans_dir, scan_id) for scan_id in scan_ids]), total=len(scan_ids)):
|
|
counts[result] += 1
|
|
for k, v in counts.items():
|
|
print(f"{k}: {v}")
|
|
|
|
if __name__ == '__main__':
|
|
main() |