Implement image upscaling service with format conversion and add unit tests

Author: LightCyan01

backend/service/image_processing.py

@@ -1,7 +1,43 @@
-from spandrel import ImageModelDescriptor
+from spandrel import ModelLoader
+from pathlib import Path
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+from config.image_formats import SUPPORTED_OUTPUT_FORMATS, SUPPORTED_INPUT_EXTENSIONS, get_supported_format
 
-class image_service():
-    def __init__(self, model: ImageModelDescriptor):
-        self._model = model
+class ImageService():
+    def __init__(self, model_path, device):
+        loader = ModelLoader()
+        model = loader.load_from_file(model_path)
+        self.model = model.to(device).eval()
 
-    
+    @staticmethod
+    def get_supported_output_formats():
+        return get_supported_format()
+        
+    
+    def upscale(self, images_path: str, save_path: str, format: str = "jpg"):
+        images_path = Path(images_path)
+        save_path = Path(save_path)
+        
+        image_files = [f for f in images_path.iterdir() if f.is_file() and f.suffix.lower() in SUPPORTED_INPUT_EXTENSIONS]
+        format_info = SUPPORTED_OUTPUT_FORMATS[format.lower()]
+        
+        for image_file in image_files:
+            image = Image.open(image_file)
+        
+            if image.mode == "RGBA":
+                image = image.convert('RGB')
+        
+         
+            to_tensor = transforms.ToTensor()
+            image_tensor = to_tensor(image).unsqueeze(0).to(self.model.device)
+        
+            with torch.no_grad():
+                result = self.model(image_tensor)
+            
+            to_pil = transforms.ToPILImage()
+            output_image = to_pil(result.cpu().squeeze(0))
+        
+            file_name = image_file.stem + "_upscaled" + format_info["extension"]
+            output_image.save(save_path / file_name, format=format_info["pil_format"])

backend/tests/test_image_processing.py

@@ -0,0 +1,71 @@
+import pytest
+import tempfile
+import shutil
+from pathlib import Path
+from PIL import Image
+from unittest.mock import Mock, patch
+from service.image_processing import ImageService
+import torch
+
+
+@pytest.fixture
+def temp_dirs():
+    """Create temporary input and output directories."""
+    input_dir = Path(tempfile.mkdtemp())
+    output_dir = Path(tempfile.mkdtemp())
+    
+    test_image = Image.new('RGB', (100, 100), color='red')
+    test_image.save(input_dir / "test.jpg", "JPEG")
+    
+    yield input_dir, output_dir
+    
+    shutil.rmtree(input_dir)
+    shutil.rmtree(output_dir)
+
+
+@pytest.fixture
+def image_service():
+    with patch('service.image_processing.ModelLoader') as mock_loader:
+        # Create a mock model that returns the same tensor (no actual upscaling)
+        mock_model = Mock()
+        mock_model.to.return_value = mock_model
+        mock_model.eval.return_value = mock_model
+        mock_model.device = "cpu"
+        mock_model.return_value = torch.rand(1, 3, 100, 100) 
+        
+        mock_loader.return_value.load_from_file.return_value = mock_model
+        
+        return ImageService("dummy_model_path.pth", "cpu")
+
+
+def test_upscale_default_format(image_service, temp_dirs):
+    """Test upscaling with default JPG format."""
+    input_dir, output_dir = temp_dirs
+    
+    image_service.upscale(str(input_dir), str(output_dir))
+    
+    output_files = list(output_dir.glob("*_upscaled.*"))
+    assert len(output_files) == 1
+    assert output_files[0].suffix == ".jpg"
+
+
+def test_upscale_format_conversion(image_service, temp_dirs):
+    """Test upscaling with format conversion."""
+    input_dir, output_dir = temp_dirs
+    
+    image_service.upscale(str(input_dir), str(output_dir), "png")
+    
+    output_files = list(output_dir.glob("*_upscaled.png"))
+    assert len(output_files) == 1
+    
+    with Image.open(output_files[0]) as img:
+        assert img.format == "PNG"
+
+
+def test_get_supported_formats():
+    """Test getting supported output formats."""
+    formats = ImageService.get_supported_output_formats()
+    
+    assert isinstance(formats, list)
+    assert len(formats) > 0
+    assert all("value" in fmt and "label" in fmt for fmt in formats)