80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import caffe
|
|
import cv2
|
|
|
|
|
|
BASE_DIR = Path(__file__).resolve().parents[1]
|
|
SUPPORTED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
|
|
|
|
|
|
def edit_fcn_proto(template_file: Path, output_file: Path, height: int, width: int) -> None:
|
|
template = template_file.read_text()
|
|
output_file.write_text(template.format(height=height, width=width))
|
|
|
|
|
|
def iter_images(input_dir: Path) -> list[Path]:
|
|
return [
|
|
path
|
|
for path in sorted(input_dir.iterdir(), key=lambda p: p.name.lower())
|
|
if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS
|
|
]
|
|
|
|
|
|
def run(input_dir: Path, output_dir: Path, max_dim: int) -> None:
|
|
caffe.set_mode_cpu()
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
template_file = BASE_DIR / "test" / "test_template.prototxt"
|
|
deploy_file = BASE_DIR / "test" / "DeployT.prototxt"
|
|
model_file = BASE_DIR / "AOD_Net.caffemodel"
|
|
|
|
images = iter_images(input_dir)
|
|
print(f"image numbers: {len(images)}")
|
|
|
|
for index, img_path in enumerate(images, start=1):
|
|
npstore = caffe.io.load_image(str(img_path))
|
|
orig_h, orig_w = npstore.shape[0], npstore.shape[1]
|
|
|
|
if orig_h > max_dim or orig_w > max_dim:
|
|
scale = max_dim / float(max(orig_h, orig_w))
|
|
new_h = int(orig_h * scale)
|
|
new_w = int(orig_w * scale)
|
|
npstore = cv2.resize(npstore, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
|
|
print(f"Resized {img_path.name} from {orig_w}x{orig_h} to {new_w}x{new_h}")
|
|
|
|
height, width = npstore.shape[0], npstore.shape[1]
|
|
edit_fcn_proto(template_file, deploy_file, height, width)
|
|
|
|
net = caffe.Net(str(deploy_file), str(model_file), caffe.TEST)
|
|
data = npstore.transpose((2, 0, 1))
|
|
net.blobs["data"].data[...] = [data]
|
|
net.forward()
|
|
|
|
result = net.blobs["sum"].data[0].transpose((1, 2, 0))
|
|
result = result[:, :, ::-1]
|
|
|
|
if height != orig_h or width != orig_w:
|
|
result = cv2.resize(result, (orig_w, orig_h), interpolation=cv2.INTER_CUBIC)
|
|
|
|
save_path = output_dir / f"{img_path.stem}_AOD-Net.png"
|
|
cv2.imwrite(str(save_path), result * 255.0, [cv2.IMWRITE_JPEG_QUALITY, 100])
|
|
print(f"[{index}/{len(images)}] saved: {save_path}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Run AOD-Net on a folder of images.")
|
|
parser.add_argument("input_dir", nargs="?", default=str(BASE_DIR / "data" / "img"))
|
|
parser.add_argument("output_dir", nargs="?", default=str(BASE_DIR / "data" / "result"))
|
|
parser.add_argument("--max-dim", type=int, default=1920)
|
|
args = parser.parse_args()
|
|
|
|
run(Path(args.input_dir), Path(args.output_dir), args.max_dim)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|