Merge pull request #36 from codelformat/fix/test

Fix/test: improve test coverage

Author: panly2003

.gitignore

@@ -21,4 +21,7 @@ model/musiq/musiq_spaq_ckpt-358bb6af.pth
 VBench
 video_detection
 watermark/gm/ckpts/model_final.pth
-model_from_hf/
+model_from_hf/
+assets
+report.html
+.coverage

detection/__init__.py

@@ -12,24 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-Detection module for MarkDiffusion.
+"""Detection module for watermark verification and extraction."""
 
-This module provides detection functionality for various watermarking algorithms.
-"""
-
-__all__ = [
-    'base',
-    'gm',
-    'gs',
-    'prc',
-    'ri',
-    'robin',
-    'seal',
-    'sfw',
-    'tr',
-    'videomark',
-    'videoshield',
-    'wind',
-]
+from .base import BaseDetector
 
+__all__ = ["BaseDetector"]

evaluation/__init__.py

@@ -12,15 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-Evaluation module for MarkDiffusion.
+"""Evaluation module for watermark quality and robustness assessment."""
 
-This module provides tools for evaluating watermarking algorithms,
-including quality analysis and detection rate calculations.
-"""
+from .dataset import BaseDataset
 
-__all__ = [
-    'dataset',
-    'pipelines',
-    'tools',
-]
+__all__ = ["BaseDataset"]

evaluation/pipelines/image_quality_analysis.py

@@ -29,6 +29,20 @@
 )
 import lpips
 
+
+class SilentProgressBar:
+    """A silent progress bar wrapper that supports set_description but shows no output."""
+
+    def __init__(self, iterable):
+        self.iterable = iterable
+
+    def __iter__(self):
+        return iter(self.iterable)
+
+    def set_description(self, desc):
+        """No-op for silent mode."""
+        pass
+
 class QualityPipelineReturnType(Enum):
     """Return type of the image quality analysis pipeline."""
     FULL = auto()
@@ -120,7 +134,7 @@ def _get_progress_bar(self, iterable):
         """Return an iterable possibly wrapped with a progress bar."""
         if self.show_progress:
             return tqdm(iterable, desc="Processing", leave=True)
-        return iterable
+        return SilentProgressBar(iterable)
 
     def _get_prompt(self, index: int) -> str:
         """Get prompt from dataset."""

evaluation/pipelines/video_quality_analysis.py

@@ -11,6 +11,20 @@
 import numpy as np
 from tqdm import tqdm
 
+
+class SilentProgressBar:
+    """A silent progress bar wrapper that supports set_description but shows no output."""
+
+    def __init__(self, iterable):
+        self.iterable = iterable
+
+    def __iter__(self):
+        return iter(self.iterable)
+
+    def set_description(self, desc):
+        """No-op for silent mode."""
+        pass
+
 class QualityPipelineReturnType(Enum):
     """Return type of the image quality analysis pipeline."""
     FULL = auto()
@@ -97,7 +111,7 @@ def _get_progress_bar(self, iterable):
         """Return an iterable possibly wrapped with a progress bar."""
         if self.show_progress:
             return tqdm(iterable, desc="Processing", leave=True)
-        return iterable
+        return SilentProgressBar(iterable)
 
     def _get_prompt(self, index: int) -> str:
         """Get prompt from dataset."""

evaluation/tools/image_editor.py

@@ -212,8 +212,8 @@ def _add_salt_pepper_noise(self, img_array, amount):
         pepper_coords_y = np.random.randint(0, h, num_pepper)
         pepper_coords_x = np.random.randint(0, w, num_pepper)
         noisy[pepper_coords_y, pepper_coords_x] = 0
-        
-        return noisy
+
+        return np.clip(noisy, 0, 255).astype(np.uint8)
 
     def _add_poisson_noise(self, img_array):
         vals = len(np.unique(img_array))