"""
Image Style Transfer demo (TensorFlow Hub + Tkinter UI)
- Pick a content image (photo) and a style image (painting).
- Apply Magenta arbitrary-image-stylization-v1-256 model.
- Preview and save the stylized result.
Requirements:
pip install tensorflow tensorflow-hub pillow opencv-python
"""
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from PIL import Image, ImageTk
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import cv2
import os
# Model URL
TFHUB_MODEL_URL = "https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2"
# --------------------------
# Helper functions
# --------------------------
def load_img(path, max_dim=512):
"""Load image with PIL, resize keeping aspect ratio so long side = max_dim."""
img = Image.open(path).convert("RGB")
# resize
long = max(img.size)
if long > max_dim:
scale = max_dim / long
new_size = (int(img.size[0]*scale), int(img.size[1]*scale))
img = img.resize(new_size, Image.ANTIALIAS)
return img
def img_to_tensor(img):
"""PIL Image -> float32 tensor shape [1, H, W, 3] in [0,1]."""
arr = np.array(img).astype(np.float32) / 255.0
# add batch dim
return tf.expand_dims(arr, axis=0)
def tensor_to_pil(tensor):
"""Tensor [1,H,W,3] in [0,1] -> PIL Image."""
arr = tensor[0].numpy()
arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
return Image.fromarray(arr)
# --------------------------
# Load model once
# --------------------------
print("Loading style transfer model from TensorFlow Hub...")
model = hub.load(TFHUB_MODEL_URL)
print("Model loaded.")
# --------------------------
# Simple CLI function (optional)
# --------------------------
def stylize_image(content_img_path, style_img_path, output_path="stylized.png", max_dim=512):
content_img = load_img(content_img_path, max_dim=max_dim)
style_img = load_img(style_img_path, max_dim=max_dim)
content_t = img_to_tensor(content_img)
style_t = img_to_tensor(style_img)
# The hub model expects float tensors in [0,1]
outputs = model(tf.constant(content_t), tf.constant(style_t))
stylized = outputs[0]
pil = tensor_to_pil(stylized)
pil.save(output_path)
print(f"Saved stylized image to {output_path}")
return output_path
# --------------------------
# Tkinter GUI
# --------------------------
class StyleTransferGUI:
def __init__(self, root):
self.root = root
root.title("AI Image Style Transfer — RootRace")
root.geometry("1000x700")
self.content_path = None
self.style_path = None
self.result_image = None # PIL Image
# Controls frame
ctrl = ttk.Frame(root)
ctrl.pack(side=tk.TOP, fill=tk.X, padx=8, pady=8)
ttk.Button(ctrl, text="Choose Content Image", command=self.choose_content).pack(side=tk.LEFT, padx=4)
ttk.Button(ctrl, text="Choose Style Image", command=self.choose_style).pack(side=tk.LEFT, padx=4)
ttk.Button(ctrl, text="Apply Style Transfer", command=self.apply_style).pack(side=tk.LEFT, padx=8)
ttk.Button(ctrl, text="Save Result", command=self.save_result).pack(side=tk.LEFT, padx=4)
ttk.Button(ctrl, text="Quit", command=root.quit).pack(side=tk.RIGHT, padx=4)
# Preview frames
preview = ttk.Frame(root)
preview.pack(fill=tk.BOTH, expand=True)
# Left: content & style thumbnails
left = ttk.Frame(preview)
left.pack(side=tk.LEFT, fill=tk.Y, padx=6, pady=6)
ttk.Label(left, text="Content Image").pack()
self.content_label = ttk.Label(left)
self.content_label.pack(padx=6, pady=6)
ttk.Label(left, text="Style Image").pack()
self.style_label = ttk.Label(left)
self.style_label.pack(padx=6, pady=6)
# Right: result canvas
right = ttk.Frame(preview)
right.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=6, pady=6)
ttk.Label(right, text="Result").pack()
self.canvas = tk.Canvas(right, bg="black")
self.canvas.pack(fill=tk.BOTH, expand=True)
# Status
self.status = ttk.Label(root, text="Load a content and style image to start.", relief=tk.SUNKEN, anchor="w")
self.status.pack(side=tk.BOTTOM, fill=tk.X)
def choose_content(self):
path = filedialog.askopenfilename(filetypes=[("Images","*.jpg *.jpeg *.png *.bmp"),("All files","*.*")])
if not path:
return
self.content_path = path
img = load_img(path, max_dim=400)
self._set_thumbnail(self.content_label, img)
self.status.config(text=f"Selected content: {os.path.basename(path)}")
def choose_style(self):
path = filedialog.askopenfilename(filetypes=[("Images","*.jpg *.jpeg *.png *.bmp"),("All files","*.*")])
if not path:
return
self.style_path = path
img = load_img(path, max_dim=200)
self._set_thumbnail(self.style_label, img)
self.status.config(text=f"Selected style: {os.path.basename(path)}")
def _set_thumbnail(self, widget, pil_img):
tk_img = ImageTk.PhotoImage(pil_img)
widget.image = tk_img # keep ref
widget.config(image=tk_img)
def apply_style(self):
if not self.content_path or not self.style_path:
messagebox.showwarning("Missing images", "Please choose both content and style images.")
return
try:
self.status.config(text="Running style transfer... (this may take a few seconds)")
self.root.update_idletasks()
content_img = load_img(self.content_path, max_dim=512)
style_img = load_img(self.style_path, max_dim=512)
content_t = img_to_tensor(content_img)
style_t = img_to_tensor(style_img)
outputs = model(tf.constant(content_t), tf.constant(style_t))
stylized = outputs[0] # [1,H,W,3]
pil = tensor_to_pil(stylized)
self.result_image = pil
# display result on canvas, scaled to fit
self._display_result_on_canvas(pil)
self.status.config(text="Done. You can save the result.")
except Exception as e:
messagebox.showerror("Error", f"Style transfer failed: {e}")
self.status.config(text="Error during style transfer")
def _display_result_on_canvas(self, pil_img):
# resize to fit canvas while preserving aspect ratio
cw = self.canvas.winfo_width() or 600
ch = self.canvas.winfo_height() or 400
w,h = pil_img.size
scale = min(cw/w, ch/h, 1.0)
new_size = (int(w*scale), int(h*scale))
disp = pil_img.resize(new_size, Image.ANTIALIAS)
tk_img = ImageTk.PhotoImage(disp)
self.canvas.image = tk_img
self.canvas.delete("all")
# center image
x = (cw - new_size[0]) // 2
y = (ch - new_size[1]) // 2
self.canvas.create_image(x, y, anchor="nw", image=tk_img)
def save_result(self):
if self.result_image is None:
messagebox.showwarning("No result", "No stylized image to save. Apply style transfer first.")
return
path = filedialog.asksaveasfilename(defaultextension=".png", filetypes=[("PNG","*.png"),("JPEG","*.jpg")])
if not path:
return
self.result_image.save(path)
messagebox.showinfo("Saved", f"Saved stylized image to {path}")
self.status.config(text=f"Saved: {path}")
# --------------------------
# CLI usage entrypoint
# --------------------------
def main():
root = tk.Tk()
app = StyleTransferGUI(root)
root.mainloop()
if __name__ == "__main__":
main()
No comments:
Post a Comment