添加光照一致化ffc程序

tests/test_ffc.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+from orisys import FFCFrameProcessor
+
+import argparse
+from pathlib import Path
+
+import cv2
+
+from roi_polygon_tool import load_roi_mask, select_polygon_roi
+
+
+def parse_source(value):
+    value = str(value).strip()
+    for prefix in ("camera:", "cam:", "webcam:"):
+        if value.lower().startswith(prefix):
+            value = value[len(prefix):].strip()
+            break
+    try:
+        return int(value)
+    except ValueError:
+        return value
+
+
+def read_image_if_possible(source):
+    if isinstance(source, int):
+        return None
+    path = Path(source)
+    if not path.exists() or not path.is_file():
+        return None
+    image = cv2.imread(str(path), cv2.IMREAD_COLOR)
+    return image
+
+
+def draw_roi_outline(image, roi_mask):
+    view = image.copy()
+    contours, _ = cv2.findContours(roi_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    cv2.drawContours(view, contours, -1, (0, 255, 255), 1)
+    return view
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Test ROI drawing and flat-field correction.")
+    parser.add_argument("--source", default="0", help="Camera id, image path, or video path.")
+    parser.add_argument("--roi-config", default=None, help="Existing ROI JSON config.")
+    parser.add_argument("--roi-output", default="roi_polygon.json", help="Path to save ROI JSON.")
+    parser.add_argument("--weights", default=None, help="Existing FFC weights .npz to load.")
+    parser.add_argument("--save-weights", default="ffc_weights.npz", help="Path to save calibrated weights.")
+    parser.add_argument("--blur-kernel", type=int, default=151)
+    parser.add_argument("--blur-sigma", type=float, default=0.0)
+    parser.add_argument("--epsilon", type=float, default=1e-6)
+    parser.add_argument("--keep-outside", action="store_true", help="Keep pixels outside ROI.")
+    return parser.parse_args()
+
+
+def prepare_corrector(frame, args):
+    ffc = FFCFrameProcessor(
+        blur_kernel=args.blur_kernel,
+        blur_sigma=args.blur_sigma,
+        epsilon=args.epsilon,
+        masked_outside=not args.keep_outside,
+    )
+
+    if args.weights:
+        ffc.load(args.weights)
+        return ffc, ffc.roi_mask
+
+    if args.roi_config:
+        roi_mask = load_roi_mask(args.roi_config, frame.shape)
+    else:
+        roi = select_polygon_roi(frame, output=args.roi_output)
+        if roi is None:
+            raise RuntimeError("ROI selection cancelled.")
+        roi_mask = roi["mask"]
+
+    ffc.calibrate(frame, roi_mask=roi_mask)
+    if args.save_weights:
+        ffc.save(args.save_weights)
+    return ffc, roi_mask
+
+
+def run_single_image(image, args):
+    ffc, roi_mask = prepare_corrector(image, args)
+    corrected = ffc.correct(image)
+    cv2.imshow("Input ROI", draw_roi_outline(image, roi_mask))
+    cv2.imshow("Corrected ROI", draw_roi_outline(corrected, roi_mask))
+    cv2.waitKey(0)
+    cv2.destroyAllWindows()
+
+
+def run_video(source, args):
+    cap = cv2.VideoCapture(source)
+    if not cap.isOpened():
+        raise RuntimeError(f"Cannot open source: {source}")
+
+    try:
+        ok, frame = cap.read()
+        if not ok:
+            raise RuntimeError("Failed to read first frame.")
+
+        ffc, roi_mask = prepare_corrector(frame, args)
+
+        while True:
+            corrected = ffc.correct(frame)
+            cv2.imshow("Input ROI", draw_roi_outline(frame, roi_mask))
+            cv2.imshow("Corrected ROI", draw_roi_outline(corrected, roi_mask))
+
+            key = cv2.waitKey(1) & 0xFF
+            if key in (27, ord("q")):
+                break
+            if key == ord("r"):
+                ffc.calibrate(frame, roi_mask=roi_mask)
+                if args.save_weights:
+                    ffc.save(args.save_weights)
+                print("Recalibrated from current frame.")
+
+            ok, frame = cap.read()
+            if not ok:
+                break
+    finally:
+        cap.release()
+        cv2.destroyAllWindows()
+
+
+def main():
+    args = parse_args()
+    source = parse_source(args.source)
+
+    image = read_image_if_possible(source)
+    if image is not None:
+        run_single_image(image, args)
+    else:
+        run_video(source, args)
+
+
+if __name__ == "__main__":
+    main()