Create train.py

train.py

@@ -0,0 +1,41 @@
+import tensorflow as tf
+import numpy as np
+import cv2
+import os
+from tqdm import tqdm
+
+def load_dataset(dataset_path, image_size=(512, 512)):
+    images = []
+    for file in tqdm(os.listdir(dataset_path)):
+        img_path = os.path.join(dataset_path, file)
+        img = cv2.imread(img_path)
+        img = cv2.resize(img, image_size)
+        img = (img / 127.5) - 1.0  # Normalize
+        images.append(img)
+    return np.array(images)
+
+def build_generator():
+    inputs = tf.keras.layers.Input(shape=(512, 512, 3))
+    x = tf.keras.layers.Conv2D(64, (7, 7), padding="same", activation="relu")(inputs)
+    x = tf.keras.layers.Conv2D(128, (3, 3), strides=2, padding="same")(x)
+    x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)
+    x = tf.keras.layers.Conv2DTranspose(64, (3, 3), strides=2, padding="same")(x)
+    x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)
+    x = tf.keras.layers.Conv2D(3, (7, 7), activation="tanh", padding="same")(x)
+    return tf.keras.models.Model(inputs, x)
+
+def train_animegan(dataset_path, epochs=100, batch_size=8):
+    dataset = load_dataset(dataset_path)
+    generator = build_generator()
+    generator.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5), loss="mse")
+    for epoch in range(epochs):
+        for i in range(0, len(dataset), batch_size):
+            batch_images = dataset[i:i+batch_size]
+            noise = np.random.normal(0, 1, (batch_size, 512, 512, 3))
+            generator.train_on_batch(noise, batch_images)
+        print(f"Epoch {epoch+1}/{epochs} completed")
+    generator.save("AnimeGANv2_Hayao.h5")
+
+if __name__ == "__main__":
+    train_animegan("path/to/dataset")
+