app.py

import os
import torch
from flask import Flask, render_template, request, redirect, url_for, send_from_directory
from flask_wtf import FlaskForm
from flask_bootstrap import Bootstrap
from werkzeug.utils import secure_filename
from wtforms import FileField, SubmitField, FloatField, HiddenField
from wtforms.validators import InputRequired
from PIL import Image
from torchvision import transforms
import io
import gc
torch.set_num_threads(1)

# Import your existing AdaIN code
from utils.models import VGGEncoder, Decoder
from utils.utils import adaptive_instance_normalization, calc_mean_std


app = Flask(__name__)
app.config['SECRET_KEY'] = os.environ.get(
    'SECRET_KEY',
    'supersecretkey'
)
app.config['UPLOAD_FOLDER'] = 'static/uploads'
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg'}
Bootstrap(app)

os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

class UploadForm(FlaskForm):
    content = FileField('Content Image')
    style = FileField('Style Image')
    content_path = HiddenField()
    style_path = HiddenField()
    alpha = FloatField('Alpha', default=1.0)
    submit = SubmitField('Transfer Style')

device = torch.device("cpu")

encoder = VGGEncoder('vgg_normalised.pth').to(device)
decoder = Decoder().to(device)
decoder.load_state_dict(
    torch.load(
        'experiment/big_dataset/decoder_final.pth',
        map_location=device
    )
)

encoder.eval()
decoder.eval()

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

def style_transfer(content_image, style_image,
                   encoder, decoder,
                   alpha, device):

    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])

    content_image = transform(
        content_image
    ).unsqueeze(0).to(device)

    style_image = transform(
        style_image
    ).unsqueeze(0).to(device)

    with torch.no_grad():

        content_feats = encoder(
            content_image,
            is_test=True
        )

        style_feats = encoder(
            style_image,
            is_test=True
        )

        stylized_feats = adaptive_instance_normalization(
            content_feats,
            style_feats
        )

        stylized_feats = (
            alpha*stylized_feats
            +(1-alpha)*content_feats
        )

        stylized_image = decoder(
            stylized_feats
        )

    del content_feats
    del style_feats
    del stylized_feats
    del content_image
    del style_image

    gc.collect()

    return stylized_image


from PIL import ImageEnhance

def save_image(image, path):
    image = image.cpu().clone()
    image = image.squeeze(0)
    image = image.clamp(0,1)

    image = transforms.ToPILImage()(image)

    image = image.resize((768,768))

    enhancer = ImageEnhance.Sharpness(image)
    image = enhancer.enhance(1.8)

    color = ImageEnhance.Color(image)
    image = color.enhance(1.15)

    contrast = ImageEnhance.Contrast(image)
    image = contrast.enhance(1.1)

    image.save(path, quality=100)
    



@app.route('/', methods=['GET', 'POST'])
def index():
    form = UploadForm()

    result_image = None
    content_filename = None
    style_filename = None
    error = None

    if request.method == 'POST':

        if form.validate_on_submit():

            if form.content.data and form.content.data.filename:
                if allowed_file(form.content.data.filename):
                    content_filename = secure_filename(
                        form.content.data.filename
                    )

                    form.content.data.save(
                        os.path.join(
                            app.config['UPLOAD_FOLDER'],
                            content_filename
                        )
                    )

                    form.content_path.data = content_filename

            else:
                content_filename = form.content_path.data


            if form.style.data and form.style.data.filename:
                if allowed_file(form.style.data.filename):

                    style_filename = secure_filename(
                        form.style.data.filename
                    )

                    form.style.data.save(
                        os.path.join(
                            app.config['UPLOAD_FOLDER'],
                            style_filename
                        )
                    )

                    form.style_path.data = style_filename

            else:
                style_filename = form.style_path.data


            if content_filename and style_filename:

                try:

                    content_path = os.path.join(
                        app.config['UPLOAD_FOLDER'],
                        content_filename
                    )

                    style_path = os.path.join(
                        app.config['UPLOAD_FOLDER'],
                        style_filename
                    )

                    content_image = Image.open(
                        content_path
                    ).convert('RGB')

                    style_image = Image.open(
                        style_path
                    ).convert('RGB')

                    alpha = min(
                        float(form.alpha.data),
                        0.75
                   )

                    stylized_image = style_transfer(
                        content_image,
                        style_image,
                        encoder,
                        decoder,
                        alpha,
                        device
             )
                    import gc

                    del content_image
                    del style_image
                    gc.collect()

                    result_filename = (
                        'stylized_' +
                        content_filename
                    )

                    result_path = os.path.join(
                        app.config['UPLOAD_FOLDER'],
                        result_filename
                    )

                    save_image(
                        stylized_image,
                        result_path
                    )

                    del stylized_image
                    gc.collect()

                    result_image = result_filename

                
                except Exception as e:
                    import traceback
                    error = traceback.format_exc()
                    print(error)

            else:
                error = "Upload both images"

    return render_template(
        'index.html',
        form=form,
        result_image=result_image,
        content_image=content_filename,
        style_image=style_filename,
        error=error
    )


@app.route('/uploads/<filename>')
def send_image(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)


@app.route('/content_data/<filename>')
def send_content_example(filename):
    return send_from_directory(
        'content_data',
        filename
    )


@app.route('/style_data/<filename>')
def send_style_example(filename):
    return send_from_directory(
        'style_data',
        filename
    )

@app.route('/experiment/<filename>')
def send_generated(filename):
    return send_from_directory(
        'experiment/big_dataset',
        filename
    )


if __name__ == '__main__':
    from werkzeug.serving import run_simple
    run_simple('localhost', 5000, app, use_reloader=True, use_debugger=True)