adding files

Allow images with broken headers to load.

.gitignore

@@ -4,4 +4,5 @@ __pycache__/
 venv/
 .env
 
-presets/
+presets/
+pictures/

Dockerfile

@@ -0,0 +1,26 @@
+FROM python:3.12-slim
+WORKDIR /usr/src/app
+COPY . .
+#RUN cp .env.exp .env
+#RUN sed -i 's|SERVER_HOST=127.0.0.1|SERVER_HOST=0.0.0.0|' .env
+
+
+RUN --mount=target=/var/lib/apt/lists,type=cache,sharing=locked \
+    --mount=target=/var/cache/apt,type=cache,sharing=locked \
+    rm -f /etc/apt/apt.conf.d/docker-clean \
+    && apt-get update \
+    && apt-get install ffmpeg libsm6 libxext6 -y
+
+
+#RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 -y
+RUN pip install --cache-dir=/tmp/pipcache  --upgrade pip 
+RUN pip list --cache-dir=/tmp/pipcache --outdated
+RUN pip install --cache-dir=/tmp/pipcache -r requirements.txt
+#RUN pip install pdm
+#RUN pdm install
+#CMD ["pdm", "run", "python3", "main.py"]
+#CMD ["python3", "run.py"]
+#CMD ["python3", "run.py" ,"--threshold 0.5" ,"--file thereneerose_0002.jpg" ,"--rawtag"]
+#CMD ["python3", "run.py" ,"--threshold","0.5" ,"--file", "thereneerose_0002.jpg" ,"--rawtag"]
+
+#CMD ["python3 run.py --threshold 0.5 --file thereneerose_0002.jpg --rawtag"]

docker-compose.yml

@@ -0,0 +1,10 @@
+---
+version: "2"
+services:
+  wd14-tagger-standalone:
+    build: .
+    container_name: wd14-tagger-standalone
+    volumes:
+      - ./pictures:/pictures
+    restart: unless-stopped
+    command: sleep 3600

requirements.txt

@@ -1,9 +1,8 @@
-numpy==2.2.2
+numpy
 pandas==2.2.3
-torchvision==0.21.0
 Pillow==11.1.0
 huggingface-hub==0.28.1
 deepdanbooru==1.0.3
-onnxruntime==1.20.1
+onnxruntime
 opencv-python==4.11.0.86
 

run.py

@@ -2,12 +2,15 @@ import sys
 
 from typing import Generator, Iterable
 from tagger.interrogator.interrogator import AbsInterrogator
-from PIL import Image
+from PIL import Image, ImageFile
 from pathlib import Path
 import argparse
 
 from tagger.interrogators import interrogators
 
+# Allow images with broken headers to load
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
 parser = argparse.ArgumentParser()
 
 group = parser.add_mutually_exclusive_group(required=True)

tagger/interrogator/camietaggerinterrogator.py

@@ -5,11 +5,8 @@ https://huggingface.co/Camais03/camie-tagger/blob/main/onnx_inference.py
 """
 import sys
 import json
-
-from typing import cast
-
 import numpy as np
-import torch
+
 from PIL import Image
 from huggingface_hub import hf_hub_download
 from onnxruntime import InferenceSession
@@ -71,11 +68,20 @@ class CamieTaggerInterrogator(AbsInterrogator):
         if self.model is None:
             raise Exception("Model not loading.")
 
-        img_tensor = preprocess_image(image)
-        img_numpy = img_tensor.unsqueeze(0).numpy()
+        # Use the NumPy preprocessing function
+        # Result is already a NumPy array with shape (C, H, W) and float32 type
+        img_array_chw = preprocess_image(image) # e.g., shape (3, 512, 512)
+
+        # ONNX models typically expect input in NCHW format (Batch, Channel, Height, Width)
+        # Add a batch dimension (N=1) at the beginning
+        img_numpy = np.expand_dims(img_array_chw, axis=0) # Now shape (1, 3, 512, 512)
 
         input_ = self.model.get_inputs()[0]
 
+        # Ensure input data type matches model expectation (usually float32)
+        if input_.type == 'tensor(float)':
+            img_numpy = img_numpy.astype(np.float32)
+
         # evaluate model
         outputs = self.model.run(None, {input_.name: img_numpy})
 
@@ -100,12 +106,25 @@ class CamieTaggerInterrogator(AbsInterrogator):
             tags_by_category[category].append((tag_name, prob))
 
         # 'year', 'rating', 'general', 'character', 'copyright', 'artist', 'meta'
-        return dict(tags_by_category['rating']), dict(tags_by_category['general'])
+        rating_tags = dict(tags_by_category.get('rating', []))
+        general_tags = dict(tags_by_category.get('general', []))
 
-def preprocess_image(img: Image.Image, image_size=512) -> torch.Tensor:
-    """Process an image for inference"""
+        return rating_tags, general_tags
 
-    # Convert RGBA or Palette images to RGB
+def preprocess_image(img: Image.Image, image_size: int = 512) -> np.ndarray:
+    """
+    Process a PIL image for inference using NumPy.
+
+    Args:
+        img: The input PIL Image.
+        image_size: The target square size (width and height) for the output array.
+
+    Returns:
+        A NumPy array representing the processed image in CHW format,
+        with pixel values scaled to [0.0, 1.0] and dtype float32.
+    """
+
+    # Convert to RGB
     if img.mode in ('RGBA', 'P'):
         img = img.convert('RGB')
 
@@ -130,14 +149,13 @@ def preprocess_image(img: Image.Image, image_size=512) -> torch.Tensor:
     paste_y = (image_size - new_height) // 2
     new_image.paste(img, (paste_x, paste_y))
 
-    import torchvision.transforms as transforms
+    # Convert PIL image to NumPy array (H, W, C) with values 0-255
+    img_array = np.array(new_image, dtype=np.float32)
 
-    # Initialize transform
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-    ])
+    # Scale pixel values from [0, 255] to [0.0, 1.0]
+    img_array /= 255.0
 
-    # Apply transforms
-    img_tensor= transform(new_image)
-    img_tensor = cast(torch.Tensor, img_tensor)
-    return img_tensor
+    # Change dimension order from HWC to CHW
+    img_array = img_array.transpose((2, 0, 1)) # C, H, W
+
+    return img_array

tagger/interrogator/interrogator.py

@@ -6,7 +6,7 @@ import pandas as pd
 from typing import Iterable, Tuple, List, Dict
 from PIL import Image
 
-from onnxruntime import InferenceSession
+from onnxruntime import InferenceSession, get_available_providers
 
 tag_escape_pattern = re.compile(r'([\\()])')
 
@@ -66,7 +66,8 @@ class AbsInterrogator:
 
     def __init__(self, name: str) -> None:
         self.name = name
-        self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        # Initialize with optimal provider selection
+        self.providers = self.get_optimal_provider()
 
     def load(self):
         raise NotImplementedError()
@@ -85,7 +86,40 @@ class AbsInterrogator:
         return unloaded
 
     def use_cpu(self) -> None:
+        """Force CPU-only execution."""
         self.providers = ['CPUExecutionProvider']
+        print(f'Forcing CPU execution for {self.name}', file=sys.stderr)
+
+    def get_available_providers(self) -> List[str]:
+        """Get list of available execution providers."""
+        return get_available_providers()
+
+    def set_providers(self, providers: List[str]) -> None:
+        """Set specific execution providers."""
+        self.providers = providers
+        #print(f'Set custom providers for {self.name}: {providers}', file=sys.stderr)
+
+    def get_optimal_provider(self) -> List[str]:
+        """Get the optimal provider based on system capabilities.
+        
+        Returns a list of providers in order of preference:
+        - CoreMLExecutionProvider (if on Apple Silicon)
+        - CUDAExecutionProvider (if NVIDIA GPU available)
+        - CPUExecutionProvider (always available as fallback)
+        """
+        available = self.get_available_providers()
+        
+        # Start with most optimal providers first
+        optimal_order = [
+            'CoreMLExecutionProvider',  # Best for Apple Silicon
+            'CUDAExecutionProvider',    # Best for NVIDIA GPUs
+            'CPUExecutionProvider'      # Fallback
+        ]
+        
+        # Return list of available providers in optimal order
+        selected = [p for p in optimal_order if p in available]
+        #print(f'Selected optimal providers for {self.name}: {selected}', file=sys.stderr)
+        return selected
 
     def interrogate(
         self,