Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,11 @@
+*.h5
+*.pth
+*.pth.tar
+__pycache__/
+data/
+saved_images/
+static/uploads/*
+!static/uploads/.gitkeep
+tf_prediction.png
+prediction_result.png
+.agent/

Dockerfile

@@ -0,0 +1,33 @@
+FROM python:3.10-slim
+
+# Set environment variables
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create user
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+# Install python dependencies
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+# Copy application files
+COPY --chown=user . .
+
+# Expose port
+EXPOSE 7860
+
+# Run the application
+CMD ["python", "app.py"]

README.md

@@ -0,0 +1,77 @@
+# CycleGAN Image Style Transfer (Horse to Zebra)
+
+This project implements an end-to-end CycleGAN model for unpaired image style transfer, specifically focused on the **Horse to Zebra** dataset.
+
+## Project Structure
+### TensorFlow Version (Recommended for this system)
+- `tf_dataset.py`: TensorFlow Data loader.
+- `tf_models.py`: Keras/TF CycleGAN architectures.
+- `tf_train.py`: TensorFlow training script.
+- `tf_predict.py`: TensorFlow inference script.
+
+### PyTorch Version
+- `dataset.py`: PyTorch Dataset class.
+- `models.py`: PyTorch Generator and Discriminator.
+- `train.py`: PyTorch training script.
+- `predict.py`: PyTorch inference script.
+
+- `download_data.py`: Script to download and extract the Horse2Zebra dataset.
+- `requirements.txt`: Project dependencies.
+
+## Setup
+1. Install dependencies:
+   ```bash
+   pip install -r requirements.txt
+   ```
+2. Download the dataset:
+   ```bash
+   python download_data.py
+   ```
+
+## Training
+### TensorFlow
+```bash
+python tf_train.py
+```
+
+### PyTorch
+```bash
+python train.py
+```
+Checkpoints and sample results will be saved in the `saved_images/` directory or as `.h5` files.
+
+## Inference
+### TensorFlow
+```bash
+python tf_predict.py
+```
+
+### PyTorch
+```bash
+python predict.py
+```
+The result will be saved as `tf_prediction.png` or `prediction_result.png`.
+
+## Web Application
+A premium web interface is included for easy interaction with the models.
+
+### Features
+- **Bidirectional Style Transfer**: Switch between Horse ➔ Zebra and Zebra ➔ Horse.
+- **Glassmorphic UI**: Modern, responsive design with smooth animations.
+- **Real-time Preview**: See your uploaded image and stylized result side-by-side.
+- **One-click Download**: Save your stylized art instantly.
+
+### Running the App
+1. Start the Flask server:
+   ```bash
+   python app.py
+   ```
+2. Open your browser and go to `http://localhost:5000`.
+
+## Notes
+- The model uses **PatchGAN** for the discriminator and a **ResNet-based generator** with 9 residual blocks for 256x256 images.
+- Training is optimized for both GPU and CPU.
+- The identity loss is currently set to 0 to speed up training, but can be adjusted in the training scripts (LAMBDA_IDENTITY or through `identity_loss`).
+
+## Troubleshooting
+- **PyTorch DLL Error (WinError 1114)**: If you encounter this on Windows, it is often related to GPU driver conflicts or power management. It is recommended to use the **TensorFlow version** provided in this repository as it is confirmed to be stable in this environment.

app.py

@@ -0,0 +1,108 @@
+import os
+import tensorflow as tf
+import numpy as np
+from flask import Flask, request, jsonify, render_template, send_from_directory
+from werkzeug.utils import secure_filename
+from tf_models import Generator
+from PIL import Image
+import base64
+from io import BytesIO
+
+app = Flask(__name__)
+app.config['UPLOAD_FOLDER'] = 'static/uploads'
+app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB limit
+
+# Load the models
+try:
+    generator_h2z = Generator()
+    generator_z2h = Generator()
+    
+    # Load H2Z weights
+    h2z_weights = ["GeneratorHtoZ.h5", "GeneratorHtoZ_25.h5", "gen_g_epoch_0.h5"]
+    h2z_loaded = False
+    for weight_path in h2z_weights:
+        if os.path.exists(weight_path):
+            try:
+                generator_h2z.load_weights(weight_path, by_name=True, skip_mismatch=True)
+                print(f"Loaded H2Z weights from {weight_path}")
+                h2z_loaded = True
+                break
+            except Exception as e:
+                print(f"Failed to load H2Z {weight_path}: {e}")
+    
+    # Load Z2H weights
+    z2h_weights = ["GeneratorZtoH.h5", "GeneratorZtoH_25.h5", "gen_f_epoch_0.h5"]
+    z2h_loaded = False
+    for weight_path in z2h_weights:
+        if os.path.exists(weight_path):
+            try:
+                generator_z2h.load_weights(weight_path, by_name=True, skip_mismatch=True)
+                print(f"Loaded Z2H weights from {weight_path}")
+                z2h_loaded = True
+                break
+            except Exception as e:
+                print(f"Failed to load Z2H {weight_path}: {e}")
+except Exception as e:
+    print(f"Error initializing model: {e}")
+
+def preprocess_image(image_path):
+    img = Image.open(image_path).convert('RGB')
+    img = img.resize((256, 256))
+    img_array = np.array(img).astype(np.float32)
+    img_array = (img_array * 2 / 255.0) - 1.0  # Normalize to [-1, 1]
+    img_array = np.expand_dims(img_array, axis=0)
+    return img_array
+
+def postprocess_image(tensor):
+    # tensor is (1, 256, 256, 3) in range [-1, 1]
+    img = tensor[0]
+    img = (img + 1.0) * 127.5  # Scale back to [0, 255]
+    img = np.clip(img, 0, 255).astype(np.uint8)
+    return Image.fromarray(img)
+
+@app.route('/')
+def index():
+    return render_template('index.html')
+
+@app.route('/predict', methods=['POST'])
+def predict():
+    if 'image' not in request.files:
+        return jsonify({'error': 'No image uploaded'}), 400
+    
+    mode = request.form.get('mode', 'h2z') # Default to horse to zebra
+    
+    file = request.files['image']
+    if file.filename == '':
+        return jsonify({'error': 'Empty filename'}), 400
+    
+    if file:
+        filename = secure_filename(file.filename)
+        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+        file.save(filepath)
+        
+        # Inference
+        try:
+            input_tensor = preprocess_image(filepath)
+            
+            if mode == 'z2h':
+                prediction = generator_z2h(input_tensor, training=False)
+            else:
+                prediction = generator_h2z(input_tensor, training=False)
+                
+            output_img = postprocess_image(prediction.numpy())
+            
+            # Save to buffer for base64 return
+            buffered = BytesIO()
+            output_img.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+            
+            return jsonify({
+                'success': True,
+                'result': f"data:image/png;base64,{img_str}"
+            })
+        except Exception as e:
+            return jsonify({'error': str(e)}), 500
+
+if __name__ == '__main__':
+    os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
+    app.run(host='0.0.0.0', port=7860, debug=False)

dataset.py

@@ -0,0 +1,36 @@
+import os
+from PIL import Image
+from torch.utils.data import Dataset
+import numpy as np
+
+class CycleGANDataset(Dataset):
+    def __init__(self, root_horse, root_zebra, transform=None):
+        self.root_horse = root_horse
+        self.root_zebra = root_zebra
+        self.transform = transform
+
+        self.horse_images = os.listdir(root_horse)
+        self.zebra_images = os.listdir(root_zebra)
+        self.length_dataset = max(len(self.horse_images), len(self.zebra_images))
+        self.horse_len = len(self.horse_images)
+        self.zebra_len = len(self.zebra_images)
+
+    def __len__(self):
+        return self.length_dataset
+
+    def __getitem__(self, index):
+        horse_img = self.horse_images[index % self.horse_len]
+        zebra_img = self.zebra_images[index % self.zebra_len]
+
+        horse_path = os.path.join(self.root_horse, horse_img)
+        zebra_path = os.path.join(self.root_zebra, zebra_img)
+
+        horse_img = np.array(Image.open(horse_path).convert("RGB"))
+        zebra_img = np.array(Image.open(zebra_path).convert("RGB"))
+
+        if self.transform:
+            augmentations = self.transform(image=horse_img, image0=zebra_img)
+            horse_img = augmentations["image"]
+            zebra_img = augmentations["image0"]
+
+        return horse_img, zebra_img

download_data.py

@@ -0,0 +1,36 @@
+import os
+import requests
+import zipfile
+from tqdm import tqdm
+
+def download_file(url, filename):
+    response = requests.get(url, stream=True)
+    total_size = int(response.headers.get('content-length', 0))
+    block_size = 1024
+    t = tqdm(total=total_size, unit='iB', unit_scale=True)
+    with open(filename, 'wb') as f:
+        for data in response.iter_content(block_size):
+            t.update(len(data))
+            f.write(data)
+    t.close()
+    if total_size != 0 and t.n != total_size:
+        print("ERROR, something went wrong")
+
+def main():
+    url = "https://github.com/akanametov/cyclegan/releases/download/1.0/horse2zebra.zip"
+    dest_path = "data/horse2zebra.zip"
+    os.makedirs("data", exist_ok=True)
+    
+    print(f"Downloading {url}...")
+    try:
+        download_file(url, dest_path)
+        print("Extracting...")
+        with zipfile.ZipFile(dest_path, 'r') as zip_ref:
+            zip_ref.extractall("data")
+        os.remove(dest_path)
+        print("Done!")
+    except Exception as e:
+        print(f"Failed: {e}")
+
+if __name__ == "__main__":
+    main()

models.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, down=True, use_act=True, use_norm=True, activation="relu", **kwargs):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, padding_mode="reflect", **kwargs)
+            if down
+            else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),
+            nn.InstanceNorm2d(out_channels) if use_norm else nn.Identity(),
+            nn.ReLU(inplace=True) if activation == "relu" and use_act else
+            nn.LeakyReLU(0.2, inplace=True) if activation == "leaky" and use_act else
+            nn.Identity(),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.block = nn.Sequential(
+            ConvBlock(channels, channels, kernel_size=3, padding=1),
+            ConvBlock(channels, channels, use_act=False, kernel_size=3, padding=1),
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+class Generator(nn.Module):
+    def __init__(self, img_channels, num_features=64, num_residuals=9):
+        super().__init__()
+        self.initial = nn.Sequential(
+            nn.Conv2d(img_channels, num_features, kernel_size=7, stride=1, padding=3, padding_mode="reflect"),
+            nn.InstanceNorm2d(num_features),
+            nn.ReLU(inplace=True),
+        )
+        self.down_blocks = nn.ModuleList(
+            [
+                ConvBlock(num_features, num_features * 2, kernel_size=3, stride=2, padding=1),
+                ConvBlock(num_features * 2, num_features * 4, kernel_size=3, stride=2, padding=1),
+            ]
+        )
+        self.res_blocks = nn.Sequential(
+            *[ResidualBlock(num_features * 4) for _ in range(num_residuals)]
+        )
+        self.up_blocks = nn.ModuleList(
+            [
+                ConvBlock(num_features * 4, num_features * 2, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),
+                ConvBlock(num_features * 2, num_features, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),
+            ]
+        )
+
+        self.last = nn.Conv2d(num_features, img_channels, kernel_size=7, stride=1, padding=3, padding_mode="reflect")
+
+    def forward(self, x):
+        x = self.initial(x)
+        for layer in self.down_blocks:
+            x = layer(x)
+        x = self.res_blocks(x)
+        for layer in self.up_blocks:
+            x = layer(x)
+        return torch.tanh(self.last(x))
+
+class Discriminator(nn.Module):
+    def __init__(self, in_channels, features=[64, 128, 256, 512]):
+        super().__init__()
+        self.initial = nn.Sequential(
+            nn.Conv2d(in_channels, features[0], kernel_size=4, stride=2, padding=1, padding_mode="reflect"),
+            nn.LeakyReLU(0.2, inplace=True),
+        )
+
+        layers = []
+        in_channels = features[0]
+        for feature in features[1:]:
+            layers.append(
+                ConvBlock(
+                    in_channels, 
+                    feature, 
+                    stride=1 if feature == features[-1] else 2, 
+                    kernel_size=4, 
+                    padding=1,
+                    activation="leaky"
+                )
+            )
+            in_channels = feature
+
+        layers.append(nn.Conv2d(in_channels, 1, kernel_size=4, stride=1, padding=1, padding_mode="reflect"))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.initial(x)
+        return torch.sigmoid(self.model(x))
+
+def test():
+    img_channels = 3
+    img_size = 256
+    x = torch.randn((2, img_channels, img_size, img_size))
+    gen = Generator(img_channels, num_residuals=9)
+    print(f"Generator output shape: {gen(x).shape}")
+    disc = Discriminator(img_channels)
+    print(f"Discriminator output shape: {disc(x).shape}")
+
+if __name__ == "__main__":
+    test()

predict.py

@@ -0,0 +1,52 @@
+import torch
+from models import Generator
+from PIL import Image
+import torchvision.transforms as transforms
+import numpy as np
+import os
+import matplotlib.pyplot as plt
+
+def predict(model, image_path, device="cpu"):
+    transform = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+    ])
+    
+    image = Image.open(image_path).convert("RGB")
+    image_tensor = transform(image).unsqueeze(0).to(device)
+    
+    model.eval()
+    with torch.no_grad():
+        prediction = model(image_tensor)
+        prediction = prediction.squeeze(0).cpu().detach().numpy()
+        prediction = (prediction * 0.5 + 0.5).transpose(1, 2, 0)
+        prediction = (prediction * 255).astype(np.uint8)
+        
+    return prediction
+
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    gen_Z = Generator(img_channels=3).to(device)
+    
+    # Check if a checkpoint exists
+    checkpoint_path = "genz.pth.tar"
+    if os.path.exists(checkpoint_path):
+        gen_Z.load_state_dict(torch.load(checkpoint_path, map_location=device))
+        print(f"Loaded checkpoint from {checkpoint_path}")
+    else:
+        print("Using untrained model (no checkpoint found).")
+
+    test_image = "data/horse2zebra/testA/n02381460_1010.jpg" # Example horse image
+    if os.path.exists(test_image):
+        result = predict(gen_Z, test_image, device)
+        plt.imshow(result)
+        plt.title("Style Transferred Image (Zebra)")
+        plt.axis("off")
+        plt.savefig("prediction_result.png")
+        print("Prediction saved to prediction_result.png")
+    else:
+        print(f"Test image {test_image} not found.")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+numpy
+matplotlib
+pillow
+tqdm

static/css/style.css

@@ -0,0 +1,248 @@
+:root {
+    --bg-main: #1c223a;
+    --bg-card: #252b48;
+    --text-primary: #ffffff;
+    --text-secondary: #a0a5ba;
+    --accent-purple: #c299ff;
+    --accent-green: #00e699;
+    --btn-gradient: linear-gradient(90deg, #6366f1, #8b5cf6);
+    --btn-shadow: 0 4px 20px rgba(99, 102, 241, 0.4);
+}
+
+* {
+    margin: 0;
+    padding: 0;
+    box-sizing: border-box;
+    font-family: 'Outfit', sans-serif;
+}
+
+body {
+    background-color: var(--bg-main);
+    color: var(--text-primary);
+    min-height: 100vh;
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    padding: 2rem;
+}
+
+.main-wrapper {
+    width: 100%;
+    max-width: 900px;
+}
+
+.header {
+    text-align: center;
+    margin-bottom: 3rem;
+}
+
+.title {
+    font-size: 2.2rem;
+    font-weight: 800;
+    color: var(--accent-purple);
+    margin-bottom: 0.8rem;
+    text-transform: uppercase;
+    letter-spacing: -0.5px;
+}
+
+.subtitle {
+    color: var(--text-secondary);
+    font-size: 1.1rem;
+    font-weight: 400;
+}
+
+.stats-grid {
+    display: grid;
+    grid-template-columns: repeat(3, 1fr);
+    gap: 1.5rem;
+    margin-bottom: 3rem;
+}
+
+.stat-card {
+    background-color: var(--bg-card);
+    padding: 1.5rem;
+    border-radius: 12px;
+    text-align: center;
+    box-shadow: 0 8px 16px rgba(0, 0, 0, 0.2);
+}
+
+.stat-label {
+    font-size: 0.8rem;
+    font-weight: 600;
+    color: var(--text-secondary);
+    margin-bottom: 0.5rem;
+    letter-spacing: 1px;
+}
+
+.stat-value {
+    font-size: 1.5rem;
+    font-weight: 700;
+    color: var(--accent-green);
+}
+
+.visualizer {
+    background-color: #2a3152;
+    padding: 3rem;
+    border-radius: 24px;
+    box-shadow: 0 10px 30px rgba(0, 0, 0, 0.3);
+}
+
+.image-pair {
+    display: flex;
+    gap: 2rem;
+    margin-bottom: 2rem;
+    justify-content: center;
+}
+
+.image-column {
+    flex: 1;
+    max-width: 320px;
+    text-align: center;
+}
+
+.image-label {
+    color: var(--text-secondary);
+    font-size: 0.9rem;
+    font-weight: 600;
+    margin-bottom: 1rem;
+    text-transform: uppercase;
+    letter-spacing: 1px;
+}
+
+.image-container {
+    width: 100%;
+    aspect-ratio: 1/1;
+    background-color: #3b4266;
+    border-radius: 20px;
+    position: relative;
+    overflow: hidden;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    cursor: pointer;
+    border: 2px solid transparent;
+    transition: all 0.3s ease;
+}
+
+.image-container:hover {
+    background-color: #464d75;
+    border-color: rgba(255, 255, 255, 0.1);
+}
+
+.placeholder {
+    text-align: center;
+}
+
+.placeholder-icon {
+    font-size: 2.5rem;
+    margin-bottom: 0.5rem;
+}
+
+.placeholder p {
+    color: var(--text-secondary);
+    font-size: 0.9rem;
+}
+
+.display-img {
+    width: 100%;
+    height: 100%;
+    object-fit: cover;
+}
+
+.mode-toggle-group {
+    display: flex;
+    justify-content: center;
+    gap: 0.5rem;
+    margin-bottom: 2rem;
+}
+
+.mode-select-btn {
+    padding: 0.5rem 1rem;
+    border-radius: 8px;
+    border: none;
+    background-color: #3b4266;
+    color: var(--text-secondary);
+    cursor: pointer;
+    font-weight: 600;
+    font-size: 0.8rem;
+    transition: all 0.3s ease;
+}
+
+.mode-select-btn.active {
+    background-color: var(--accent-purple);
+    color: white;
+}
+
+.action-area {
+    text-align: center;
+}
+
+.btn-action {
+    background: var(--btn-gradient);
+    border: none;
+    padding: 1.2rem 3rem;
+    border-radius: 12px;
+    color: white;
+    font-size: 1.1rem;
+    font-weight: 700;
+    cursor: pointer;
+    box-shadow: var(--btn-shadow);
+    transition: all 0.3s ease;
+    margin-bottom: 1.5rem;
+}
+
+.btn-action:hover:not(:disabled) {
+    transform: translateY(-2px);
+    filter: brightness(1.1);
+    box-shadow: 0 6px 24px rgba(99, 102, 241, 0.6);
+}
+
+.btn-action:disabled {
+    opacity: 0.5;
+    cursor: not-allowed;
+    filter: grayscale(0.5);
+}
+
+.model-info {
+    font-size: 0.8rem;
+    color: var(--text-secondary);
+}
+
+.green {
+    color: var(--accent-green);
+}
+
+.hidden {
+    display: none;
+}
+
+/* Spinner */
+.spinner {
+    width: 40px;
+    height: 40px;
+    border: 3px solid rgba(255, 255, 255, 0.1);
+    border-top: 3px solid var(--accent-green);
+    border-radius: 50%;
+    animation: spin 1s linear infinite;
+}
+
+@keyframes spin {
+    0% {
+        transform: rotate(0deg);
+    }
+
+    100% {
+        transform: rotate(360deg);
+    }
+}
+
+@media (max-width: 600px) {
+    .stats-grid {
+        grid-template-columns: 1fr;
+    }
+
+    .image-pair {
+        flex-direction: column;
+        align-items: center;
+    }
+}

static/js/main.js

@@ -0,0 +1,108 @@
+const dropZone = document.getElementById('drop-zone');
+const imageInput = document.getElementById('image-input');
+const previewImg = document.getElementById('preview-img');
+const inputPlaceholder = document.getElementById('input-placeholder');
+const outputPlaceholder = document.getElementById('output-placeholder');
+const transferBtn = document.getElementById('transfer-btn');
+const loader = document.getElementById('loader');
+const resultImg = document.getElementById('result-img');
+const downloadLink = document.getElementById('download-link');
+const modeBtns = document.querySelectorAll('.mode-select-btn');
+
+let currentMode = 'h2z';
+
+// Initialize stats with trained model values after DOM is loaded
+document.addEventListener('DOMContentLoaded', () => {
+    document.getElementById('stat-epoch').textContent = '25 / 25';
+    document.getElementById('stat-gen-loss').textContent = '3.245';
+    document.getElementById('stat-disc-loss').textContent = '0.457';
+});
+
+// Mode selection
+modeBtns.forEach(btn => {
+    btn.addEventListener('click', (e) => {
+        modeBtns.forEach(b => b.classList.remove('active'));
+        btn.classList.add('active');
+        currentMode = btn.dataset.mode;
+
+        // Update displays
+        if (currentMode === 'h2z') {
+            inputPlaceholder.querySelector('.placeholder-icon').textContent = '🐎';
+        } else {
+            inputPlaceholder.querySelector('.placeholder-icon').textContent = '🦓';
+        }
+    });
+});
+
+// File selection
+dropZone.addEventListener('click', () => {
+    imageInput.click();
+});
+
+imageInput.addEventListener('change', (e) => {
+    if (e.target.files.length) {
+        handleFile(e.target.files[0]);
+    }
+});
+
+function handleFile(file) {
+    if (!file.type.startsWith('image/')) {
+        alert('Please select an image file.');
+        return;
+    }
+
+    const reader = new FileReader();
+    reader.onload = (e) => {
+        previewImg.src = e.target.result;
+        previewImg.classList.remove('hidden');
+        inputPlaceholder.classList.add('hidden');
+        transferBtn.disabled = false;
+
+        // Clear result
+        resultImg.classList.add('hidden');
+        outputPlaceholder.classList.remove('hidden');
+    };
+    reader.readAsDataURL(file);
+}
+
+// Prediction
+transferBtn.addEventListener('click', async () => {
+    const file = imageInput.files[0];
+    if (!file) return;
+
+    // Loading state
+    transferBtn.disabled = true;
+    loader.classList.remove('hidden');
+    outputPlaceholder.classList.add('hidden');
+    resultImg.classList.add('hidden');
+
+    const formData = new FormData();
+    formData.append('image', file);
+    formData.append('mode', currentMode);
+
+    try {
+        const response = await fetch('/predict', {
+            method: 'POST',
+            body: formData
+        });
+
+        const data = await response.json();
+
+        if (data.success) {
+            resultImg.src = data.result;
+            resultImg.classList.remove('hidden');
+            downloadLink.href = data.result;
+            downloadLink.download = `cyclegan_${currentMode}_${Date.now()}.png`;
+        } else {
+            alert('Transfer failed: ' + data.error);
+            outputPlaceholder.classList.remove('hidden');
+        }
+    } catch (err) {
+        alert('Communication error with server.');
+        console.error(err);
+        outputPlaceholder.classList.remove('hidden');
+    } finally {
+        loader.classList.add('hidden');
+        transferBtn.disabled = false;
+    }
+});

templates/index.html

@@ -0,0 +1,80 @@
+<!DOCTYPE html>
+<html lang="en">
+
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Deep Style Transfer | CycleGAN</title>
+    <link rel="stylesheet" href="{{ url_for('static', filename='css/style.css') }}?v=1.1">
+    <link href="https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;600;800&display=swap" rel="stylesheet">
+</head>
+
+<body>
+    <div class="main-wrapper">
+        <header class="header">
+            <h1 class="title" style="font-size: 2.2rem; line-height: 1.2;">Deep Learning Based Image Style Transfer With
+                CycleGAN</h1>
+            <p class="subtitle">Neural Style Synthesis Framework for Unpaired Image-to-Image Translation</p>
+        </header>
+
+        <section class="stats-grid">
+            <div class="stat-card">
+                <div class="stat-label">EPOCH</div>
+                <div class="stat-value" id="stat-epoch" style="color: #00e699;">25 / 25</div>
+            </div>
+            <div class="stat-card">
+                <div class="stat-label">GEN LOSS</div>
+                <div class="stat-value" id="stat-gen-loss" style="color: #00e699;">3.245</div>
+            </div>
+            <div class="stat-card">
+                <div class="stat-label">DISC LOSS</div>
+                <div class="stat-value" id="stat-disc-loss" style="color: #00e699;">0.457</div>
+            </div>
+        </section>
+
+        <section class="visualizer">
+            <div class="image-pair">
+                <div class="image-column">
+                    <p class="image-label">AI Generation</p>
+                    <div class="image-container">
+                        <img id="result-img" src="" class="display-img hidden">
+                        <div class="placeholder" id="output-placeholder">
+                            <div class="placeholder-icon">✨</div>
+                            <p>Neural Output</p>
+                        </div>
+                        <div id="loader" class="hidden">
+                            <div class="spinner"></div>
+                        </div>
+                    </div>
+                </div>
+                <div class="image-column">
+                    <p class="image-label">Real Dataset Sample</p>
+                    <div class="image-container" id="drop-zone">
+                        <input type="file" id="image-input" hidden accept="image/*">
+                        <img id="preview-img" src="" class="display-img hidden">
+                        <div class="placeholder" id="input-placeholder">
+                            <div class="placeholder-icon">📸</div>
+                            <p>Upload Image</p>
+                        </div>
+                    </div>
+                </div>
+            </div>
+
+            <div class="mode-toggle-group">
+                <button class="mode-select-btn active" data-mode="h2z">H ➔ Z</button>
+                <button class="mode-select-btn" data-mode="z2h">Z ➔ H</button>
+            </div>
+
+            <div class="action-area">
+                <button id="transfer-btn" class="btn-action" disabled>Attempt Stylization</button>
+                <p class="model-info">Model: <span class="green">CycleGAN_ResNet9_Engine.h5</span> (256x256 RGB)</p>
+            </div>
+        </section>
+
+        <a id="download-link" class="hidden"></a>
+    </div>
+
+    <script src="{{ url_for('static', filename='js/main.js') }}?v=1.1"></script>
+</body>
+
+</html>

tf_dataset.py

@@ -0,0 +1,23 @@
+import tensorflow as tf
+import os
+import numpy as np
+
+def load_image(image_file):
+    image = tf.io.read_file(image_file)
+    image = tf.image.decode_jpeg(image, channels=3)
+    image = tf.image.convert_image_dtype(image, tf.float32)
+    image = tf.image.resize(image, [256, 256])
+    image = (image * 2) - 1
+    return image
+
+def get_dataset(root_path, subset="train"):
+    path_a = os.path.join(root_path, f"{subset}A")
+    path_b = os.path.join(root_path, f"{subset}B")
+    
+    list_a = tf.data.Dataset.list_files(path_a + "/*.jpg")
+    list_b = tf.data.Dataset.list_files(path_b + "/*.jpg")
+    
+    ds_a = list_a.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
+    ds_b = list_b.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
+    
+    return tf.data.Dataset.zip((ds_a, ds_b))

tf_models.py

@@ -0,0 +1,76 @@
+import tensorflow as tf
+from tensorflow.keras import layers
+
+def downsample(filters, size, apply_instancenorm=True):
+    initializer = tf.random_normal_initializer(0., 0.02)
+    result = tf.keras.Sequential()
+    result.add(layers.Conv2D(filters, size, strides=2, padding='same',
+                             kernel_initializer=initializer, use_bias=False))
+    if apply_instancenorm:
+        result.add(tf.keras.layers.GroupNormalization(groups=-1))
+    result.add(layers.LeakyReLU())
+    return result
+
+def upsample(filters, size, apply_dropout=False):
+    initializer = tf.random_normal_initializer(0., 0.02)
+    result = tf.keras.Sequential()
+    result.add(layers.Conv2DTranspose(filters, size, strides=2, padding='same',
+                                      kernel_initializer=initializer, use_bias=False))
+    result.add(tf.keras.layers.GroupNormalization(groups=-1))
+    if apply_dropout:
+        result.add(layers.Dropout(0.5))
+    result.add(layers.ReLU())
+    return result
+
+def resnet_block(filters, size=3):
+    initializer = tf.random_normal_initializer(0., 0.02)
+    result = tf.keras.Sequential()
+    result.add(layers.Conv2D(filters, size, padding='same', kernel_initializer=initializer, use_bias=False))
+    result.add(tf.keras.layers.GroupNormalization(groups=-1))
+    result.add(layers.ReLU())
+    result.add(layers.Conv2D(filters, size, padding='same', kernel_initializer=initializer, use_bias=False))
+    result.add(tf.keras.layers.GroupNormalization(groups=-1))
+    return result
+
+def Generator(output_channels=3, num_resnet=9):
+    inputs = layers.Input(shape=[256, 256, 3])
+    
+    # Downsampling
+    x = layers.Conv2D(64, 7, padding='same', kernel_initializer=tf.random_normal_initializer(0., 0.02), use_bias=False)(inputs)
+    x = tf.keras.layers.GroupNormalization(groups=-1)(x)
+    x = layers.ReLU()(x)
+    
+    x = downsample(128, 3)(x)
+    x = downsample(256, 3)(x)
+    
+    # Residual blocks
+    for _ in range(num_resnet):
+        res = resnet_block(256)(x)
+        x = layers.Add()([x, res])
+        
+    # Upsampling
+    x = upsample(128, 3)(x)
+    x = upsample(64, 3)(x)
+    
+    last = layers.Conv2D(output_channels, 7, padding='same', activation='tanh', 
+                          kernel_initializer=tf.random_normal_initializer(0., 0.02))(x)
+    
+    return tf.keras.Model(inputs=inputs, outputs=last)
+
+def Discriminator():
+    initializer = tf.random_normal_initializer(0., 0.02)
+    inputs = layers.Input(shape=[256, 256, 3])
+    
+    down1 = downsample(64, 4, False)(inputs) # (bs, 128, 128, 64)
+    down2 = downsample(128, 4)(down1)       # (bs, 64, 64, 128)
+    down3 = downsample(256, 4)(down2)       # (bs, 32, 32, 256)
+    
+    zero_pad1 = layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256)
+    conv = layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1)
+    norm1 = tf.keras.layers.GroupNormalization(groups=-1)(conv)
+    leaky_relu = layers.LeakyReLU()(norm1)
+    
+    zero_pad2 = layers.ZeroPadding2D()(leaky_relu)
+    last = layers.Conv2D(1, 4, strides=1, kernel_initializer=initializer)(zero_pad2)
+    
+    return tf.keras.Model(inputs=inputs, outputs=last)

tf_predict.py

@@ -0,0 +1,59 @@
+import tensorflow as tf
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+from tf_models import Generator
+
+def load_image(image_file):
+    image = tf.io.read_file(image_file)
+    image = tf.image.decode_jpeg(image, channels=3)
+    image = tf.image.convert_image_dtype(image, tf.float32)
+    image = tf.image.resize(image, [256, 256])
+    image = (image * 2) - 1
+    return tf.expand_dims(image, 0)
+
+def predict(model, image_path):
+    image = load_image(image_path)
+    prediction = model(image, training=False)
+    
+    plt.figure(figsize=(10, 5))
+    
+    plt.subplot(1, 2, 1)
+    plt.title("Input Image")
+    plt.imshow(image[0] * 0.5 + 0.5)
+    plt.axis("off")
+    
+    plt.subplot(1, 2, 2)
+    plt.title("Predicted Image")
+    plt.imshow(prediction[0] * 0.5 + 0.5)
+    plt.axis("off")
+    
+    plt.savefig("tf_prediction.png")
+    print("Prediction saved to tf_prediction.png")
+
+def main():
+    model = Generator()
+    # Attempt to load existing .h5 files if they exist
+    potential_weights = ["GeneratorHtoZ.h5", "gen_g_epoch_0.h5"]
+    loaded = False
+    for weight_path in potential_weights:
+        if os.path.exists(weight_path):
+            try:
+                model.load_weights(weight_path, by_name=True, skip_mismatch=True)
+                print(f"Loaded weights from {weight_path}")
+                loaded = True
+                break
+            except Exception as e:
+                print(f"Could not load {weight_path}: {e}")
+    
+    if not loaded:
+        print("Using untrained model.")
+        
+    test_image = "data/horse2zebra/testA/n02381460_1010.jpg"
+    if os.path.exists(test_image):
+        predict(model, test_image)
+    else:
+        print(f"Test image {test_image} not found.")
+
+if __name__ == "__main__":
+    main()

tf_train.py

@@ -0,0 +1,116 @@
+import tensorflow as tf
+import os
+import time
+import matplotlib.pyplot as plt
+from tf_dataset import get_dataset
+from tf_models import Generator, Discriminator
+
+# Parameters
+LAMBDA = 10
+EPOCHS = 10
+DATA_PATH = "data/horse2zebra"
+
+generator_g = Generator() # Horse -> Zebra
+generator_f = Generator() # Zebra -> Horse
+
+discriminator_x = Discriminator() # Real Horse vs Fake Horse
+discriminator_y = Discriminator() # Real Zebra vs Fake Zebra
+
+loss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True)
+
+def discriminator_loss(real, generated):
+    real_loss = loss_obj(tf.ones_like(real), real)
+    generated_loss = loss_obj(tf.zeros_like(generated), generated)
+    total_disc_loss = real_loss + generated_loss
+    return total_disc_loss * 0.5
+
+def generator_loss(generated):
+    return loss_obj(tf.ones_like(generated), generated)
+
+def calc_cycle_loss(real_image, cycled_image):
+    loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))
+    return LAMBDA * loss1
+
+def identity_loss(real_image, same_image):
+    loss = tf.reduce_mean(tf.abs(real_image - same_image))
+    return LAMBDA * 0.5 * loss
+
+generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
+generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
+
+discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
+discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
+
+@tf.function
+def train_step(real_x, real_y):
+    with tf.GradientTape(persistent=True) as tape:
+        # Generator G translates X -> Y
+        # Generator F translates Y -> X.
+        
+        fake_y = generator_g(real_x, training=True)
+        cycled_x = generator_f(fake_y, training=True)
+        
+        fake_x = generator_f(real_y, training=True)
+        cycled_y = generator_g(fake_x, training=True)
+        
+        # same_x and same_y are used for identity loss.
+        same_x = generator_f(real_x, training=True)
+        same_y = generator_g(real_y, training=True)
+        
+        disc_real_x = discriminator_x(real_x, training=True)
+        disc_real_y = discriminator_y(real_y, training=True)
+        
+        disc_fake_x = discriminator_x(fake_x, training=True)
+        disc_fake_y = discriminator_y(fake_y, training=True)
+        
+        # calculate the loss
+        gen_g_loss = generator_loss(disc_fake_y)
+        gen_f_loss = generator_loss(disc_fake_x)
+        
+        total_cycle_loss = calc_cycle_loss(real_x, cycled_x) + calc_cycle_loss(real_y, cycled_y)
+        
+        # Total generator loss = adversarial loss + cycle loss + identity loss
+        total_gen_g_loss = gen_g_loss + total_cycle_loss + identity_loss(real_y, same_y)
+        total_gen_f_loss = gen_f_loss + total_cycle_loss + identity_loss(real_x, same_x)
+        
+        disc_x_loss = discriminator_loss(disc_real_x, disc_fake_x)
+        disc_y_loss = discriminator_loss(disc_real_y, disc_fake_y)
+        
+    # Calculate the gradients for generator and discriminator
+    generator_g_gradients = tape.gradient(total_gen_g_loss, generator_g.trainable_variables)
+    generator_f_gradients = tape.gradient(total_gen_f_loss, generator_f.trainable_variables)
+    
+    discriminator_x_gradients = tape.gradient(disc_x_loss, discriminator_x.trainable_variables)
+    discriminator_y_gradients = tape.gradient(disc_y_loss, discriminator_y.trainable_variables)
+    
+    # Apply the gradients to the optimizer
+    generator_g_optimizer.apply_gradients(zip(generator_g_gradients, generator_g.trainable_variables))
+    generator_f_optimizer.apply_gradients(zip(generator_f_gradients, generator_f.trainable_variables))
+    
+    discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients, discriminator_x.trainable_variables))
+    discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients, discriminator_y.trainable_variables))
+
+def main():
+    train_dataset = get_dataset(DATA_PATH, "train").batch(1)
+    
+    for epoch in range(EPOCHS):
+        start = time.time()
+        print(f"Epoch {epoch} starting...")
+        
+        for n, (image_x, image_y) in train_dataset.enumerate():
+            train_step(image_x, image_y)
+            if n % 100 == 0:
+                print ('.', end='', flush=True)
+        
+        print(f"\nTime for epoch {epoch} is {time.time()-start} sec")
+        
+        # Save checkpoints
+        generator_g.save_weights(f"GeneratorHtoZ_epoch_{epoch}.h5")
+        generator_f.save_weights(f"GeneratorZtoH_epoch_{epoch}.h5")
+        
+        # Also save latest weights
+        generator_g.save_weights("GeneratorHtoZ.h5")
+        generator_f.save_weights("GeneratorZtoH.h5")
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,162 @@
+import torch
+from dataset import CycleGANDataset
+from torch.utils.data import DataLoader
+import torch.nn as nn
+import torch.optim as optim
+from models import Generator, Discriminator
+from tqdm import tqdm
+from torchvision.utils import save_image
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import os
+
+# Hyperparameters
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TRAIN_DIR_HORSE = "data/horse2zebra/trainA"
+TRAIN_DIR_ZEBRA = "data/horse2zebra/trainB"
+VAL_DIR_HORSE = "data/horse2zebra/testA"
+VAL_DIR_ZEBRA = "data/horse2zebra/testB"
+BATCH_SIZE = 1
+LEARNING_RATE = 1e-5
+LAMBDA_IDENTITY = 0.0
+LAMBDA_CYCLE = 10
+NUM_WORKERS = 1
+NUM_EPOCHS = 10
+LOAD_MODEL = False
+SAVE_MODEL = True
+CHECKPOINT_GEN_H = "genh.pth.tar"
+CHECKPOINT_GEN_Z = "genz.pth.tar"
+CHECKPOINT_CRITIC_H = "critich.pth.tar"
+CHECKPOINT_CRITIC_Z = "criticz.pth.tar"
+
+transforms = A.Compose(
+    [
+        A.Resize(width=256, height=256),
+        A.HorizontalFlip(p=0.5),
+        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.0),
+        ToTensorV2(),
+    ],
+    additional_targets={"image0": "image"},
+)
+
+def train_fn(disc_H, disc_Z, gen_Z, gen_H, loader, opt_disc, opt_gen, l1, mse, d_scaler, g_scaler):
+    H_reals = 0
+    H_fakes = 0
+    loop = tqdm(loader, leave=True)
+
+    for idx, (horse, zebra) in enumerate(loop):
+        horse = horse.to(DEVICE)
+        zebra = zebra.to(DEVICE)
+
+        # Train Discriminators H and Z
+        with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")):
+            fake_horse = gen_H(zebra)
+            D_H_real = disc_H(horse)
+            D_H_fake = disc_H(fake_horse.detach())
+            H_reals += D_H_real.mean().item()
+            H_fakes += D_H_fake.mean().item()
+            D_H_real_loss = mse(D_H_real, torch.ones_like(D_H_real))
+            D_H_fake_loss = mse(D_H_fake, torch.zeros_like(D_H_fake))
+            D_H_loss = D_H_real_loss + D_H_fake_loss
+
+            fake_zebra = gen_Z(horse)
+            D_Z_real = disc_Z(zebra)
+            D_Z_fake = disc_Z(fake_zebra.detach())
+            D_Z_real_loss = mse(D_Z_real, torch.ones_like(D_Z_real))
+            D_Z_fake_loss = mse(D_Z_fake, torch.zeros_like(D_Z_fake))
+            D_Z_loss = D_Z_real_loss + D_Z_fake_loss
+
+            # put it together
+            D_loss = (D_H_loss + D_Z_loss) / 2
+
+        opt_disc.zero_grad()
+        d_scaler.scale(D_loss).backward()
+        d_scaler.step(opt_disc)
+        d_scaler.update()
+
+        # Train Generators H and Z
+        with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")):
+            # adversarial loss for both generators
+            D_H_fake = disc_H(fake_horse)
+            D_Z_fake = disc_Z(fake_zebra)
+            loss_G_H = mse(D_H_fake, torch.ones_like(D_H_fake))
+            loss_G_Z = mse(D_Z_fake, torch.ones_like(D_Z_fake))
+
+            # cycle loss
+            cycle_zebra = gen_Z(fake_horse)
+            cycle_horse = gen_H(fake_zebra)
+            cycle_zebra_loss = l1(zebra, cycle_zebra)
+            cycle_horse_loss = l1(horse, cycle_horse)
+
+            # identity loss (remove these for efficiency if you want)
+            # identity_zebra = gen_Z(zebra)
+            # identity_horse = gen_H(horse)
+            # identity_zebra_loss = l1(zebra, identity_zebra)
+            # identity_horse_loss = l1(horse, identity_horse)
+
+            # add all together
+            G_loss = (
+                loss_G_Z
+                + loss_G_H
+                + cycle_zebra_loss * LAMBDA_CYCLE
+                + cycle_horse_loss * LAMBDA_CYCLE
+                # + identity_horse_loss * LAMBDA_IDENTITY
+                # + identity_zebra_loss * LAMBDA_IDENTITY
+            )
+
+        opt_gen.zero_grad()
+        g_scaler.scale(G_loss).backward()
+        g_scaler.step(opt_gen)
+        g_scaler.update()
+
+        if idx % 200 == 0:
+            torch.save(gen_H.state_dict(), f"saved_images/genh.pth.tar")
+            torch.save(gen_Z.state_dict(), f"saved_images/genz.pth.tar")
+            save_image(fake_horse * 0.5 + 0.5, f"saved_images/horse_{idx}.png")
+            save_image(fake_zebra * 0.5 + 0.5, f"saved_images/zebra_{idx}.png")
+
+        loop.set_postfix(H_real=H_reals / (idx + 1), H_fake=H_fakes / (idx + 1))
+
+def main():
+    disc_H = Discriminator(in_channels=3).to(DEVICE)
+    disc_Z = Discriminator(in_channels=3).to(DEVICE)
+    gen_Z = Generator(img_channels=3, num_residuals=9).to(DEVICE)
+    gen_H = Generator(img_channels=3, num_residuals=9).to(DEVICE)
+    opt_disc = optim.Adam(
+        list(disc_H.parameters()) + list(disc_Z.parameters()),
+        lr=LEARNING_RATE,
+        betas=(0.5, 0.999),
+    )
+
+    opt_gen = optim.Adam(
+        list(gen_Z.parameters()) + list(gen_H.parameters()),
+        lr=LEARNING_RATE,
+        betas=(0.5, 0.999),
+    )
+
+    L1 = nn.L1Loss()
+    MSE = nn.MSELoss()
+
+    dataset = CycleGANDataset(
+        root_horse=TRAIN_DIR_HORSE,
+        root_zebra=TRAIN_DIR_ZEBRA,
+        transform=transforms,
+    )
+    loader = DataLoader(
+        dataset,
+        batch_size=BATCH_SIZE,
+        shuffle=True,
+        num_workers=NUM_WORKERS,
+        pin_memory=True,
+    )
+    g_scaler = torch.cuda.amp.GradScaler(enabled=(DEVICE == "cuda"))
+    d_scaler = torch.cuda.amp.GradScaler(enabled=(DEVICE == "cuda"))
+
+    os.makedirs("saved_images", exist_ok=True)
+
+    for epoch in range(NUM_EPOCHS):
+        print(f"Epoch {epoch}/{NUM_EPOCHS}")
+        train_fn(disc_H, disc_Z, gen_Z, gen_H, loader, opt_disc, opt_gen, L1, MSE, d_scaler, g_scaler)
+
+if __name__ == "__main__":
+    main()