 
import os
import warnings
from optimizer import Optimizer
from loss import Loss
from preprocess import Preprocessor
from feature_maps import FeatureMaps
import numpy as np
import time 
import torch
import argparse
import torchvision.models as models
import torch.optim as optim
from torchvision.utils import save_image
warnings.filterwarnings('ignore')
from logger import Logger
LOGGER = Logger().logger()
LOGGER.info("Started Style Transfer")
class StyleTransfer:
    '''
    Style Transfer Base Class
    '''
    def __init__(self) -> None:
        pass

    @staticmethod
    def pipeline():
        '''
        Pipeline for style transfer
        @params: None
        Author: @gaurangathavale
        '''

        device = torch.device( "cuda" if (torch.cuda.is_available()) else 'cpu')

        content_img_path = 'test/content.jpg'
        style_img_path = 'test/style.jpg'

        content_img = Preprocessor.process(content_img_path)
        style_img = Preprocessor.process(style_img_path)

        content_img_clone = content_img.clone().requires_grad_(True)

        Optimizer.gradient_descent(content_img, style_img, content_img_clone)

if __name__ == "__main__":
    stf = StyleTransfer()
    stf.pipeline()
        

    