127 lines
4.3 KiB
Python
127 lines
4.3 KiB
Python
import os
|
|
import sys
|
|
import glob
|
|
import argparse
|
|
import numpy as np
|
|
import cv2
|
|
|
|
ROOT = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.append(ROOT)
|
|
|
|
from model_data_processing.Crop_And_Fill import (
|
|
process_complete,
|
|
smart_crop,
|
|
remove_noise,
|
|
masked_hist_equalize_color_bgr,
|
|
)
|
|
from Load_process.file_processing import Process_File
|
|
|
|
|
|
def list_images(path):
|
|
exts = ("*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tif", "*.tiff")
|
|
files = []
|
|
for ext in exts:
|
|
files.extend(glob.glob(os.path.join(path, ext)))
|
|
return sorted(files)
|
|
|
|
|
|
def make_synthetic_images(save_dir, n=3):
|
|
pf = Process_File()
|
|
pf.JudgeRoot_MakeDir(save_dir)
|
|
rng = np.random.default_rng(123)
|
|
paths = []
|
|
for i in range(n):
|
|
h, w = 300, 300
|
|
img = np.zeros((h, w, 3), dtype=np.uint8)
|
|
border = int(rng.integers(15, 40))
|
|
content = rng.integers(40, 180, size=(h - 2*border, w - 2*border, 3), dtype=np.uint8)
|
|
img[border:h-border, border:w-border] = content
|
|
path = os.path.join(save_dir, f"synthetic_{i+1}.png")
|
|
cv2.imwrite(path, img)
|
|
paths.append(path)
|
|
return paths
|
|
|
|
|
|
def run_single(image_path, output_dir, args):
|
|
pf = Process_File()
|
|
pf.JudgeRoot_MakeDir(output_dir)
|
|
img = cv2.imread(image_path)
|
|
if img is None:
|
|
raise FileNotFoundError(f"無法讀取: {image_path}")
|
|
|
|
denoised = img if args.no_denoise else remove_noise(img, method=args.denoise_method, strength=args.denoise_strength)
|
|
eq = masked_hist_equalize_color_bgr(
|
|
denoised,
|
|
ignore_below=(args.ignore_below if args.ignore_below is not None else args.threshold),
|
|
space=args.equalize_space,
|
|
)
|
|
cropped, borders, success = smart_crop(eq, args.threshold, args.min_content_ratio, keep_margin=0, show_info=False)
|
|
|
|
result = process_complete(
|
|
img,
|
|
threshold=args.threshold,
|
|
min_content_ratio=args.min_content_ratio,
|
|
fill_method=args.fill_method,
|
|
denoise=(not args.no_denoise),
|
|
denoise_method=args.denoise_method,
|
|
denoise_strength=args.denoise_strength,
|
|
equalize=(not args.no_equalize),
|
|
equalize_ignore_below=(args.ignore_below if args.ignore_below is not None else None),
|
|
equalize_space=args.equalize_space,
|
|
)
|
|
|
|
base = os.path.basename(image_path)
|
|
name, _ = os.path.splitext(base)
|
|
p_cropped = os.path.join(output_dir, f"{name}_cropped.png")
|
|
p_equalized = os.path.join(output_dir, f"{name}_equalized.png")
|
|
p_processed = os.path.join(output_dir, f"{name}_processed.png")
|
|
cv2.imwrite(p_cropped, cropped)
|
|
cv2.imwrite(p_equalized, eq)
|
|
cv2.imwrite(p_processed, result)
|
|
print(p_cropped)
|
|
print(p_equalized)
|
|
print(p_processed)
|
|
|
|
|
|
def run_batch(input_dir, output_dir, args):
|
|
images = list_images(input_dir)
|
|
if len(images) == 0:
|
|
raise FileNotFoundError(f"資料夾中沒有影像: {input_dir}")
|
|
for p in images:
|
|
run_single(p, output_dir, args)
|
|
print(output_dir)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="測試 Crop_And_Fill.py 流程並輸出裁切、等化、處理後影像")
|
|
parser.add_argument("--input", type=str, default=None)
|
|
parser.add_argument("--output", type=str, default=os.path.join(ROOT, "outputs", "crop_fill_test"))
|
|
parser.add_argument("--threshold", type=int, default=30)
|
|
parser.add_argument("--min_content_ratio", type=float, default=0.3)
|
|
parser.add_argument("--fill_method", type=str, default="inpaint_ns")
|
|
parser.add_argument("--no_denoise", action="store_true")
|
|
parser.add_argument("--denoise_method", type=str, default="median")
|
|
parser.add_argument("--denoise_strength", type=str, default="medium")
|
|
parser.add_argument("--no_equalize", action="store_true")
|
|
parser.add_argument("--ignore_below", type=int, default=None)
|
|
parser.add_argument("--equalize_space", type=str, default="LAB")
|
|
args = parser.parse_args()
|
|
|
|
pf = Process_File()
|
|
pf.JudgeRoot_MakeDir(args.output)
|
|
|
|
if args.input is None:
|
|
synth_dir = os.path.join(ROOT, "tests", "sample_images")
|
|
paths = make_synthetic_images(synth_dir, n=3)
|
|
for p in paths:
|
|
run_single(p, args.output, args)
|
|
else:
|
|
if os.path.isdir(args.input):
|
|
run_batch(args.input, args.output, args)
|
|
else:
|
|
run_single(args.input, args.output, args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|