File size: 1,824 Bytes
fc115ba e75d927 fc115ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
from abc import ABC, abstractmethod
import logging
from pathlib import Path
from typing import Optional, Any
from PIL import Image
import pillow_jxl # type: ignore
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}")
|