|
from abc import ABC, abstractmethod |
|
import logging |
|
from pathlib import Path |
|
from typing import Optional, Any |
|
from PIL import Image |
|
import pillow_jxl |
|
import torch |
|
from .imgproc_utils import ProcessingOptions, get_output_path, should_process_file |
|
|
|
class ImageProcessor(ABC): |
|
"""Base class for image processors""" |
|
|
|
def __init__(self, opts: ProcessingOptions): |
|
self.opts = opts |
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
@abstractmethod |
|
def load_models(self) -> None: |
|
"""Load required models""" |
|
pass |
|
|
|
@abstractmethod |
|
def process_image( |
|
self, |
|
image: Image.Image, |
|
image_path: Path, |
|
**kwargs: Any |
|
) -> str: |
|
"""Process a single image and return the result""" |
|
pass |
|
|
|
def save_result( |
|
self, |
|
result: str, |
|
image_path: Path, |
|
output_dir: Optional[Path] = None |
|
) -> None: |
|
"""Save the processing result""" |
|
output_path = get_output_path(image_path, self.opts, output_dir) |
|
output_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
|
with open(output_path, 'w', encoding='utf-8') as f: |
|
f.write(result) |
|
|
|
def process_file( |
|
self, |
|
image_path: Path, |
|
output_dir: Optional[Path] = None, |
|
**kwargs: Any |
|
) -> None: |
|
"""Process a single file""" |
|
if not should_process_file(image_path, self.opts, output_dir): |
|
return |
|
|
|
try: |
|
image = Image.open(image_path).convert('RGB') |
|
result = self.process_image(image, image_path, **kwargs) |
|
self.save_result(result, image_path, output_dir) |
|
except Exception as e: |
|
logging.error(f"Error processing {image_path}: {e}") |
|
|
