【Hackathon 10th Spring No.14】SFIN模型复现

README.md

@@ -33,6 +33,7 @@
 ## 📑 Task
 - [MLIP-Machine Learning Interatomic Potential](interatomic_potentials/README.md)
 - [MLES-Machine Learning Electronic Structure](electronic_structure/README.md)
+- [CM-Crystal Materials](crystal_materials/README.md)
 - [PP-Property Prediction](property_prediction/README.md)
 - [SG-Structure Generation](structure_generation/README.md)
 - [SE-Spectrum Elucidation](spectrum_elucidation/README.md)

crystal_materials/README.md

@@ -0,0 +1,30 @@
+# CM-Crystal Materials
+
+## 1.Introduction
+
+Crystal Materials (CM) is a task category in PaddleMaterials for crystal-related model workflows.
+Current support includes microscopy enhancement for crystalline materials with SFIN.
+
+The supported model is:
+
+- **SFIN**: *Noise Calibration and Spatial-Frequency Interactive Network for STEM Image Enhancement* (CVPR 2025)  
+  Paper: https://arxiv.org/pdf/2504.02555
+
+## 2.Models Matrix
+
+| **Supported Functions**                           | **[SFIN](./configs/sfin/README.md)** |
+| ------------------------------------------------- | :-----------------------------------: |
+| **Crystal Microscopy**                            |                                       |
+| STEM image denoising/enhancement                  |                  ✅                   |
+| **ML Capabilities · Training**                    |                                       |
+| Single-GPU                                        |                  ✅                   |
+| Distributed training                              |                  ✅                   |
+| Mixed precision (AMP)                             |                  —                    |
+| Fine-tuning                                       |                  ✅                   |
+| **ML Capabilities · Evaluation**                  |                                       |
+| PSNR                                              |                  ✅                   |
+| SSIM                                              |                  ✅                   |
+| **Datasets**                                      |                                       |
+| Paired STEM `noisy` / `gt_enhance` image datasets |                  ✅                   |
+
+**Notice**:🌟 represent originate research work published from paddlematerials toolkit

crystal_materials/configs/sfin/README.md

@@ -0,0 +1,116 @@
+# SFIN
+
+[Noise Calibration and Spatial-Frequency Interactive Network for STEM Image Enhancement](https://arxiv.org/pdf/2504.02555)
+
+## 1.Introduction
+
+SFIN is a CNN model for STEM image enhancement in crystal materials microscopy.
+Input is a noisy grayscale STEM image (`noisy`), and output is the enhanced grayscale image (`gt_enhance` target during training).
+
+## 2.Dataset Format
+
+The default config `sfin_tem_enhance.yaml` expects paired images under:
+
+```text
+../sfin/
+  data/
+    noisy/
+      0.png
+      ...
+    gt_enhance/
+      0.png
+      ...
+  data_test/
+    noisy/
+      0.png
+      ...
+    gt_enhance/
+      0.png
+      ...
+  checkpoints/
+    sfin_he_500.pdparams
+```
+
+Notes:
+- `STEMImageDataset` uses `strict_index_naming: True` by default, so paired files should follow indexed names (such as `0.png`, `1.png`).
+- `Dataset.train.dataset.__init_params__.data_path` points to `../sfin/data`.
+- `Dataset.val/test.dataset.__init_params__.data_path` points to `../sfin/data_test`.
+
+## 3.Command
+
+Run commands from the `PaddleMaterials` root directory.  
+If `ppmat` is not installed in your environment, install once:
+
+```bash
+pip install -e . --no-build-isolation
+```
+
+Single GPU training:
+
+```bash
+python crystal_materials/train.py \
+  -c crystal_materials/configs/sfin/sfin_tem_enhance.yaml
+```
+
+Multi GPU training:
+
+```bash
+python -m paddle.distributed.launch --gpus="0,1,2,3" \
+  crystal_materials/train.py \
+  -c crystal_materials/configs/sfin/sfin_tem_enhance.yaml
+```
+
+Tiny smoke training with 2 images (for quick validation):
+
+```bash
+python crystal_materials/train.py \
+  -c crystal_materials/configs/sfin/sfin_tem_enhance.yaml \
+  Trainer.max_epochs=1 \
+  Trainer.save_freq=1 \
+  Trainer.log_freq=1 \
+  Trainer.eval_freq=1 \
+  Global.do_test=False \
+  Dataset.train.dataset.__init_params__.data_path='../sfin/data_test' \
+  Dataset.train.dataset.__init_params__.data_count=2 \
+  Dataset.train.sampler.__init_params__.batch_size=1 \
+  Dataset.train.sampler.__init_params__.drop_last=False \
+  Dataset.val.dataset.__init_params__.data_count=2
+```
+
+Eval only with a checkpoint:
+
+```bash
+python crystal_materials/train.py \
+  -c crystal_materials/configs/sfin/sfin_tem_enhance.yaml \
+  Global.do_train=False Global.do_eval=True Global.do_test=False \
+  Trainer.pretrained_model_path='path/to/model.pdparams'
+```
+
+Predict with provided SFIN checkpoint:
+
+```bash
+python crystal_materials/predict.py \
+  --config_path crystal_materials/configs/sfin/sfin_tem_enhance.yaml \
+  --checkpoint_path ../sfin/checkpoints/sfin_he_500.pdparams \
+  --input_dir ../sfin/data_test/noisy \
+  --output_dir ./output/sfin_tem_enhance/predictions
+```
+
+Evaluate PSNR + SSIM on prediction folder:
+
+```bash
+python crystal_materials/evaluate.py \
+  --gt_dir ../sfin/data_test/gt_enhance \
+  --pred_dir ./output/sfin_tem_enhance/predictions
+```
+
+## 4.Citation
+
+```bibtex
+@article{Li2025SFIN,
+  title={Noise Calibration and Spatial-Frequency Interactive Network for STEM Image Enhancement},
+  author={Li, Hesong and Wu, Ziqi and Shao, Ruiwen and Zhang, Tao and Fu, Ying},
+  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
+  year={2025}
+}
+```

crystal_materials/configs/sfin/sfin_tem_enhance.yaml

@@ -0,0 +1,128 @@
+Global:
+  do_train: True
+  do_eval: True
+  do_test: True
+  prim_eager_enabled: False
+
+Trainer:
+  max_epochs: 500
+  seed: 42
+  output_dir: ./output/sfin_tem_enhance
+  save_freq: 100
+  log_freq: 100
+  start_eval_epoch: 1
+  eval_freq: 1
+  pretrained_model_path: null
+  pretrained_weight_name: null
+  resume_from_checkpoint: null
+  use_amp: False
+  amp_level: "O1"
+  eval_with_no_grad: True
+  gradient_accumulation_steps: 1
+
+  best_metric_indicator: "eval_metric"
+  name_for_best_metric: "gt_enhance"
+  greater_is_better: True
+
+  compute_metric_during_train: True
+  metric_strategy_during_eval: "epoch"
+
+  use_visualdl: False
+  use_wandb: False
+  use_tensorboard: False
+
+Model:
+  __class_name__: SFIN
+  __init_params__:
+    in_channels: 1
+    base_channels: 64
+    num_blocks: 8
+    input_name: "noisy"
+    target_name: "gt_enhance"
+    loss_type: "l1"
+    loss_weight: 1.0
+
+Optimizer:
+  __class_name__: Adam
+  __init_params__:
+    lr:
+      __class_name__: MultiStepDecay
+      __init_params__:
+        learning_rate: 2.0e-4
+        milestones: [250, 400, 425, 450, 475]
+        gamma: 0.5
+        by_epoch: True
+    beta1: 0.9
+    beta2: 0.999
+    epsilon: 1.0e-8
+    weight_decay: 0.0
+
+Metric:
+  gt_enhance:
+    __class_name__: PSNRMetric
+    __init_params__:
+      data_range: 255.0
+
+Dataset:
+  train:
+    dataset:
+      __class_name__: STEMImageDataset
+      __init_params__:
+        data_path: "../sfin/data"
+        data_count: 1000
+        noisy_subdir: "noisy"
+        target_subdir: "gt_enhance"
+        strict_index_naming: True
+        scale_to_unit: False
+    sampler:
+      __class_name__: BatchSampler
+      __init_params__:
+        shuffle: True
+        drop_last: True
+        batch_size: 8
+    loader:
+      num_workers: 0
+      use_shared_memory: False
+      collate_fn: DefaultCollator
+
+  val:
+    dataset:
+      __class_name__: STEMImageDataset
+      __init_params__:
+        data_path: "../sfin/data_test"
+        data_count: 100
+        noisy_subdir: "noisy"
+        target_subdir: "gt_enhance"
+        strict_index_naming: True
+        scale_to_unit: False
+    sampler:
+      __class_name__: BatchSampler
+      __init_params__:
+        shuffle: False
+        drop_last: False
+        batch_size: 1
+    loader:
+      num_workers: 0
+      use_shared_memory: False
+      collate_fn: DefaultCollator
+
+  test:
+    dataset:
+      __class_name__: STEMImageDataset
+      __init_params__:
+        data_path: "../sfin/data_test"
+        data_count: 100
+        noisy_subdir: "noisy"
+        target_subdir: "gt_enhance"
+        strict_index_naming: True
+        scale_to_unit: False
+    sampler:
+      __class_name__: BatchSampler
+      __init_params__:
+        shuffle: False
+        drop_last: False
+        batch_size: 1
+    loader:
+      num_workers: 0
+      use_shared_memory: False
+      collate_fn: DefaultCollator

crystal_materials/evaluate.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import paddle
+from PIL import Image
+
+from ppmat.metrics import PSNRMetric
+from ppmat.metrics import SSIMMetric
+
+
+def load_gray_tensor(path: Path) -> paddle.Tensor:
+    arr = np.asarray(Image.open(path).convert("L"), dtype=np.float32)
+    return paddle.to_tensor(arr).unsqueeze(0).unsqueeze(0)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--gt_dir",
+        type=str,
+        default="../sfin/data_test/gt_enhance",
+        help="Ground truth image directory.",
+    )
+    parser.add_argument(
+        "--pred_dir",
+        type=str,
+        required=True,
+        help="Predicted image directory.",
+    )
+    parser.add_argument(
+        "--file_suffix",
+        type=str,
+        default=".png",
+        help="Image file suffix.",
+    )
+    args = parser.parse_args()
+
+    gt_dir = Path(args.gt_dir)
+    pred_dir = Path(args.pred_dir)
+
+    gt_files = sorted([p for p in gt_dir.glob(f"*{args.file_suffix}") if p.is_file()])
+
+    psnr_metric = PSNRMetric(data_range=255.0)
+    ssim_metric = SSIMMetric(data_range=255.0)
+
+    psnr_sum = 0.0
+    ssim_sum = 0.0
+    count = 0
+    for gt_path in gt_files:
+        pred_path = pred_dir / gt_path.name
+        if not pred_path.exists():
+            continue
+        gt = load_gray_tensor(gt_path)
+        pred = load_gray_tensor(pred_path)
+        psnr_sum += float(psnr_metric(pred, gt))
+        ssim_sum += float(ssim_metric(pred, gt))
+        count += 1
+
+    if count == 0:
+        raise RuntimeError("No matched prediction files found for evaluation.")
+
+    print(f"Matched images: {count}")
+    print(f"Average PSNR: {psnr_sum / count:.6f}")
+    print(f"Average SSIM: {ssim_sum / count:.6f}")