47 lines
1.1 KiB
Python
47 lines
1.1 KiB
Python
|
from sdsv_dewarp.api import AlignImage
|
||
|
|
import cv2
|
||
|
|
import glob
|
||
|
|
import os
|
||
|
|
import tqdm
|
||
|
|
import time
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--input")
|
||
|
|
parser.add_argument("--out")
|
||
|
|
parser.add_argument("--device", type=str, default="cuda:1")
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
model = AlignImage(device=args.device)
|
||
|
|
|
||
|
|
|
||
|
|
img_dir = args.input
|
||
|
|
out_dir = args.out
|
||
|
|
if not os.path.exists(out_dir):
|
||
|
|
os.makedirs(out_dir)
|
||
|
|
|
||
|
|
img_paths = glob.glob(img_dir + "/*")
|
||
|
|
|
||
|
|
times = []
|
||
|
|
for img_path in tqdm.tqdm(img_paths):
|
||
|
|
t1 = time.time()
|
||
|
|
img = cv2.imread(img_path)
|
||
|
|
if img is None:
|
||
|
|
print(img_path)
|
||
|
|
continue
|
||
|
|
|
||
|
|
aligned_img, is_blank, angle_align = model(img)
|
||
|
|
|
||
|
|
times.append(time.time() - t1)
|
||
|
|
|
||
|
|
if not is_blank:
|
||
|
|
cv2.imwrite(os.path.join(out_dir, os.path.basename(img_path)), aligned_img)
|
||
|
|
else:
|
||
|
|
cv2.imwrite(os.path.join(out_dir, os.path.basename(img_path)), img)
|
||
|
|
|
||
|
|
|
||
|
|
times = times[1:]
|
||
|
|
print("Avg time: ", sum(times) / len(times))
|