Spaces:
Build error
Build error
| import streamlit as st | |
| import os | |
| # Set the title of the app | |
| import torch | |
| import torch.nn as nn | |
| import torchvision.transforms as transforms | |
| from torchvision.models import vgg19,vgg16 | |
| from PIL import Image | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import os | |
| import torch.nn.functional as F | |
| import pandas as pd | |
| import torch.optim as optim | |
| import neural_style | |
| import argparse | |
| import os | |
| import sys | |
| import time | |
| import re | |
| # from PIL import ImageFilter, ImageResampling | |
| import numpy as np | |
| import torch | |
| from torch.optim import Adam | |
| from torch.utils.data import DataLoader | |
| from torchvision import datasets | |
| from torchvision import transforms | |
| import torch.onnx | |
| import utils | |
| from transformer_net import TransformerNet | |
| from vgg import Vgg16 | |
| from torch.utils.data import Dataset, DataLoader | |
| st.set_page_config(layout='wide') | |
| st.title("Neural Style transfer ") | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| def stylize_onnx(content_image, args): | |
| """ | |
| Read ONNX model and run it using onnxruntime | |
| """ | |
| # assert not args["export_onnx"] | |
| import onnxruntime | |
| ort_session = onnxruntime.InferenceSession(args["model"]) | |
| def to_numpy(tensor): | |
| return ( | |
| tensor.detach().cpu().numpy() | |
| if tensor.requires_grad | |
| else tensor.cpu().numpy() | |
| ) | |
| ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(content_image)} | |
| ort_outs = ort_session.run(None, ort_inputs) | |
| img_out_y = ort_outs[0] | |
| return torch.from_numpy(img_out_y) | |
| def stylize(device,args): | |
| content_image = utils.load_image(args['content_image'], scale=args['content_scale']) | |
| content_transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Lambda(lambda x: x.mul(255)) | |
| ]) | |
| content_image = content_transform(content_image) | |
| content_image = content_image.unsqueeze(0).to(device) | |
| if args['model'].endswith(".onnx"): | |
| output = stylize_onnx(content_image, args) | |
| else: | |
| with torch.no_grad(): | |
| style_model = TransformerNet() | |
| state_dict = torch.load( args['model']) | |
| # remove saved deprecated running_* keys in InstanceNorm from the checkpoint | |
| for k in list(state_dict.keys()): | |
| if re.search(r'in\d+\.running_(mean|var)$', k): | |
| del state_dict[k] | |
| style_model.load_state_dict(state_dict) | |
| style_model.to(device) | |
| style_model.eval() | |
| if args['export_onnx'] != False: | |
| assert args['export_onnx_name'].endswith(".onnx"), "Export model file should end with .onnx" | |
| output = torch.onnx.export( | |
| style_model, content_image, args['export_onnx_name'], opset_version=11, | |
| ) | |
| else: | |
| output = style_model(content_image).cpu() | |
| utils.save_image(args['output_image'], output[0]) | |
| return True | |
| # args = { | |
| # "content_image":"hemant_content3.jpg", | |
| # "content_scale":0.95, | |
| # "model":"epoch_100_Fri_Apr_19_00_29_00_2024_10_100000.model", | |
| # "export_onnx": False, | |
| # "export_onnx_name":"natural_art_style.onnx", | |
| # "output_image":"./kavitha/hemanth_style13.jpg" | |
| # } | |
| # stylize(device,args) | |
| def run_model(args): | |
| # Function to run the model | |
| # This is where you can place your model execution code | |
| st.write("Running the model...") | |
| # Simulate model running (you can replace this with your actual model code) | |
| import time | |
| result = stylize(device,args) # Simulate model running for 2 seconds | |
| st.success("Model has run successfully!") | |
| if result: | |
| st.subheader('Output Image') | |
| width = 400 | |
| height = 400 | |
| st.image(args["output_image"], caption='Final output',width=width) | |
| # if result: | |
| # width = 400 | |
| # height = 400 | |
| # st.image(args["output_image"], caption='Final output',width=width) | |
| # def app(): | |
| # # Streamlit app title | |
| # args = { | |
| # } | |
| # # Allow user to upload an image file | |
| # uploaded_file = st.file_uploader('Upload a JPEG image', type=['jpg', 'jpeg']) | |
| # # Check if a file has been uploaded | |
| # if uploaded_file is not None: | |
| # # Verify the uploaded file's format (extension) | |
| # file_extension = uploaded_file.name.split('.')[-1].lower() | |
| # if file_extension in ['jpg', 'jpeg']: | |
| # # Specify the directory path where the image will be stored | |
| # directory_path = 'C:\\Users\\Kavitha padala\\Desktop\\HemanthDL\\examples\\fast_neural_style\\neural_style\\images\\content_images' | |
| # # Create the directory if it doesn't exist | |
| # if not os.path.exists(directory_path): | |
| # os.makedirs(directory_path) | |
| # file_name = "content_image" | |
| # # Generate the full path for the image file | |
| # file_path = os.path.join(directory_path, file_name+"."+ file_extension) | |
| # # Save the uploaded image to the specified directory | |
| # with open(file_path, 'wb') as f: | |
| # f.write(uploaded_file.getbuffer()) | |
| # model_sets = ['women_style_art_model_best1.model', 'women_style_art_model_best1.model'] | |
| # # Create a select box to allow users to choose one model set | |
| # selected_model_set = st.selectbox( | |
| # 'Choose a model set:', # Label for the select box | |
| # model_sets # List of model sets | |
| # ) | |
| # output_image_name = "style_image.jpg" | |
| # args["content_image"] = file_path | |
| # args["model"] = selected_model_set | |
| # args["content_scale"] = 0.95 | |
| # args["output_image"] = "C:\\Users\\Kavitha padala\\Desktop\\HemanthDL\\examples\\fast_neural_style\\neural_style\\images\\style_images\\" + output_image_name | |
| # args["export_onnx"] = False | |
| # if st.button('Run Model'): | |
| # run_model(args) | |
| # # Display the selected model set | |
| # st.write(f'You selected: {selected_model_set}') | |
| # else: | |
| # # If the file is not in the correct format, display an error message | |
| # st.error('Please upload a JPG image file.') | |
| # if __name__ == "__main__": | |
| # app() | |
| # def display_image(): | |
| # def run_style_transfer(content_image, style_image, model): | |
| # # Placeholder function to run the style transfer model | |
| # # Replace this function with your actual model execution code | |
| # st.write(f"Running model: {model}") | |
| # # Simulate processing (you can replace this with your actual model code) | |
| # import time | |
| # time.sleep(3) | |
| # # For demonstration, we'll just return the style image as the output | |
| # # (You would return the styled image as the final output in your actual model code) | |
| # return style_image | |
| def app(): | |
| # Streamlit app title | |
| # Create a layout with three columns | |
| left_col, middle_col, right_col = st.columns([1, 1, 2]) | |
| args = {} | |
| result = False; | |
| # Middle column: Display title | |
| # Left column: Ask user to upload an image | |
| with left_col: | |
| st.subheader('Upload Image') | |
| uploaded_image = st.file_uploader('Choose an image', type=['jpg', 'jpeg']) | |
| if uploaded_image is not None: | |
| # Verify the uploaded file's format (extension) | |
| file_extension = uploaded_image.name.split('.')[-1].lower() | |
| if file_extension in ['jpg', 'jpeg']: | |
| # Specify the directory path where the image will be stored | |
| directory_path = 'uploads' | |
| # Create the directory if it doesn't exist | |
| if not os.path.exists(directory_path): | |
| os.makedirs(directory_path) | |
| file_name = "content_image" | |
| # Generate the full path for the image file | |
| file_path = os.path.join(directory_path, file_name+"."+ file_extension) | |
| args["content_image"] = file_path | |
| # Save the uploaded image to the specified directory | |
| with open(file_path, 'wb') as f: | |
| f.write(uploaded_image.getbuffer()) | |
| st.image(uploaded_image,width=256) | |
| # Right column top: Allow user to select a model and display style image | |
| with middle_col: | |
| # Define a list of models | |
| style_images = ['Art Style', 'Women Art','glass art'] | |
| st.subheader('Style model') | |
| # Create a select box to choose a model | |
| selected_image = st.selectbox('Choose a model will be selected :', style_images) | |
| models_dict = { | |
| 'Art Style':'art_style.model', | |
| 'Women Art':'women_style.model', | |
| 'glass art':'glass_art.model' | |
| } | |
| output_image_name = "style_image.jpg" | |
| args["model"] = models_dict[selected_image] | |
| args["content_scale"] = 0.95 | |
| args["output_image"] = output_image_name | |
| args["export_onnx"] = False | |
| # Left column: After clicking "Run Model" button, display the final output style transfer image | |
| with right_col: | |
| if uploaded_image is not None and selected_image in models_dict: | |
| # Get the selected style image | |
| # style_image_path = style_images[selected_model] | |
| # # Load the uploaded image and style image | |
| # content_image = uploaded_image | |
| # style_image = style_image_path | |
| if st.button('Run Model'): | |
| result = run_model(args) | |
| # Create a button to run the model | |
| # if st.button('Run Model'): | |
| # # Run style transfer model | |
| # final_output = run_style_transfer(content_image, style_image, selected_model) | |
| # # Resize the final output image | |
| # # You can specify the desired width and height for the resized image | |
| # width = 400 # Replace with your desired width | |
| # height = 400 # Replace with your desired height | |
| # # Display the final output image with specified width and height | |
| # st.subheader('Final Output') | |
| # st.image(final_output, caption='Styled Image', width=width, height=height) | |
| if __name__ == "__main__": | |
| app() | |