Updated core modules, made some unit tests for them.

core/ImageProcessor.py

@@ -3,12 +3,14 @@
 class ImageProcessor:
     @staticmethod
     def preprocess_image(image_path, size=512):
+        '''Resize and normalize image for input.'''
         image = Image.open(image_path).convert("RGB")
         image = image.resize((size, size))
         return image 
 
     @staticmethod
     def save_image(image, output_path):
+        '''Save the processed image.'''
         image.save(output_path)
 
 

core/StyleTransferModel.py

@@ -1,11 +1,20 @@
 import torch 
 import torchvision.transforms as transforms
 from PIL import Image 
+from torchvision.models import vgg19
+
 
 class StyleTransferModel:
-    def __init__(self, model_path):
+    def __init__(self):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model = torch.load(model_path).to(self.device)
+        self.model = self._load_pretrained_model()
+
+    def _load_pretrained_model(self):
+        '''Using VGG-19 for feature extraction'''
+        model = vgg19(pretrained=True).features
+        for param in model.parameters():
+            param.requires_grad = False 
+        return model.to(self.device)
 
     def apply_style(self, content_image, style_image):
         '''Style transfer logic'''

tests/test_artify.py

@@ -0,0 +1,21 @@
+from core.StyleTransferModel import StyleTransferModel
+from core.ImageProcessor import ImageProcessor
+
+if __name__ == "__main__":
+    content_image_path = ""
+    style_image_path = ""
+    output_image_path = ""
+
+    processor = ImageProcessor()
+    model = StyleTransferModel()
+
+    content_image = processor.preprocess_image(content_image_path)
+    style_image = processor.preprocess_image(style_image_path)
+
+    # Apply style (placeholder for now)
+    # styled_image = model.apply_style(content_image, style_image)
+
+    # Save the output (placeholder for now)
+    # processor.save_image(styled_image, output_image_path)
+
+    print("Pipeline tested successfully")

tests/test_image_processor.py

@@ -0,0 +1,13 @@
+import os
+from core.ImageProcessor import ImageProcessor
+
+def test_preprocess_image():
+    processor = ImageProcessor()
+    processed_image = processor.preprocess_image("")
+    assert processed_image.size == (512, 512)
+
+def test_save_image():
+    processor = ImageProcessor()
+    sample_image = processor.preprocess_image("")
+    processor.save_image(sample_image, "")
+    assert os.path.exists("")

tests/test_style_transfer_model.py

@@ -0,0 +1,5 @@
+from core.StyleTransferModel import StyleTransferModel
+
+def test_model_initialization():
+    model = StyleTransferModel()
+    assert model.model is not None