Stomach_Cancer_Pytorch/test_one_image.py

93 lines
3.2 KiB
Python

import os
import sys
import argparse
import cv2
import numpy as np
ROOT = os.path.dirname(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 make_synth(path):
h, w = 300, 300
img = np.zeros((h, w, 3), dtype=np.uint8)
b = 30
content = np.random.default_rng(7).integers(40, 180, size=(h - 2*b, w - 2*b, 3), dtype=np.uint8)
img[b:h-b, b:w-b] = content
cv2.imwrite(path, img)
def hconcat(a, b):
h = max(a.shape[0], b.shape[0])
wa = int(a.shape[1] * h / a.shape[0])
wb = int(b.shape[1] * h / b.shape[0])
ra = cv2.resize(a, (wa, h))
rb = cv2.resize(b, (wb, h))
return cv2.hconcat([ra, rb])
def main():
p = argparse.ArgumentParser()
p.add_argument("--input", type=str, default=None)
p.add_argument("--output", type=str, default=os.path.join(ROOT, "outputs", "crop_fill_single"))
p.add_argument("--threshold", type=int, default=30)
p.add_argument("--min_content_ratio", type=float, default=0.3)
p.add_argument("--fill_method", type=str, default="inpaint_ns")
p.add_argument("--no_denoise", action="store_true")
p.add_argument("--denoise_method", type=str, default="median")
p.add_argument("--denoise_strength", type=str, default="medium")
p.add_argument("--no_equalize", action="store_true")
p.add_argument("--ignore_below", type=int, default=None)
p.add_argument("--equalize_space", type=str, default="LAB")
args = p.parse_args()
pf = Process_File()
pf.JudgeRoot_MakeDir(args.output)
src = "C1706.jpg"
img = cv2.imread(src)
if img is None:
raise FileNotFoundError(f"無法讀取: {src}")
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)
res = 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.splitext(os.path.basename(src))[0]
out_cropped = f"{base}_cropped.png"
out_equalized = f"{base}_equalized.png"
out_proc = f"{base}_processed.png"
out_cmp = f"{base}_compare.png"
cv2.imwrite(out_cropped, cropped)
cv2.imwrite(out_equalized, eq)
cv2.imwrite(out_proc, res)
cmp = hconcat(img, res)
cv2.imwrite(out_cmp, cmp)
print(out_cropped)
print(out_equalized)
print(out_proc)
print(out_cmp)
if __name__ == "__main__":
main()