Update app.py

app.py

@@ -3,28 +3,50 @@ import tensorflow as tf
 import numpy as np
 from PIL import Image
 
-# Load your Keras model
-model = tf.keras.models.load_model("NeuralStyleTransferModel.keras")
+# Load VGG for feature extraction
+vgg = tf.keras.applications.VGG19(include_top=False, weights="imagenet")
+vgg.trainable = False
 
-def style_transfer(content, style):
-    content = Image.fromarray(content).resize((256, 256))
-    style = Image.fromarray(style).resize((256, 256))
-    
-    # preprocess (example: normalize, expand dims, etc.)
-    content_arr = np.expand_dims(np.array(content) / 255.0, axis=0)
-    style_arr = np.expand_dims(np.array(style) / 255.0, axis=0)
+STYLE_LAYERS = [...]  # same layers as in your notebook
+CONTENT_LAYER = [...] # e.g. [('block5_conv4', 1)]
 
-    # run through your model
-    output = model.predict([content_arr, style_arr])
+def preprocess(img):
+    img = Image.fromarray(img).resize((256, 256))
+    arr = np.expand_dims(np.array(img) / 255.0, axis=0).astype(np.float32)
+    return tf.convert_to_tensor(arr)
 
-    # postprocess
-    output_img = (output[0] * 255).astype("uint8")
-    return Image.fromarray(output_img)
+def style_transfer(content, style, steps=50):
+    content_tensor = preprocess(content)
+    style_tensor = preprocess(style)
+
+    # Extract features
+    a_C = vgg(content_tensor)
+    a_S = vgg(style_tensor)
+
+    # Init generated image
+    generated_image = tf.Variable(content_tensor)
+
+    # Optimize
+    opt = tf.keras.optimizers.Adam(learning_rate=0.01)
+    for i in range(steps):
+        with tf.GradientTape() as tape:
+            a_G = vgg(generated_image)
+            J_style = compute_style_cost(a_S, a_G)
+            J_content = compute_content_cost(a_C, a_G)
+            J = total_cost(J_content, J_style, alpha=10, beta=40)
+
+        grad = tape.gradient(J, generated_image)
+        opt.apply_gradients([(grad, generated_image)])
+        generated_image.assign(tf.clip_by_value(generated_image, 0.0, 1.0))
+
+    out_img = (generated_image[0].numpy() * 255).astype("uint8")
+    return Image.fromarray(out_img)
 
 demo = gr.Interface(
     fn=style_transfer,
-    inputs=[gr.Image(type="numpy"), gr.Image(type="numpy")],
-    outputs=gr.Image(),
+    inputs=[gr.Image(type="numpy", label="Content Image"), 
+            gr.Image(type="numpy", label="Style Image")],
+    outputs=gr.Image(type="pil", label="Stylized Image"),
 )
 
 demo.launch()