Initial Upload (attempt 2)

.gitattributes

@@ -40,3 +40,10 @@ FastSplatStyler_huggingface/style_ims/style0.jpg filter=lfs diff=lfs merge=lfs -
 FastSplatStyler_huggingface/style_ims/style2.jpg filter=lfs diff=lfs merge=lfs -text
 FastSplatStyler_huggingface/style_ims/style44.jpg filter=lfs diff=lfs merge=lfs -text
 FastSplatStyler_huggingface/style_ims/style6.jpg filter=lfs diff=lfs merge=lfs -text
+example-broche-rose-gold.splat filter=lfs diff=lfs merge=lfs -text
+output.splat filter=lfs diff=lfs merge=lfs -text
+style_ims/style-10.jpg filter=lfs diff=lfs merge=lfs -text
+style_ims/style0.jpg filter=lfs diff=lfs merge=lfs -text
+style_ims/style2.jpg filter=lfs diff=lfs merge=lfs -text
+style_ims/style44.jpg filter=lfs diff=lfs merge=lfs -text
+style_ims/style6.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 ECU Computer Vision Lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,14 +1,36 @@
----
-title: FastSplatStyler
-emoji: 🐨
-colorFrom: blue
-colorTo: yellow
-sdk: gradio
-sdk_version: 6.9.0
-app_file: app.py
-pinned: false
-license: mit
-short_description: Optimization-Free Style Transfer for 3D Gaussian Splats
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FastSplatStyler
+Official Implementation of "Optimization-Free Style Transfer of 3D Gaussian Splats"
+
+[arXiv Paper](https://arxiv.org/abs/2508.05813)
+
+![](example.jpg)
+
+## Example Outputs
+
+Example Outputs can be visualized using the [Antimatter WebGL viewer](https://antimatter15.com/splat/) at the following links.
+
+- Broche: [Original](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/broche-rose-gold_original.splat) and [Stylized](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/broche-rose-gold_style3.splat)
+- Crystal Lamp: [Original](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/crystal-lamp-original.splat) and [Stylized](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/crystal-lamp-style2.splat)
+- Family Statue: [Original](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/family.ply) and [Stylized](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/family-style-6.splat)
+- M60 Tanks: [Original](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/m60.ply) and [Stylized](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/m60-style-31.splat)
+- Table: [Original](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/Table.ply) and [Stylized](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/Table_style5.splat)
+- Train: [Original](https://antimatter15.com/splat/) and [Stylized](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/train_style1.splat)
+- Truck: [Original](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/truck.splat) and [Stylized](https://antimatter15.com/splat/?url=https://huggingface.co/datasets/incrl/fast-splat-styler/resolve/main/truck-style-21.splat)
+
+Example inputs and outputs can be downloaded from [Google Drive](https://drive.google.com/drive/folders/10YmtcCOKGosXfPEi84ho1AfYRYioYo12?usp=drive_link) or [Hugging Face](https://huggingface.co/datasets/incrl/fast-splat-styler/tree/main)
+
+## Demo
+
+**Coming soon**
+
+## Install
+
+This work relies heavily on the [Pytorch](https://pytorch.org/) and [Pytorch Geometric](https://www.pyg.org/) libraries.
+
+This code was tested with Python 3.12, Pytorch 2.9.1 (CUDA Toolkit 12.8), and Pytorch Geometric 2.8 on Windows. It was also tested on Mac with the same setup (without Cuda). The provided requirements.txt file comes from the Mac configuration.
+
+This repository relies on a graph networks library that was presented in a [previous work](https://github.com/davidmhart/interpolated-selectionconv/tree/main). The library can be downloaded, with included model weights, at this [google drive link](https://drive.google.com/drive/folders/10YmtcCOKGosXfPEi84ho1AfYRYioYo12?usp=drive_link). Place the "graph_networks" folder in the main directory.
+
+To stylize an example splat, run `python styletransfer_splat.py example-broche-rose-gold.splat --stylePath style_ims/style0.jpg --samplingRate 1.5`. You can change the input splat and style image to your specific use case.
+
+Supports `.splat` and `.ply` files.

app.py

@@ -0,0 +1,302 @@
+import gradio as gr
+import torch
+import os
+import tempfile
+import shutil
+from pathlib import Path
+from time import time
+
+# ── Core style-transfer logic (adapted from styletransfer_splat.py) ──────────
+import pointCloudToMesh as ply2M
+import utils
+import graph_io as gio
+from clusters import *
+import splat_mesh_helpers as splt
+import clusters as cl
+from torch_geometric.data import Data
+from scipy.interpolate import NearestNDInterpolator
+
+from graph_networks.LinearStyleTransfer_vgg import encoder, decoder
+from graph_networks.LinearStyleTransfer_matrix import TransformLayer
+from graph_networks.LinearStyleTransfer.libs.Matrix import MulLayer
+from graph_networks.LinearStyleTransfer.libs.models import encoder4, decoder4
+
+
+# ── Example assets (place your own files in ./examples/) ─────────────────────
+EXAMPLE_SPLATS = [
+    ["example-broche-rose-gold.splat",  "style_ims/style2.jpg"],
+    ["example-broche-rose-gold.splat",  "style_ims/style6.jpg"],
+]
+
+
+# ── Style-transfer function called by Gradio ─────────────────────────────────
+def run_style_transfer(
+    splat_file,
+    style_image,
+    threshold: float,
+    sampling_rate: float,
+    device_choice: str,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if splat_file is None:
+        raise gr.Error("Please upload a 3D Gaussian Splat file (.ply or .splat).")
+    if style_image is None:
+        raise gr.Error("Please upload a style image.")
+
+    device = device_choice if device_choice == "cpu" else f"cuda:{device_choice}"
+
+    # ── Parameters ────────────────────────────────────────────────────────────
+    n = 25
+    ratio = 0.25
+    depth = 3
+    style_shape = (512, 512)
+
+    logs = []
+
+    def log(msg):
+        logs.append(msg)
+        print(msg)
+        return "\n".join(logs)
+
+    # ── 1. Load splat ─────────────────────────────────────────────────────────
+    progress(0.05, desc="Loading splat…")
+    splat_path = splat_file.name if hasattr(splat_file, "name") else splat_file
+    log(f"Loading splat: {splat_path}")
+
+    pos3D_Original, _, colors_Original, opacity_Original, scales_Original, rots_Original, fileType = \
+        splt.splat_unpacker_with_threshold(n, splat_path, threshold)
+
+    # ── 2. Gaussian super-sampling ────────────────────────────────────────────
+    progress(0.15, desc="Super-sampling…")
+    t0 = time()
+    if sampling_rate > 1:
+        GaussianSamples = int(pos3D_Original.shape[0] * sampling_rate)
+        pos3D, colors = splt.splat_GaussianSuperSampler(
+            pos3D_Original.clone(), colors_Original.clone(),
+            opacity_Original.clone(), scales_Original.clone(), rots_Original.clone(),
+            GaussianSamples,
+        )
+    else:
+        pos3D, colors = pos3D_Original, colors_Original
+    log(f"Nodes in graph: {pos3D.shape[0]}  ({time()-t0:.1f}s)")
+
+    # ── 3. Graph construction ─────────────────────────────────────────────────
+    progress(0.30, desc="Building surface graph…")
+    t0 = time()
+    style_ref = utils.loadImage(style_image, shape=style_shape)
+
+    normalsNP = ply2M.Estimate_Normals(pos3D, threshold)
+    normals = torch.from_numpy(normalsNP)
+
+    up_vector = torch.tensor([[1, 1, 1]], dtype=torch.float)
+    up_vector = up_vector / torch.linalg.norm(up_vector, dim=1)
+
+    pos3D = pos3D.to(device)
+    colors = colors.to(device)
+    normals = normals.to(device)
+    up_vector = up_vector.to(device)
+
+    edge_index, directions = gh.surface2Edges(pos3D, normals, up_vector, k_neighbors=16)
+    edge_index, selections, interps = gh.edges2Selections(edge_index, directions, interpolated=True)
+
+    clusters, edge_indexes, selections_list, interps_list = cl.makeSurfaceClusters(
+        pos3D, normals, edge_index, selections, interps,
+        ratio=ratio, up_vector=up_vector, depth=depth, device=device,
+    )
+    log(f"Graph built  ({time()-t0:.1f}s)")
+
+    # ── 4. Load networks ──────────────────────────────────────────────────────
+    progress(0.50, desc="Loading networks…")
+    t0 = time()
+
+    enc_ref = encoder4()
+    dec_ref = decoder4()
+    matrix_ref = MulLayer("r41")
+
+    enc_ref.load_state_dict(torch.load("graph_networks/LinearStyleTransfer/models/vgg_r41.pth",    map_location=device))
+    dec_ref.load_state_dict(torch.load("graph_networks/LinearStyleTransfer/models/dec_r41.pth",    map_location=device))
+    matrix_ref.load_state_dict(torch.load("graph_networks/LinearStyleTransfer/models/r41.pth",     map_location=device))
+
+    enc = encoder(padding_mode="replicate")
+    dec = decoder(padding_mode="replicate")
+    matrix = TransformLayer()
+
+    with torch.no_grad():
+        enc.copy_weights(enc_ref)
+        dec.copy_weights(dec_ref)
+        matrix.copy_weights(matrix_ref)
+
+    content = Data(
+        x=colors, clusters=clusters,
+        edge_indexes=edge_indexes,
+        selections_list=selections_list,
+        interps_list=interps_list,
+    ).to(device)
+
+    style, _ = gio.image2Graph(style_ref, depth=3, device=device)
+
+    enc = enc.to(device)
+    dec = dec.to(device)
+    matrix = matrix.to(device)
+    log(f"Networks loaded  ({time()-t0:.1f}s)")
+
+    # ── 5. Style transfer ─────────────────────────────────────────────────────
+    progress(0.70, desc="Running style transfer…")
+    t0 = time()
+
+    with torch.no_grad():
+        cF = enc(content)
+        sF = enc(style)
+        feature, _ = matrix(
+            cF["r41"], sF["r41"],
+            content.edge_indexes[3], content.selections_list[3],
+            style.edge_indexes[3],  style.selections_list[3],
+            content.interps_list[3] if hasattr(content, "interps_list") else None,
+        )
+        result = dec(feature, content).clamp(0, 1)
+
+    colors[:, 0:3] = result
+    log(f"Stylization done  ({time()-t0:.1f}s)")
+
+    # ── 6. Interpolate back to original resolution ────────────────────────────
+    progress(0.88, desc="Interpolating back to original splat…")
+    t0 = time()
+
+    interp2 = NearestNDInterpolator(pos3D.cpu(), colors.cpu())
+    results_OriginalNP = interp2(pos3D_Original)
+    results_Original = torch.from_numpy(results_OriginalNP).to(torch.float32)
+    colors_and_opacity_Original = torch.cat(
+        (results_Original, opacity_Original.unsqueeze(1)), dim=1
+    )
+    log(f"Interpolation done  ({time()-t0:.1f}s)")
+
+    # ── 7. Save output ────────────────────────────────────────────────────────
+    progress(0.95, desc="Saving output splat…")
+    suffix = ".splat" if fileType == "splat" else ".ply"
+    out_dir = tempfile.mkdtemp()
+    out_path = os.path.join(out_dir, f"stylized{suffix}")
+
+    splt.splat_save(
+        pos3D_Original.numpy(),
+        scales_Original.numpy(),
+        rots_Original.numpy(),
+        colors_and_opacity_Original.numpy(),
+        out_path,
+        fileType,
+    )
+    log(f"Saved to: {out_path}")
+    progress(1.0, desc="Done!")
+
+    return out_path, "\n".join(logs)
+
+
+# ── Gradio UI ─────────────────────────────────────────────────────────────────
+def build_ui():
+    available_devices = (
+        [str(i) for i in range(torch.cuda.device_count())] + ["cpu"]
+        if torch.cuda.is_available()
+        else ["cpu"]
+    )
+
+    with gr.Blocks(
+        title="3DGS Style Transfer",
+        theme=gr.themes.Soft(primary_hue="violet"),
+        css="""
+        #title { text-align: center; }
+        #subtitle { text-align: center; color: #666; margin-bottom: 1rem; }
+        .panel { border-radius: 12px; }
+        #run-btn { font-size: 1.1rem; }
+        """,
+    ) as demo:
+
+        gr.Markdown("# 🎨 3D Gaussian Splat Style Transfer", elem_id="title")
+        gr.Markdown(
+            "Upload a 3DGS scene and a style image — the app will repaint the splat "
+            "with the artistic style of the image and give you a stylized splat to download. "
+            "After downloading, you can view your splat with an [online viewer](https://antimatter15.com/splat/).",
+            elem_id="subtitle",
+        )
+
+        with gr.Row():
+            # ── Left column: inputs ───────────────────────────────────────────
+            with gr.Column(scale=1, elem_classes="panel"):
+                gr.Markdown("### 📂 Inputs")
+
+                splat_input = gr.File(
+                    label="3D Gaussian Splat (.ply or .splat)",
+                    file_types=[".ply", ".splat"],
+                    type="filepath",
+                )
+
+                style_input = gr.Image(
+                    label="Style Image",
+                    type="filepath",
+                    height=240,
+                )
+
+                with gr.Accordion("⚙️ Advanced Settings", open=False):
+                    threshold_slider = gr.Slider(
+                        minimum=90.0, maximum=100.0, value=99.8, step=0.1,
+                        label="Opacity threshold (percentile)",
+                        info="Points below this opacity percentile are removed.",
+                    )
+                    sampling_slider = gr.Slider(
+                        minimum=0.5, maximum=3.0, value=1.5, step=0.1,
+                        label="Gaussian super-sampling rate",
+                        info="Values > 1 add extra samples; 1.0 = no super-sampling.",
+                    )
+                    device_radio = gr.Radio(
+                        choices=available_devices,
+                        value=available_devices[0],
+                        label="Device",
+                    )
+
+                run_btn = gr.Button("🚀 Run Style Transfer", variant="primary", elem_id="run-btn")
+
+            # ── Right column: outputs ─────────────────────────────────────────
+            with gr.Column(scale=1, elem_classes="panel"):
+                gr.Markdown("### 📥 Output")
+
+                output_file = gr.File(
+                    label="Download Stylized Splat",
+                    interactive=False,
+                )
+
+                log_box = gr.Textbox(
+                    label="Progress log",
+                    lines=12,
+                    max_lines=20,
+                    interactive=False,
+                    placeholder="Logs will appear here once processing starts…",
+                )
+
+        # ── Examples ─────────────────────────────────────────────────────────
+        example_splat_paths = [row[0] for row in EXAMPLE_SPLATS]
+        example_style_paths = [row[1] for row in EXAMPLE_SPLATS]
+
+        valid_examples = [
+            row for row in EXAMPLE_SPLATS
+            if os.path.exists(row[0]) and os.path.exists(row[1])
+        ]
+
+        if valid_examples:
+            gr.Markdown("### 🖼️ Examples")
+            gr.Examples(
+                examples=valid_examples,
+                inputs=[splat_input, style_input],
+                label="Click an example to load it",
+            )
+
+        # ── Event wiring ──────────────────────────────────────────────────────
+        run_btn.click(
+            fn=run_style_transfer,
+            inputs=[splat_input, style_input, threshold_slider, sampling_slider, device_radio],
+            outputs=[output_file, log_box],
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = build_ui()
+    demo.launch(share=False)

clusters.py

@@ -0,0 +1,234 @@
+import torch
+from torch_scatter import scatter
+from torch_geometric.nn.pool.consecutive import consecutive_cluster
+from torch_geometric.utils import add_self_loops, add_remaining_self_loops, remove_self_loops
+from torch_geometric.nn import fps, knn
+from torch_sparse import coalesce
+import graph_helpers as gh
+import sphere_helpers as sh
+import mesh_helpers as mh
+
+import math
+from math import pi,sqrt
+
+
+from warnings import warn
+
+def makeImageClusters(pos2D,Nx,Ny,edge_index,selections,depth=1,device='cpu',stride=2):
+    clusters = []
+    edge_indexes = [torch.clone(edge_index).to(device)]
+    selections_list = [torch.clone(selections).to(device)]
+    
+    for _ in range(depth):
+        Nx = Nx//stride
+        Ny = Ny//stride
+        cx,cy = getGrid(pos2D,Nx,Ny)
+        cluster, pos2D = gridCluster(pos2D,cx,cy,Nx)
+        edge_index, selections = selectionAverage(cluster, edge_index, selections)
+        
+        clusters.append(torch.clone(cluster).to(device))
+        edge_indexes.append(torch.clone(edge_index).to(device))
+        selections_list.append(torch.clone(selections).to(device))
+
+    return clusters, edge_indexes, selections_list
+        
+def makeSphereClusters(pos3D,edge_index,selections,interps,rows,cols,cluster_method="layering",stride=2,bary_d=None,depth=1,device='cpu'):
+    clusters = []
+    edge_indexes = [torch.clone(edge_index).to(device)]
+    selections_list = [torch.clone(selections).to(device)]
+    interps_list = [torch.clone(interps).to(device)]
+    
+    for _ in range(depth):
+    
+        rows = rows//stride
+        cols = cols//stride
+        
+        if bary_d is not None:
+            bary_d = bary_d*stride
+    
+        if cluster_method == "equirec":
+            centroids, _ = sh.sampleSphere_Equirec(rows,cols)
+    
+        elif cluster_method == "layering":
+            centroids, _ = sh.sampleSphere_Layering(rows)
+            
+        elif cluster_method == "spiral":
+            centroids, _ = sh.sampleSphere_Spiral(rows,cols)
+            
+        elif cluster_method == "icosphere":
+            centroids, _ = sh.sampleSphere_Icosphere(rows)
+            
+        elif cluster_method == "random":
+            centroids, _ = sh.sampleSphere_Random(rows,cols)
+            
+        elif cluster_method == "random_nodes":
+            index = torch.multinomial(torch.ones(len(pos3D)),N) # close equivalent to np.random.choice
+            centroids = pos3D[index]
+
+        elif cluster_method == "fps":
+            # Farthest Point Search used in PointNet++
+            index = fps(pos3D, ratio=ratio)
+            centroids = pos3D[index]
+        else:
+            raise ValueError("Sphere cluster_method unknown")
+        
+        
+        # Find closest centriod to each current point
+        cluster = knn(centroids,pos3D,1)[1]
+        cluster, _ = consecutive_cluster(cluster)
+        pos3D = scatter(pos3D, cluster, dim=0, reduce='mean')
+
+        # Regenerate surface graph
+        normals = pos3D/torch.linalg.norm(pos3D,dim=1,keepdims=True) # Make sure normals are unit vectors
+        edge_index,directions = gh.surface2Edges(pos3D,normals)
+        edge_index,selections,interps = gh.edges2Selections(edge_index,directions,interpolated=True,bary_d=bary_d)
+        
+        clusters.append(torch.clone(cluster).to(device))
+        edge_indexes.append(torch.clone(edge_index).to(device))
+        selections_list.append(torch.clone(selections).to(device))
+        interps_list.append(torch.clone(interps).to(device))
+    
+    return clusters, edge_indexes, selections_list, interps_list
+
+def makeSurfaceClusters(pos3D,normals,edge_index,selections,interps,cluster_method="random",ratio=.25,up_vector=None,depth=1,device='cpu'):
+    clusters = []
+    edge_indexes = [torch.clone(edge_index).to(device)]
+    selections_list = [torch.clone(selections).to(device)]
+    interps_list = [torch.clone(interps).to(device)]
+    
+    for _ in range(depth):
+    
+        #Desired number of clusters in the next level
+        N = int(len(pos3D) * ratio)
+            
+        if cluster_method == "random":
+            index = torch.multinomial(torch.ones(len(pos3D)),N) # close equivalent to np.random.choice
+            centroids = pos3D[index]
+
+        elif cluster_method == "fps":
+            # Farthest Point Search used in PointNet++
+            index = fps(pos3D, ratio=ratio)
+            centroids = pos3D[index]
+        
+        # Find closest centriod to each current point
+        cluster = knn(centroids,pos3D,1)[1]
+        cluster, _ = consecutive_cluster(cluster)
+        pos3D = scatter(pos3D, cluster, dim=0, reduce='mean')
+        normals = scatter(normals, cluster, dim=0, reduce='mean')
+
+        # Regenerate surface graph
+        normals = normals/torch.linalg.norm(normals,dim=1,keepdims=True) # Make sure normals are unit vectors
+        edge_index,directions = gh.surface2Edges(pos3D,normals,up_vector,k_neighbors=16)
+        edge_index,selections,interps = gh.edges2Selections(edge_index,directions,interpolated=True)
+        
+        clusters.append(torch.clone(cluster).to(device))
+        edge_indexes.append(torch.clone(edge_index).to(device))
+        selections_list.append(torch.clone(selections).to(device))
+        interps_list.append(torch.clone(interps).to(device))
+    
+    return clusters, edge_indexes, selections_list, interps_list
+
+
+def makeMeshClusters(pos3D,mesh,edge_index,selections,interps,ratio=.25,up_vector=None,depth=1,device='cpu'):
+    clusters = []
+    edge_indexes = [torch.clone(edge_index).to(device)]
+    selections_list = [torch.clone(selections).to(device)]
+    interps_list = [torch.clone(interps).to(device)]
+    
+    for _ in range(depth):
+    
+        #Desired number of clusters in the next level
+        N = int(len(pos3D) * ratio)
+        
+        # Generate new point cloud from downsampled version of texture map
+        centroids, normals = mh.sampleSurface(mesh,N,return_x=False)
+        
+        # Find closest centriod to each current point
+        cluster = knn(centroids,pos3D,1)[1]
+        cluster, _ = consecutive_cluster(cluster)
+        pos3D = scatter(pos3D, cluster, dim=0, reduce='mean')
+
+        # Regenerate surface graph
+        #normals = normals/torch.linalg.norm(normals,dim=1,keepdims=True) # Make sure normals are unit vectors
+        edge_index,directions = gh.surface2Edges(pos3D,normals,up_vector)
+        edge_index,selections,interps = gh.edges2Selections(edge_index,directions,interpolated=True)
+        
+        clusters.append(torch.clone(cluster).to(device))
+        edge_indexes.append(torch.clone(edge_index).to(device))
+        selections_list.append(torch.clone(selections).to(device))
+        interps_list.append(torch.clone(interps).to(device))
+    
+    return clusters, edge_indexes, selections_list, interps_list
+    
+    
+    
+def getGrid(pos,Nx,Ny,xrange=None,yrange=None):
+    xmin = torch.min(pos[:,0]) if xrange is None else xrange[0]
+    ymin = torch.min(pos[:,1]) if yrange is None else yrange[0]
+    xmax = torch.max(pos[:,0]) if xrange is None else xrange[1]
+    ymax = torch.max(pos[:,1]) if yrange is None else yrange[1]
+
+    cx = torch.clamp(torch.floor((pos[:,0] - xmin)/(xmax-xmin) * Nx),0,Nx-1)
+    cy = torch.clamp(torch.floor((pos[:,1] - ymin)/(ymax-ymin) * Ny),0,Ny-1)
+    return cx, cy
+    
+def gridCluster(pos,cx,cy,xmax):
+    cluster = cx + cy*xmax
+    cluster = cluster.type(torch.long) # Cast appropriately
+    cluster, _ = consecutive_cluster(cluster)
+    pos = scatter(pos, cluster, dim=0, reduce='mean')
+    
+    return cluster, pos
+
+def selectionAverage(cluster, edge_index, selections):
+    num_nodes = cluster.size(0)
+    edge_index = cluster[edge_index.contiguous().view(1, -1)].view(2, -1)
+    edge_index, selections = remove_self_loops(edge_index, selections)
+    if edge_index.numel() > 0:
+
+        # To avoid means over discontinuities, do mean for two selections at at a time
+        final_edge_index, selections_check = coalesce(edge_index, selections.type(torch.float), num_nodes, num_nodes, op="mean")
+        selections_check = torch.round(selections_check).type(torch.long)
+
+        final_selections = torch.zeros_like(selections_check).to(selections.device)
+
+        final_selections[torch.where(selections_check==4)] = 4
+        final_selections[torch.where(selections_check==5)] = 5
+
+        #Rotate selection kernel
+        selections += 2
+        selections = selections % 9 + torch.div(selections, 9, rounding_mode="floor")
+
+        _, selections_check = coalesce(edge_index, selections.type(torch.float), num_nodes, num_nodes, op="mean")
+        selections_check = torch.round(selections_check).type(torch.long)
+        final_selections[torch.where(selections_check==4)] = 2
+        final_selections[torch.where(selections_check==5)] = 3
+
+        #Rotate selection kernel
+        selections += 2
+        selections = selections % 9 + torch.div(selections, 9, rounding_mode="floor")
+
+        _, selections_check = coalesce(edge_index, selections.type(torch.float), num_nodes, num_nodes, op="mean")
+        selections_check = torch.round(selections_check).type(torch.long)
+        final_selections[torch.where(selections_check==4)] = 8
+        final_selections[torch.where(selections_check==5)] = 1
+
+        #Rotate selection kernel
+        selections += 2
+        selections = selections % 9 + torch.div(selections, 9, rounding_mode="floor")
+
+        _, selections_check = coalesce(edge_index, selections.type(torch.float), num_nodes, num_nodes, op="mean")
+        selections_check = torch.round(selections_check).type(torch.long)
+        final_selections[torch.where(selections_check==4)] = 6
+        final_selections[torch.where(selections_check==5)] = 7
+
+        #print(torch.min(final_selections), torch.max(final_selections))
+        #print(torch.mean(final_selections.type(torch.float)))
+
+        edge_index, selections = add_remaining_self_loops(final_edge_index,final_selections,fill_value=torch.tensor(0,dtype=torch.long))
+
+    else:
+        edge_index, selections = add_remaining_self_loops(edge_index,selections,fill_value=torch.tensor(0,dtype=torch.long))
+        print("Warning: Edge Pool found no edges")
+
+    return edge_index, selections

example-broche-rose-gold.splat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62cccf0adfb4e5713985e8ee874fa30a6a37ae222a23ead7c6e4639a5802ab62
+size 4157728

graph_helpers.py

@@ -0,0 +1,400 @@
+import torch
+from torch_geometric.nn import radius_graph, knn_graph
+import torch_geometric as tg
+from torch_geometric.utils import subgraph
+import utils
+from math import sqrt
+
+def getImPos(rows,cols,start_row=0,start_col=0):
+    row_space = torch.arange(start_row,rows+start_row)
+    col_space = torch.arange(start_col,cols+start_col)
+    col_image,row_image = torch.meshgrid(col_space,row_space,indexing='xy')
+    im_pos = torch.reshape(torch.stack((row_image,col_image),dim=-1),(rows*cols,2))
+    return im_pos
+
+def convertImPos(im_pos,flip_y=True):
+    
+    # Cast to float for clustering based methods
+    pos2D = im_pos.float()
+    
+    # Switch rows,cols to x,y
+    pos2D[:,[1,0]] = pos2D[:,[0,1]]
+    
+    if flip_y:
+        
+        # Flip to y-axis to match mathematical definition and edges2Selections settings
+        pos2D[:,1] = torch.amax(pos2D[:,1]) - pos2D[:,1]
+    
+    return pos2D
+
+
+def grid2Edges(locs):
+    # Assume locs are already spaced at a distance of 1 structure
+    edge_index = radius_graph(locs,1.44,loop=True)
+    return edge_index
+    
+def radius2Edges(locs,r=1.0):
+    edge_index = radius_graph(locs,r,loop=True)
+    return edge_index
+    
+def knn2Edges(locs,knn=9):
+    edge_index = knn_graph(locs,knn,loop=True)
+    return edge_index
+
+def surface2Edges(pos3D,normals,up_vector=None,k_neighbors=9):
+
+    if up_vector is None:
+        up_vector = torch.tensor([[0.0,1.0,0.0]]).to(pos3D.device)
+    
+    # K Nearest Neighbors graph
+    edge_index = knn_graph(pos3D,k_neighbors,loop=True)
+
+    # Cull neighbors based on normals (dot them together)
+    culling = torch.sum(torch.multiply(normals[edge_index[1]],normals[edge_index[0]]),dim=1)
+    edge_index = edge_index[:,torch.where(culling>0)[0]]
+
+    # For each node, rotate based on Grahm-Schmidt Orthognalization
+    norms = normals[edge_index[0]]
+    
+    z_dir = norms
+    z_dir = z_dir/torch.linalg.norm(z_dir,dim=1,keepdims=True) # Make sure it is a unit vector
+    #x_dir = torch.cross(up_vector,norms,dim=1)
+    x_dir = utils.cross(up_vector,norms) # torch.cross doesn't broadcast properly in some versions of torch
+    x_dir = x_dir/torch.linalg.norm(x_dir,dim=1,keepdims=True)
+    #y_dir = torch.cross(norms,x_dir,dim=1)
+    y_dir = utils.cross(norms,x_dir)
+    y_dir = y_dir/torch.linalg.norm(y_dir,dim=1,keepdims=True)
+
+    directions = (pos3D[edge_index[1]] - pos3D[edge_index[0]])
+    
+    # Perform rotation by multiplying out rotation matrix
+    temp = torch.clone(directions) # Buffer
+    directions[:,0] = temp[:,0] * x_dir[:,0] + temp[:,1] * x_dir[:,1] + temp[:,2] * x_dir[:,2]
+    directions[:,1] = temp[:,0] * y_dir[:,0] + temp[:,1] * y_dir[:,1] + temp[:,2] * y_dir[:,2]
+    #directions[:,2] = temp[:,0] * z_dir[:,0] + temp[:,1] * z_dir[:,1] + temp[:,2] * z_dir[:,2]
+    
+    # Drop z coordinate
+    directions = directions[:,:2]
+    
+    return edge_index, directions
+      
+def edges2Selections(edge_index,directions,interpolated=True,bary_d=None,y_down=False):
+    
+    # Current Ordering
+    # 4  3  2
+    # 5  0  1
+    # 6  7  8
+    if y_down:
+        vectorList = torch.tensor([[1,0],[sqrt(2)/2,-sqrt(2)/2],[0,-1],[-sqrt(2)/2,-sqrt(2)/2],[-1,0],[-sqrt(2)/2,sqrt(2)/2],[0,1],[sqrt(2)/2,sqrt(2)/2]],dtype=torch.float).transpose(1,0)
+    else:
+        vectorList = torch.tensor([[1,0],[sqrt(2)/2,sqrt(2)/2],[0,1],[-sqrt(2)/2,sqrt(2)/2],[-1,0],[-sqrt(2)/2,-sqrt(2)/2],[0,-1],[sqrt(2)/2,-sqrt(2)/2]],dtype=torch.float).transpose(1,0)
+    
+    if interpolated:    
+        
+        if bary_d is None:
+            edge_index,selections,interps = interpolateSelections(edge_index,directions,vectorList)
+        else:
+            edge_index,selections,interps = interpolateSelections_barycentric(edge_index,directions,bary_d,vectorList)
+        interps = normalizeEdges(edge_index,selections,interps)
+        return edge_index,selections,interps
+    
+    else:
+        selections = torch.argmax(torch.matmul(directions,vectorList),dim=1) + 1
+        selections[torch.where(torch.sum(torch.abs(directions),axis=1) == 0)] = 0 # Same cell selection
+        return selections
+
+    
+def makeEdges(prev_sources,prev_targets,prev_selections,sources,targets,selection,reverse=True):
+    
+    sources = sources.flatten()
+    targets = targets.flatten()
+    
+    prev_sources += sources.tolist()
+    prev_targets += targets.tolist()
+    prev_selections += len(sources)*[selection]
+    
+    if reverse:
+        prev_sources += targets
+        prev_targets += sources
+        prev_selections += len(sources)*[utils.reverse_selection(selection)]
+        
+    return prev_sources,prev_targets,prev_selections
+        
+def maskNodes(mask,x):
+    node_mask = torch.where(mask)
+    x = x[node_mask]
+    return x
+
+def maskPoints(mask,x,y):
+    
+    mask = torch.squeeze(mask)
+    
+    x0 = torch.floor(x).long()
+    x1 = x0 + 1
+    y0 = torch.floor(y).long()
+    y1 = y0 + 1
+    
+    x0 = torch.clip(x0, 0, mask.shape[1]-1);
+    x1 = torch.clip(x1, 0, mask.shape[1]-1);
+    y0 = torch.clip(y0, 0, mask.shape[0]-1);
+    y1 = torch.clip(y1, 0, mask.shape[0]-1);
+
+    Ma = mask[ y0, x0 ]
+    Mb = mask[ y1, x0 ]
+    Mc = mask[ y0, x1 ]
+    Md = mask[ y1, x1 ]
+    
+    node_mask = torch.where(torch.logical_and(torch.logical_and(torch.logical_and(Ma,Mb),Mc),Md))[0]
+    
+    return node_mask
+    
+
+def maskGraph(mask,edge_index,selections,interps=None):
+    
+    edge_index,_,edge_mask = subgraph(mask,edge_index,relabel_nodes=True,return_edge_mask=True)
+    selections = selections[edge_mask]
+    
+    if interps:
+        interps = interps[edge_mask]
+        return edge_index, selections, interps
+    else:
+        return edge_index, selections
+    
+def interpolateSelections(edge_index,directions,vectorList=None):
+
+    if vectorList is None:
+        # Current Ordering
+        # 4  3  2
+        # 5  0  1
+        # 6  7  8
+        vectorList = torch.tensor([[1,0],[sqrt(2)/2,sqrt(2)/2],[0,1],[-sqrt(2)/2,sqrt(2)/2],[-1,0],[-sqrt(2)/2,-sqrt(2)/2],[0,-1],[sqrt(2)/2,-sqrt(2)/2]],dtype=torch.float).transpose(1,0)
+
+    # Normalize directions for simplicity of calculations
+    dir_norm = torch.linalg.norm(directions,dim=1,keepdims=True)
+    directions = directions/dir_norm
+    #locs = torch.where(dir_norm > 1)[0]
+    #directions[locs] = directions[locs]/dir_norm[locs]
+    
+    values = torch.matmul(directions,vectorList)
+    best = torch.unsqueeze(torch.argmax(values,dim=1),1)
+    
+    best_val = torch.take_along_dim(values,best,dim=1)
+    
+    # Look at both neighbors to see who is closer
+    lower_val = torch.take_along_dim(values,(best-1) % 8,dim=1)
+    upper_val = torch.take_along_dim(values,(best+1) % 8,dim=1)
+    
+    comp_vals = torch.cat((lower_val,upper_val),dim=1)
+    
+    second_best_vals = torch.amax(comp_vals,dim=1)
+    second_best = torch.argmax(comp_vals,dim=1)
+    
+    # Find the interpolation value (in terms of angles)
+    best_val = torch.minimum(best_val[:,0],torch.tensor(1,device=directions.device)) # Prep for arccos function
+    angle_best = torch.arccos(best_val)
+    angle_second_best = torch.arccos(second_best_vals)
+    
+    angle_vals = angle_best/(angle_second_best + angle_best)
+    
+    # Use negative values for clockwise selections
+    clockwise = torch.where(second_best == 0)[0]
+    angle_vals[clockwise] = -angle_vals[clockwise] 
+    
+    # Handle computation problems at the poles
+    angle_vals = torch.nan_to_num(angle_vals)
+    
+    # Make Selections
+    selections = best[:,0] + 1
+    
+    # Same cell selection
+    same_locs = torch.where(edge_index[0] == edge_index[1])
+    selections[same_locs] = 0
+    angle_vals[same_locs] = 0
+    
+    # Make starting interp_values
+    interps = torch.ones_like(angle_vals)
+    interps -= torch.abs(angle_vals)
+    
+    # Add new edges
+    pos_interp_locs = torch.where(angle_vals > 1e-2)[0]
+    pos_interps = angle_vals[pos_interp_locs]
+    pos_edges = edge_index[:,pos_interp_locs]
+    pos_selections = selections[pos_interp_locs] + 1
+    pos_selections[torch.where(pos_selections>8)] = 1 # Account for wrap around
+    
+    neg_interp_locs = torch.where(angle_vals < -1e-2)[0]
+    neg_interps = torch.abs(angle_vals[neg_interp_locs])
+    neg_edges = edge_index[:,neg_interp_locs]
+    neg_selections = selections[neg_interp_locs] - 1
+    neg_selections[torch.where(neg_selections<1)] = 8 # Account for wrap around
+    
+    edge_index = torch.cat((edge_index,pos_edges,neg_edges),dim=1)
+    selections = torch.cat((selections,pos_selections,neg_selections),dim=0)
+    interps = torch.cat((interps,pos_interps,neg_interps),dim=0)
+    
+    return edge_index,selections,interps
+
+def interpolateSelections_barycentric(edge_index,directions,d,vectorList=None):
+
+    if vectorList is None:
+        # Current Ordering
+        # 4  3  2
+        # 5  0  1
+        # 6  7  8
+        vectorList = torch.tensor([[1,0],[sqrt(2)/2,-sqrt(2)/2],[0,-1],[-sqrt(2)/2,-sqrt(2)/2],[-1,0],[-sqrt(2)/2,sqrt(2)/2],[0,1],[sqrt(2)/2,sqrt(2)/2]],dtype=torch.float).transpose(1,0).to(directions.device)
+
+    # Preprune central selections and reappend them at the end
+    same_locs = torch.where(edge_index[0] == edge_index[1])
+    same_edges = edge_index[:,same_locs[0]]
+    
+    different_locs = torch.where(edge_index[0] != edge_index[1])
+    edge_index = edge_index[:,different_locs[0]]
+    directions = directions[different_locs[0]]
+        
+    # Normalize directions for simplicity of calculations
+    dir_norm = torch.linalg.norm(directions,dim=1,keepdims=True)
+    unit_directions = directions/dir_norm
+    #locs = torch.where(dir_norm > 1)[0]
+    #directions[locs] = directions[locs]/dir_norm[locs]
+    
+    values = torch.matmul(unit_directions,vectorList)
+    best = torch.unsqueeze(torch.argmax(values,dim=1),1)
+    #best_val = torch.take_along_dim(values,best,dim=1)
+    
+    # Look at both neighbors to see who is closer
+    lower_val = torch.take_along_dim(values,(best-1) % 8,dim=1)
+    upper_val = torch.take_along_dim(values,(best+1) % 8,dim=1)
+    
+    comp_vals = torch.cat((lower_val,upper_val),dim=1)
+    
+    second_best = torch.argmax(comp_vals,dim=1)
+    #second_best_vals = torch.amax(comp_vals,dim=1)
+    
+    # Convert into uv cooridnates for barycentric interpolation calculation
+    #     /|
+    #    / |v
+    #   /__|
+    #    u 
+    
+    scaled_directions = torch.abs(directions/d)
+    u = torch.amax(scaled_directions,dim=1)
+    v = torch.amin(scaled_directions,dim=1)
+    
+    # Force coordinates to be within the triangle
+    boundary_check = torch.where(u > d)
+    v[boundary_check] /= u[boundary_check]
+    u[boundary_check] = 1.0
+    
+    # Precalculated barycentric values from linear matrix solve
+    I0 = 1 - u
+    I1 = u - v
+    I2 = v
+    
+    # Make first selections and proper interps
+    selections = best[:,0] + 1
+    interps = I1
+    even_sels = torch.where(selections % 2 == 0)
+    interps[even_sels] = I2[even_sels] # Corners get different weights
+    
+    # Make new edges for the central selections
+    central_edges = torch.clone(edge_index).to(edge_index.device)
+    central_selections = torch.zeros_like(selections)
+    central_interps = I0
+    
+    # Make new edges for the last selection
+    pos_locs = torch.where(second_best==1)[0]
+    pos_edges = edge_index[:,pos_locs]
+    pos_selections = selections[pos_locs] + 1
+    pos_selections[torch.where(pos_selections>8)] = 1 #Account for wrap around
+    pos_interps = I1[pos_locs]
+    even_sels = torch.where(pos_selections % 2 == 0)
+    pos_interps[even_sels] = I2[pos_locs][even_sels]
+    
+    neg_locs = torch.where(second_best==0)[0]
+    neg_edges = edge_index[:,neg_locs]
+    neg_selections = selections[neg_locs] - 1
+    neg_selections[torch.where(neg_selections<1)] = 8 # Account for wrap around
+    neg_interps = I1[neg_locs]
+    even_sels = torch.where(neg_selections % 2 == 0)
+    neg_interps[even_sels] = I2[neg_locs][even_sels]
+
+    # Account for the previously pruned same node edges
+    same_selections = torch.zeros(same_edges.shape[1],dtype=torch.long)
+    same_interps = torch.ones(same_edges.shape[1],dtype=torch.float)
+    
+    # Combine
+    edge_index = torch.cat((edge_index,central_edges,pos_edges,neg_edges,same_edges),dim=1)
+    selections = torch.cat((selections,central_selections,pos_selections,neg_selections,same_selections),dim=0)
+    interps = torch.cat((interps,central_interps,pos_interps,neg_interps,same_interps),dim=0)
+    
+    #edge_index = torch.cat((edge_index,central_edges,pos_edges,neg_edges),dim=1)
+    #selections = torch.cat((selections,central_selections,pos_selections,neg_selections),dim=0)
+    #interps = torch.cat((interps,central_interps,pos_interps,neg_interps),dim=0)
+        
+    # Account for edges to the same node
+    #same_locs = torch.where(edge_index[0] == edge_index[1])
+    #selections[same_locs] = 0
+    #interps[same_locs] = 1
+    
+    return edge_index,selections,interps
+
+def normalizeEdges(edge_index,selections,interps=None,kernel_norm=False):
+    '''Given an edge_index and selections, normalize the edges for each node so that 
+    aggregation of edges with interps = 1. If interps is given, use a weighted average.
+    if kernel_norm = True, account for missing selections by increasing weight on other selections.'''
+    
+    N = torch.max(edge_index) + 1
+    S = torch.max(selections) + 1
+    
+    total_weight = torch.zeros((N,S),dtype=torch.float).to(edge_index.device)
+    
+    if interps is None:
+        interps = torch.ones(len(selections),dtype=torch.float).to(edge_index.device)
+    
+    # Aggregate all edges to determine normalizations per selection
+    nodes = edge_index[0]
+    #total_weight[nodes,selections] += interps
+    total_weight.index_put_((nodes,selections),interps,accumulate=True)
+    
+    # Reassign interps accordingly
+    if kernel_norm:
+        row_totals = torch.sum(total_weight,dim=1)
+        interps = interps * S/row_totals[nodes]
+    else:
+        norms = total_weight[nodes,selections]
+        norms[torch.where(norms < 1e-6)] = 1e-6 # Avoid divide by zero error
+        interps = interps/norms
+    
+    return interps
+
+def simplifyGraph(edge_index,selections,edge_lengths):
+    # Take the shortest edge for the set of the same selections on a given node
+    num_edges = edge_index.shape[1]
+
+    # Keep track of which nodes have been visited
+    keep_edges = torch.zeros(num_edges,dtype=torch.bool).to(edge_index.device)
+
+    previous_best_distance = 100000*torch.ones((torch.amax(edge_index)+1,torch.amax(selections)+1),dtype=torch.long).to(edge_index.device)
+    previous_best_edge = -1*torch.ones((torch.amax(edge_index)+1,torch.amax(selections)+1),dtype=torch.long).to(edge_index.device)
+
+    for i in range(num_edges):
+        start_node = edge_index[0,i]
+        #end_node = edge_index[1,i]
+        selection = selections[i]
+        distance = edge_lengths[i]
+
+        if distance < previous_best_distance[start_node,selection]:
+            previous_best_distance[start_node,selection] = distance
+            keep_edges[i] = True
+
+            prev = previous_best_edge[start_node,selection]
+            if prev != -1:
+                keep_edges[prev] = False
+
+            previous_best_edge[start_node,selection] = i
+
+    edge_index = edge_index[:,torch.where(keep_edges)[0]]
+    selections = selections[torch.where(keep_edges)]
+
+    return edge_index, selections
+    

graph_io.py

@@ -0,0 +1,306 @@
+import numpy as np
+import torch
+from torch_geometric.nn import knn
+from torch_geometric.data import Data
+from torch_geometric.nn import radius_graph, knn_graph
+
+import graph_helpers as gh
+import sphere_helpers as sh
+import mesh_helpers as mh
+import clusters as cl
+import utils
+
+from torch_scatter import scatter
+
+import math
+from math import pi, sqrt
+
+from warnings import warn
+
+def image2Graph(data, gt = None, mask = None, depth = 1, x_only = False, device = 'cpu'):
+    
+    _,ch,rows,cols = data.shape
+    
+    x = torch.reshape(data,(ch,rows*cols)).permute((1,0)).to(device)
+    
+    if mask is not None:
+        # Mask out nodes
+        node_mask = torch.where(mask.flatten())
+        x = x[node_mask]
+    
+    if gt is not None:
+        y = gt.flatten().to(device)
+        if mask is not None:
+            y = y[node_mask]
+    
+    if x_only:
+        if gt is not None:
+            return x,y
+        else:
+            return x
+    
+    im_pos = gh.getImPos(rows,cols)
+    
+    if mask is not None:
+        im_pos = im_pos[node_mask]
+    
+    # Make "point cloud" for clustering
+    pos2D = gh.convertImPos(im_pos,flip_y=False)
+    
+    # Generate initial graph
+    edge_index = gh.grid2Edges(pos2D)
+    directions = pos2D[edge_index[1]] - pos2D[edge_index[0]]
+    selections = gh.edges2Selections(edge_index,directions,interpolated=False,y_down=True)
+    
+    # Generate info for downsampled versions of the graph
+    clusters, edge_indexes, selections_list = cl.makeImageClusters(pos2D,cols,rows,edge_index,selections,depth=depth,device=device)
+    
+    # Make final graph and metadata needed for mapping the result after going through the network
+    graph = Data(x=x,clusters=clusters,edge_indexes=edge_indexes,selections_list=selections_list,interps_list=None)
+    metadata = Data(original=data,im_pos=im_pos.long(),rows=rows,cols=cols,ch=ch)
+    
+    if gt is not None:
+        graph.y = y
+    
+    return graph,metadata
+
+def graph2Image(result,metadata,canvas=None):
+    
+    x = utils.toNumpy(result,permute=False)
+    im_pos = utils.toNumpy(metadata.im_pos,permute=False)
+    if canvas is None:
+        canvas = utils.makeCanvas(x,metadata.original)
+    
+    # Paint over the original image (neccesary for masked images)
+    canvas[im_pos[:,0],im_pos[:,1]] = x
+    
+    return canvas
+
+### Begin Interpolated Methods ###
+
+def sphere2Graph(data, structure="layering", cluster_method="layering", scale=1.0, stride=2, interpolation_mode = "angle", gt = None, mask = None, depth = 1, x_only = False, device = 'cpu'):
+    
+    _,ch,rows,cols = data.shape
+    
+    if structure == "equirec":
+        # Use the original data to start with
+        cartesian, spherical = sh.sampleSphere_Equirec(scale*rows,scale*cols)
+    elif structure == "layering":
+        cartesian, spherical = sh.sampleSphere_Layering(scale*rows)
+    elif structure == "spiral":
+        cartesian, spherical = sh.sampleSphere_Spiral(scale*rows,scale*cols)
+    elif structure == "icosphere":
+        cartesian, spherical = sh.sampleSphere_Icosphere(scale*rows)
+    elif structure == "random":
+        cartesian, spherical = sh.sampleSphere_Random(scale*rows,scale*cols)
+    else:
+        raise ValueError("Sphere structure unknown")
+    
+    if interpolation_mode == "bary":            
+        bary_d = pi/(scale*rows)
+    else:
+        bary_d = None
+    
+    # Get the landing point for each node
+    sample_x, sample_y = sh.spherical2equirec(spherical[:,0],spherical[:,1],rows,cols)
+    
+    if mask is not None:
+
+        node_mask = gh.maskPoints(mask,sample_x,sample_y)
+        sample_x = sample_x[node_mask]
+        sample_y = sample_y[node_mask]
+        spherical = spherical[node_mask]
+        cartesian = cartesian[node_mask]
+    
+    features = utils.bilinear_interpolate(data, sample_x, sample_y).to(device)
+    
+    if gt is not None:
+        features_y = utils.bilinear_interpolate(gt.unsqueeze(0), sample_x, sample_y).to(device)
+    
+    if x_only:
+        if gt is not None:
+            return features,features_y
+        else:
+            return features
+        
+    # Build initial graph
+    edge_index,directions = gh.surface2Edges(cartesian,cartesian)
+    edge_index,selections,interps = gh.edges2Selections(edge_index,directions,interpolated=True,bary_d=bary_d)
+    
+    # Generate info for downsampled versions of the graph
+    clusters, edge_indexes, selections_list, interps_list = cl.makeSphereClusters(cartesian,edge_index,selections,interps,rows*scale,cols*scale,cluster_method,stride=stride,bary_d=bary_d,depth=depth,device=device)
+    
+    # Make final graph and metadata needed for mapping the result after going through the network
+    graph = Data(x=features,clusters=clusters,edge_indexes=edge_indexes,selections_list=selections_list,interps_list=interps_list)
+    metadata = Data(original=data,pos3D=cartesian,mask=mask,rows=rows,cols=cols,ch=ch)
+
+    if gt is not None:
+        graph.y = features_y   
+    
+    return graph, metadata
+    
+def graph2Sphere(features,metadata):
+    
+    # Generate equirectangular points and their 3D locations
+    theta, phi = sh.equirec2spherical(metadata.rows, metadata.cols)
+    x,y,z = sh.spherical2xyz(theta,phi)
+    
+    v = torch.stack((x,y,z),dim=1)
+    
+    # Find closest 3D point to each equirectangular point
+    nearest = torch.reshape(knn(metadata.pos3D,v,3)[1],(len(v),3))
+    
+    #Interpolate based on proximty to each node
+    w0 = 1/torch.linalg.norm((v - metadata.pos3D[nearest[:,0]]),dim=1, keepdim=True).to(features.device)
+    w1 = 1/torch.linalg.norm((v - metadata.pos3D[nearest[:,1]]),dim=1, keepdim=True).to(features.device)
+    w2 = 1/torch.linalg.norm((v - metadata.pos3D[nearest[:,2]]),dim=1, keepdim=True).to(features.device)
+    
+    w0 = torch.nan_to_num(w0, nan=1e6)
+    w1 = torch.nan_to_num(w1, nan=1e6)
+    w2 = torch.nan_to_num(w2, nan=1e6)
+    
+    w0 = torch.clamp(w0,0,1e6)
+    w1 = torch.clamp(w1,0,1e6)
+    w2 = torch.clamp(w2,0,1e6)
+    
+    total = w0 + w1 + w2
+    
+    #w0,w1,w2 = mh.getBarycentricWeights(v,metadata.pos3D[nearest[:,0]],metadata.pos3D[nearest[:,1]],metadata.pos3D[nearest[:,2]])
+    
+    #w0 = w0.unsqueeze(1).to(features.device)
+    #w1 = w1.unsqueeze(1).to(features.device)
+    #w2 = w2.unsqueeze(1).to(features.device)
+        
+    result = (w0*features[nearest[:,0]] + w1*features[nearest[:,1]] + w2*features[nearest[:,2]])/total
+    
+    #result = result.clamp(0,1)
+
+    if hasattr(metadata,"mask"):
+        mask = utils.toNumpy(metadata.mask.squeeze(),permute=False)
+        canvas = utils.makeCanvas(result,metadata.original)
+        result = np.reshape(result.data.cpu().numpy(),(metadata.rows,metadata.cols,features.shape[1]))
+        canvas[np.where(mask)] = result[np.where(mask)]
+        return canvas
+    else:
+        return np.reshape(result.data.cpu().numpy(),(metadata.rows,metadata.cols,features.shape[1]))
+
+
+
+def splat2Graph(data, mesh, up_vector = None, N = 100000, ratio=.25, depth = 1, device = 'cpu'):
+    """ Sample mesh faces to determine graph """
+
+    if up_vector == None:
+        up_vector = torch.tensor([[1,1,1]],dtype=torch.float)
+        #up_vector = 2*torch.rand((1,3))-1
+        up_vector = up_vector/torch.linalg.norm(up_vector,dim=1)
+
+
+    #position, normal vector, uv coordinates in the texture map, x is color
+    pos3D, normals = mh.sampleSurface(mesh,N)
+        
+    # Build initial graph
+    #edge_index are neighbors of a point, directions are the directions from that point
+    edge_index,directions = gh.surface2Edges(pos3D,normals,up_vector,k_neighbors=16)
+    #directions need to be turned into selections "W sub n" from the star-like coordinate system from Dr. Hart's github interpolated-selectionconv
+    edge_index,selections,interps = gh.edges2Selections(edge_index,directions,interpolated=True)
+
+    # Generate info for downsampled versions of the graph
+    clusters, edge_indexes, selections_list, interps_list = cl.makeSurfaceClusters(pos3D,normals,edge_index,selections,interps,ratio=ratio,up_vector=up_vector,depth=depth,device=device)
+    #clusters, edge_indexes, selections_list, interps_list = cl.makeMeshClusters(pos3D,mesh,edge_index,selections,interps,ratio=ratio,up_vector=up_vector,depth=depth,device=device)
+    
+    # Make final graph and metadata needed for mapping the result after going through the network
+    graph = Data(clusters=clusters,edge_indexes=edge_indexes,selections_list=selections_list,interps_list=interps_list)
+    metadata = Data(original=data,pos3D=pos3D,mesh=mesh)
+    
+    return graph,metadata
+
+def mesh2Graph(data, mesh, up_vector = None, N = 100000, ratio=.25, mask = None, depth = 1, x_only = False, device = 'cpu'):
+    """ Sample mesh faces to determine graph """
+
+    if up_vector == None:
+        up_vector = torch.tensor([[1,1,1]],dtype=torch.float)
+        #up_vector = 2*torch.rand((1,3))-1
+        up_vector = up_vector/torch.linalg.norm(up_vector,dim=1)
+
+    if mask is not None:
+        warn("Masks are not currently implemented for mesh graphs")      
+    
+    #position, normal vector, uv coordinates in the texture map, x is color
+    pos3D, normals, uvs, x = mh.sampleSurface(mesh,N,return_x=True)
+    
+    x = x.to(device)
+    
+    if x_only:
+        warn("x_only returns randomly selected points for mesh2Graph. Do not use with previous graph structures")
+        return x
+    
+    # Build initial graph
+    #edge_index are neighbors of a point, directions are the directions from that point
+    edge_index,directions = gh.surface2Edges(pos3D,normals,up_vector,k_neighbors=16)
+    #directions need to be turned into selections "W sub n" from the star-like coordinate system from Dr. Hart's github interpolated-selectionconv
+    edge_index,selections,interps = gh.edges2Selections(edge_index,directions,interpolated=True)
+
+    # Generate info for downsampled versions of the graph
+    clusters, edge_indexes, selections_list, interps_list = cl.makeSurfaceClusters(pos3D,normals,edge_index,selections,interps,ratio=ratio,up_vector=up_vector,depth=depth,device=device)
+    #clusters, edge_indexes, selections_list, interps_list = cl.makeMeshClusters(pos3D,mesh,edge_index,selections,interps,ratio=ratio,up_vector=up_vector,depth=depth,device=device)
+    
+    # Make final graph and metadata needed for mapping the result after going through the network
+    graph = Data(x=x,clusters=clusters,edge_indexes=edge_indexes,selections_list=selections_list,interps_list=interps_list)
+    metadata = Data(original=data,pos3D=pos3D,uvs=uvs,mesh=mesh)
+    
+    return graph,metadata
+
+def graph2Splat(features,metadata,view3D=False):
+
+    features = features.cpu().numpy()
+    
+    canvas = utils.toNumpy(metadata.original)
+    rows,cols,ch = canvas.shape
+    
+    # Get 2D positions by scaling uv
+    pos2D = metadata.uvs.cpu().numpy()
+    pos2D[:,0] = pos2D[:,0]*cols
+    pos2D[:,1] = 1-pos2D[:,1] # UV puts y=0 at the bottom
+    pos2D[:,1] = pos2D[:,1]*rows
+    
+    # Generate desired points
+    row_space = np.arange(rows)
+    col_space = np.arange(cols)
+    col_image,row_image = np.meshgrid(col_space,row_space)
+    
+    canvas = utils.interpolatePointCloud2D(pos2D,features,col_image,row_image)
+    canvas = np.clip(canvas,0,1)
+
+    if view3D:
+        mesh = mh.setTexture(metadata.mesh,canvas)
+        mesh.show()
+    
+    return canvas
+
+
+def graph2Mesh(features,metadata,view3D=False):
+    
+    features = features.cpu().numpy()
+    
+    canvas = utils.toNumpy(metadata.original)
+    rows,cols,ch = canvas.shape
+    
+    # Get 2D positions by scaling uv
+    pos2D = metadata.uvs.cpu().numpy()
+    pos2D[:,0] = pos2D[:,0]*cols
+    pos2D[:,1] = 1-pos2D[:,1] # UV puts y=0 at the bottom
+    pos2D[:,1] = pos2D[:,1]*rows
+    
+    # Generate desired points
+    row_space = np.arange(rows)
+    col_space = np.arange(cols)
+    col_image,row_image = np.meshgrid(col_space,row_space)
+    
+    canvas = utils.interpolatePointCloud2D(pos2D,features,col_image,row_image)
+    canvas = np.clip(canvas,0,1)
+
+    if view3D:
+        mesh = mh.setTexture(metadata.mesh,canvas)
+        mesh.show()
+    
+    return canvas

graph_networks/LinearStyleTransfer/LICENSE

@@ -0,0 +1,25 @@
+BSD 2-Clause License
+
+Copyright (c) 2018, SunshineAtNoon 
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

graph_networks/LinearStyleTransfer/README.md

@@ -0,0 +1,102 @@
+## Learning Linear Transformations for Fast Image and Video Style Transfer
+**[[Paper]](http://openaccess.thecvf.com/content_CVPR_2019/papers/Li_Learning_Linear_Transformations_for_Fast_Image_and_Video_Style_Transfer_CVPR_2019_paper.pdf)** **[[Project Page]](https://sites.google.com/view/linear-style-transfer-cvpr19/)**
+
+<img src="doc/images/chicago_paste.png" height="149" hspace="5"><img src="doc/images/photo_content.png" height="150" hspace="5"><img src="doc/images/content.gif" height="150" hspace="5">
+<img src="doc/images/chicago_27.png" height="150" hspace="5"><img src="doc/images/in5_result.png" height="150" hspace="5"><img src="doc/images/test.gif" height="150" hspace="5">
+
+## Prerequisites
+- [Pytorch](http://pytorch.org/)
+- [torchvision](https://github.com/pytorch/vision)
+- [opencv](https://opencv.org/) for video generation
+
+**All code tested on Ubuntu 16.04, pytorch 0.4.1, and opencv 3.4.2**
+
+## Style Transfer
+- Clone from github: `git clone https://github.com/sunshineatnoon/LinearStyleTransfer`
+- Download pre-trained models from [google drive](https://drive.google.com/file/d/1H9T5rfXGlGCUh04DGkpkMFbVnmscJAbs/view?usp=sharing).
+- Uncompress to root folder :
+```
+cd LinearStyleTransfer
+unzip models.zip
+rm models.zip
+```
+
+#### Artistic style transfer
+```
+python TestArtistic.py
+```
+or conduct style transfer on relu_31 features
+```
+python TestArtistic.py --vgg_dir models/vgg_r31.pth --decoder_dir models/dec_r31.pth --matrixPath models/r31.pth --layer r31
+```
+
+#### Photo-realistic style transfer
+For photo-realistic style transfer, we need first compile the [pytorch_spn](https://github.com/Liusifei/pytorch_spn) repository.
+```
+cd libs/pytorch_spn
+sh make.sh
+cd ../..
+```
+Then:
+```
+python TestPhotoReal.py
+```
+Note: images with `_filtered.png` as postfix are images filtered by the SPN after style transfer, images with `_smooth.png` as postfix are images post process by a [smooth filter](https://github.com/LouieYang/deep-photo-styletransfer-tf/blob/master/smooth_local_affine.py).
+
+#### Video style transfer
+```
+python TestVideo.py
+```
+
+#### Real-time video demo
+```
+python real-time-demo.py --vgg_dir models/vgg_r31.pth --decoder_dir models/dec_r31.pth --matrixPath models/r31.pth --layer r31
+```
+
+## Model Training
+### Data Preparation
+- MSCOCO
+```
+wget http://msvocds.blob.core.windows.net/coco2014/train2014.zip
+```
+- WikiArt
+  - Either manually download from [kaggle](https://www.kaggle.com/c/painter-by-numbers).
+  - Or install [kaggle-cli](https://github.com/floydwch/kaggle-cli) and download by running:
+  ```
+  kg download -u <username> -p <password> -c painter-by-numbers -f train.zip
+  ```
+
+### Training
+#### Train a style transfer model
+To train a model that transfers relu4_1 features, run:
+```
+python Train.py --vgg_dir models/vgg_r41.pth --decoder_dir models/dec_r41.pth --layer r41 --contentPath PATH_TO_MSCOCO --stylePath PATH_TO_WikiArt --outf OUTPUT_DIR
+```
+or train a model that transfers relu3_1 features:
+```
+python Train.py --vgg_dir models/vgg_r31.pth --decoder_dir models/dec_r31.pth --layer r31 --contentPath PATH_TO_MSCOCO --stylePath PATH_TO_WikiArt --outf OUTPUT_DIR
+```
+Key hyper-parameters:
+- style_layers: which features to compute style loss.
+- style_weight: larger style weight leads to heavier style in transferred images.
+
+Intermediate results and weight will be stored in `OUTPUT_DIR`
+
+#### Train a SPN model to cancel distortions for photo-realistic style transfer
+Run:
+```
+python TrainSPN.py --contentPath PATH_TO_MSCOCO
+```
+
+### Acknowledgement
+- We use the [smooth filter](https://github.com/LouieYang/deep-photo-styletransfer-tf/blob/master/smooth_local_affine.py) by [LouieYang](https://github.com/LouieYang) in the photo-realistic style transfer.
+
+### Citation
+```
+@inproceedings{li2018learning,
+    author = {Li, Xueting and Liu, Sifei and Kautz, Jan and Yang, Ming-Hsuan},
+    title = {Learning Linear Transformations for Fast Arbitrary Style Transfer},
+    booktitle = {IEEE Conference on Computer Vision and Pattern Recognition},
+    year = {2019}
+}
+```

graph_networks/LinearStyleTransfer/TestArtistic.py

@@ -0,0 +1,98 @@
+import os
+import torch
+import argparse
+from libs.Loader import Dataset
+from libs.Matrix import MulLayer
+import torchvision.utils as vutils
+import torch.backends.cudnn as cudnn
+from libs.utils import print_options
+from libs.models import encoder3,encoder4, encoder5
+from libs.models import decoder3,decoder4, decoder5
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--vgg_dir", default='models/vgg_r41.pth',
+                    help='pre-trained encoder path')
+parser.add_argument("--decoder_dir", default='models/dec_r41.pth',
+                    help='pre-trained decoder path')
+parser.add_argument("--matrixPath", default='models/r41.pth',
+                    help='pre-trained model path')
+parser.add_argument("--stylePath", default="./data/style/",
+                    help='path to style image')
+parser.add_argument("--contentPath", default="./data/content/",
+                    help='path to frames')
+parser.add_argument("--outf", default="Artistic/",
+                    help='path to transferred images')
+parser.add_argument("--batchSize", type=int,default=1,
+                    help='batch size')
+parser.add_argument('--loadSize', type=int, default=256,
+                    help='scale image size')
+parser.add_argument('--fineSize', type=int, default=256,
+                    help='crop image size')
+parser.add_argument("--layer", default="r41",
+                    help='which features to transfer, either r31 or r41')
+
+################# PREPARATIONS #################
+opt = parser.parse_args()
+opt.cuda = torch.cuda.is_available()
+print_options(opt)
+
+os.makedirs(opt.outf,exist_ok=True)
+cudnn.benchmark = True
+
+################# DATA #################
+content_dataset = Dataset(opt.contentPath,opt.loadSize,opt.fineSize,test=True)
+content_loader = torch.utils.data.DataLoader(dataset=content_dataset,
+                                             batch_size = opt.batchSize,
+                                             shuffle = False)
+                                             #num_workers = 1)
+style_dataset = Dataset(opt.stylePath,opt.loadSize,opt.fineSize,test=True)
+style_loader = torch.utils.data.DataLoader(dataset=style_dataset,
+                                           batch_size = opt.batchSize,
+                                           shuffle = False)
+                                           #num_workers = 1)
+
+################# MODEL #################
+if(opt.layer == 'r31'):
+    vgg = encoder3()
+    dec = decoder3()
+elif(opt.layer == 'r41'):
+    vgg = encoder4()
+    dec = decoder4()
+matrix = MulLayer(opt.layer)
+vgg.load_state_dict(torch.load(opt.vgg_dir))
+dec.load_state_dict(torch.load(opt.decoder_dir))
+matrix.load_state_dict(torch.load(opt.matrixPath))
+
+################# GLOBAL VARIABLE #################
+contentV = torch.Tensor(opt.batchSize,3,opt.fineSize,opt.fineSize)
+styleV = torch.Tensor(opt.batchSize,3,opt.fineSize,opt.fineSize)
+
+################# GPU  #################
+if(opt.cuda):
+    vgg.cuda()
+    dec.cuda()
+    matrix.cuda()
+    contentV = contentV.cuda()
+    styleV = styleV.cuda()
+
+for ci,(content,contentName) in enumerate(content_loader):
+    contentName = contentName[0]
+    contentV.resize_(content.size()).copy_(content)
+    for sj,(style,styleName) in enumerate(style_loader):
+        styleName = styleName[0]
+        styleV.resize_(style.size()).copy_(style)
+
+        # forward
+        with torch.no_grad():
+            sF = vgg(styleV)
+            cF = vgg(contentV)
+
+            if(opt.layer == 'r41'):
+                feature,transmatrix = matrix(cF[opt.layer],sF[opt.layer])
+            else:
+                feature,transmatrix = matrix(cF,sF)
+            transfer = dec(feature)
+
+        transfer = transfer.clamp(0,1)
+        vutils.save_image(transfer,'%s/%s_%s.png'%(opt.outf,contentName,styleName),normalize=True,scale_each=True,nrow=opt.batchSize)
+        print('Transferred image saved at %s%s_%s.png'%(opt.outf,contentName,styleName))

graph_networks/LinearStyleTransfer/TestPhotoReal.py

@@ -0,0 +1,118 @@
+import os
+import cv2
+import time
+import torch
+import argparse
+import numpy as np
+from PIL import Image
+from libs.SPN import SPN
+import torchvision.utils as vutils
+from libs.utils import print_options
+from libs.MatrixTest import MulLayer
+import torch.backends.cudnn as cudnn
+from libs.LoaderPhotoReal import Dataset
+from libs.models import encoder3,encoder4
+from libs.models import decoder3,decoder4
+import torchvision.transforms as transforms
+from libs.smooth_filter import smooth_filter
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--vgg_dir", default='models/vgg_r41.pth',
+                    help='pre-trained encoder path')
+parser.add_argument("--decoder_dir", default='models/dec_r41.pth',
+                    help='pre-trained decoder path')
+parser.add_argument("--matrixPath", default='models/r41.pth',
+                    help='pre-trained model path')
+parser.add_argument("--stylePath", default="data/photo_real/style/images/",
+                    help='path to style image')
+parser.add_argument("--styleSegPath", default="data/photo_real/styleSeg/",
+                    help='path to style image masks')
+parser.add_argument("--contentPath", default="data/photo_real/content/images/",
+                    help='path to content image')
+parser.add_argument("--contentSegPath", default="data/photo_real/contentSeg/",
+                    help='path to content image masks')
+parser.add_argument("--outf", default="PhotoReal/",
+                    help='path to save output images')
+parser.add_argument("--batchSize", type=int,default=1,
+                    help='batch size')
+parser.add_argument('--fineSize', type=int, default=512,
+                    help='image size')
+parser.add_argument("--layer", default="r41",
+                    help='features of which layer to transform, either r31 or r41')
+parser.add_argument("--spn_dir", default='models/r41_spn.pth',
+                    help='path to pretrained SPN model')
+
+################# PREPARATIONS #################
+opt = parser.parse_args()
+opt.cuda = torch.cuda.is_available()
+print_options(opt)
+
+os.makedirs(opt.outf, exist_ok=True)
+
+cudnn.benchmark = True
+
+################# DATA #################
+dataset = Dataset(opt.contentPath,opt.stylePath,opt.contentSegPath,opt.styleSegPath,opt.fineSize)
+loader = torch.utils.data.DataLoader(dataset=dataset,
+                                     batch_size=1,
+                                     shuffle=False)
+
+################# MODEL #################
+if(opt.layer == 'r31'):
+    vgg = encoder3()
+    dec = decoder3()
+elif(opt.layer == 'r41'):
+    vgg = encoder4()
+    dec = decoder4()
+matrix = MulLayer(opt.layer)
+vgg.load_state_dict(torch.load(opt.vgg_dir))
+dec.load_state_dict(torch.load(opt.decoder_dir))
+matrix.load_state_dict(torch.load(opt.matrixPath))
+spn = SPN()
+spn.load_state_dict(torch.load(opt.spn_dir))
+
+################# GLOBAL VARIABLE #################
+contentV = torch.Tensor(opt.batchSize,3,opt.fineSize,opt.fineSize)
+styleV = torch.Tensor(opt.batchSize,3,opt.fineSize,opt.fineSize)
+whitenV = torch.Tensor(opt.batchSize,3,opt.fineSize,opt.fineSize)
+
+################# GPU  #################
+if(opt.cuda):
+    vgg.cuda()
+    dec.cuda()
+    spn.cuda()
+    matrix.cuda()
+    contentV = contentV.cuda()
+    styleV = styleV.cuda()
+    whitenV = whitenV.cuda()
+
+for i,(contentImg,styleImg,whitenImg,cmasks,smasks,imname) in enumerate(loader):
+    imname = imname[0]
+    contentV.resize_(contentImg.size()).copy_(contentImg)
+    styleV.resize_(styleImg.size()).copy_(styleImg)
+    whitenV.resize_(whitenImg.size()).copy_(whitenImg)
+
+    # forward
+    sF = vgg(styleV)
+    cF = vgg(contentV)
+
+    with torch.no_grad():
+        if(opt.layer == 'r41'):
+            feature = matrix(cF[opt.layer],sF[opt.layer],cmasks,smasks)
+        else:
+            feature = matrix(cF,sF,cmasks,smasks)
+        transfer = dec(feature)
+        filtered = spn(transfer,whitenV)
+    vutils.save_image(transfer,os.path.join(opt.outf,'%s_transfer.png'%(imname.split('.')[0])))
+
+    filtered = filtered.clamp(0,1)
+    filtered = filtered.cpu()
+    vutils.save_image(filtered,'%s/%s_filtered.png'%(opt.outf,imname.split('.')[0]))
+    out_img = filtered.squeeze(0).mul(255).clamp(0,255).byte().permute(1,2,0).cpu().numpy()
+    content = contentImg.squeeze(0).mul(255).clamp(0,255).byte().permute(1,2,0).cpu().numpy()
+    content = content.copy()
+    out_img = out_img.copy()
+    smoothed = smooth_filter(out_img, content, f_radius=15, f_edge=1e-1)
+    smoothed.save('%s/%s_smooth.png'%(opt.outf,imname.split('.')[0]))
+    print('Transferred image saved at %s%s, filtered image saved at %s%s_filtered.png' \
+          %(opt.outf,imname,opt.outf,imname.split('.')[0]))

graph_networks/LinearStyleTransfer/TestVideo.py

@@ -0,0 +1,108 @@
+import os
+import torch
+import argparse
+from PIL import Image
+from libs.Loader import Dataset
+from libs.Matrix import MulLayer
+import torch.backends.cudnn as cudnn
+from libs.models import encoder3,encoder4
+from libs.models import decoder3,decoder4
+import torchvision.transforms as transforms
+from libs.utils import makeVideo, print_options
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--vgg_dir", default='models/vgg_r31.pth',
+                    help='pre-trained encoder path')
+parser.add_argument("--decoder_dir", default='models/dec_r31.pth',
+                    help='pre-trained decoder path')
+parser.add_argument("--matrix_dir", default="models/r31.pth",
+                    help='path to pre-trained model')
+parser.add_argument("--style", default="data/style/in2.jpg",
+                    help='path to style image')
+parser.add_argument("--content_dir", default="data/videos/content/mountain_2/",
+                    help='path to video frames')
+parser.add_argument('--loadSize', type=int, default=512,
+                    help='scale image size')
+parser.add_argument('--fineSize', type=int, default=512,
+                    help='crop image size')
+parser.add_argument("--name",default="transferred_video",
+                    help="name of generated video")
+parser.add_argument("--layer",default="r31",
+                    help="features of which layer to transform")
+parser.add_argument("--outf",default="videos",
+                    help="output folder")
+
+################# PREPARATIONS #################
+opt = parser.parse_args()
+opt.cuda = torch.cuda.is_available()
+print_options(opt)
+
+os.makedirs(opt.outf,exist_ok=True)
+cudnn.benchmark = True
+
+################# DATA #################
+def loadImg(imgPath):
+    img = Image.open(imgPath).convert('RGB')
+    transform = transforms.Compose([
+                transforms.Scale(opt.fineSize),
+                transforms.ToTensor()])
+    return transform(img)
+styleV = loadImg(opt.style).unsqueeze(0)
+
+content_dataset = Dataset(opt.content_dir,
+                          loadSize = opt.loadSize,
+                          fineSize = opt.fineSize,
+                          test     = True,
+                          video    = True)
+content_loader = torch.utils.data.DataLoader(dataset    = content_dataset,
+					                         batch_size = 1,
+				 	                         shuffle    = False)
+
+################# MODEL #################
+if(opt.layer == 'r31'):
+    vgg = encoder3()
+    dec = decoder3()
+elif(opt.layer == 'r41'):
+    vgg = encoder4()
+    dec = decoder4()
+matrix = MulLayer(layer=opt.layer)
+vgg.load_state_dict(torch.load(opt.vgg_dir))
+dec.load_state_dict(torch.load(opt.decoder_dir))
+matrix.load_state_dict(torch.load(opt.matrix_dir))
+
+################# GLOBAL VARIABLE #################
+contentV = torch.Tensor(1,3,opt.fineSize,opt.fineSize)
+
+################# GPU  #################
+if(opt.cuda):
+    vgg.cuda()
+    dec.cuda()
+    matrix.cuda()
+
+    styleV = styleV.cuda()
+    contentV = contentV.cuda()
+
+result_frames = []
+contents = []
+style = styleV.squeeze(0).cpu().numpy()
+sF = vgg(styleV)
+
+for i,(content,contentName) in enumerate(content_loader):
+    print('Transfer frame %d...'%i)
+    contentName = contentName[0]
+    contentV.resize_(content.size()).copy_(content)
+    contents.append(content.squeeze(0).float().numpy())
+    # forward
+    with torch.no_grad():
+        cF = vgg(contentV)
+
+        if(opt.layer == 'r41'):
+            feature,transmatrix = matrix(cF[opt.layer],sF[opt.layer])
+        else:
+            feature,transmatrix = matrix(cF,sF)
+        transfer = dec(feature)
+
+    transfer = transfer.clamp(0,1)
+    result_frames.append(transfer.squeeze(0).cpu().numpy())
+
+makeVideo(contents,style,result_frames,opt.outf)

graph_networks/LinearStyleTransfer/Train.py

@@ -0,0 +1,185 @@
+import os
+import torch
+import argparse
+import torch.nn as nn
+import torch.optim as optim
+from libs.Loader import Dataset
+from libs.Matrix import MulLayer
+import torchvision.utils as vutils
+import torch.backends.cudnn as cudnn
+from libs.utils import print_options
+from libs.Criterion import LossCriterion
+from libs.models import encoder3,encoder4
+from libs.models import decoder3,decoder4
+from libs.models import encoder5 as loss_network
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--vgg_dir", default='models/vgg_r41.pth',
+                    help='pre-trained encoder path')
+parser.add_argument("--loss_network_dir", default='models/vgg_r51.pth',
+                    help='used for loss network')
+parser.add_argument("--decoder_dir", default='models/dec_r41.pth',
+                    help='pre-trained decoder path')
+parser.add_argument("--stylePath", default="/home/xtli/DATA/wikiArt/train/images/",
+                    help='path to wikiArt dataset')
+parser.add_argument("--contentPath", default="/home/xtli/DATA/MSCOCO/train2014/images/",
+                    help='path to MSCOCO dataset')
+parser.add_argument("--outf", default="trainingOutput/",
+                    help='folder to output images and model checkpoints')
+parser.add_argument("--content_layers", default="r41",
+                    help='layers for content')
+parser.add_argument("--style_layers", default="r11,r21,r31,r41",
+                    help='layers for style')
+parser.add_argument("--batchSize", type=int,default=8,
+                    help='batch size')
+parser.add_argument("--niter", type=int,default=100000,
+                    help='iterations to train the model')
+parser.add_argument('--loadSize', type=int, default=300,
+                    help='scale image size')
+parser.add_argument('--fineSize', type=int, default=256,
+                    help='crop image size')
+parser.add_argument("--lr", type=float, default=1e-4,
+                    help='learning rate')
+parser.add_argument("--content_weight", type=float, default=1.0,
+                    help='content loss weight')
+parser.add_argument("--style_weight", type=float, default=0.02,
+                    help='style loss weight')
+parser.add_argument("--log_interval", type=int, default=500,
+                    help='log interval')
+parser.add_argument("--gpu_id", type=int, default=0,
+                    help='which gpu to use')
+parser.add_argument("--save_interval", type=int, default=5000,
+                    help='checkpoint save interval')
+parser.add_argument("--layer", default="r41",
+                    help='which features to transfer, either r31 or r41')
+
+################# PREPARATIONS #################
+opt = parser.parse_args()
+opt.content_layers = opt.content_layers.split(',')
+opt.style_layers = opt.style_layers.split(',')
+opt.cuda = torch.cuda.is_available()
+if(opt.cuda):
+    torch.cuda.set_device(opt.gpu_id)
+
+os.makedirs(opt.outf,exist_ok=True)
+cudnn.benchmark = True
+print_options(opt)
+
+################# DATA #################
+content_dataset = Dataset(opt.contentPath,opt.loadSize,opt.fineSize)
+content_loader_ = torch.utils.data.DataLoader(dataset     = content_dataset,
+                                              batch_size  = opt.batchSize,
+                                              shuffle     = True,
+                                              num_workers = 1,
+                                              drop_last   = True)
+content_loader = iter(content_loader_)
+style_dataset = Dataset(opt.stylePath,opt.loadSize,opt.fineSize)
+style_loader_ = torch.utils.data.DataLoader(dataset     = style_dataset,
+                                            batch_size  = opt.batchSize,
+                                            shuffle     = True,
+                                            num_workers = 1,
+                                            drop_last   = True)
+style_loader = iter(style_loader_)
+
+################# MODEL #################
+vgg5 = loss_network()
+if(opt.layer == 'r31'):
+    matrix = MulLayer('r31')
+    vgg = encoder3()
+    dec = decoder3()
+elif(opt.layer == 'r41'):
+    matrix = MulLayer('r41')
+    vgg = encoder4()
+    dec = decoder4()
+vgg.load_state_dict(torch.load(opt.vgg_dir))
+dec.load_state_dict(torch.load(opt.decoder_dir))
+vgg5.load_state_dict(torch.load(opt.loss_network_dir))
+
+for param in vgg.parameters():
+    param.requires_grad = False
+for param in vgg5.parameters():
+    param.requires_grad = False
+for param in dec.parameters():
+    param.requires_grad = False
+
+################# LOSS & OPTIMIZER #################
+criterion = LossCriterion(opt.style_layers,
+                          opt.content_layers,
+                          opt.style_weight,
+                          opt.content_weight)
+optimizer = optim.Adam(matrix.parameters(), opt.lr)
+
+################# GLOBAL VARIABLE #################
+contentV = torch.Tensor(opt.batchSize,3,opt.fineSize,opt.fineSize)
+styleV = torch.Tensor(opt.batchSize,3,opt.fineSize,opt.fineSize)
+
+################# GPU  #################
+if(opt.cuda):
+    vgg.cuda()
+    dec.cuda()
+    vgg5.cuda()
+    matrix.cuda()
+    contentV = contentV.cuda()
+    styleV = styleV.cuda()
+
+################# TRAINING #################
+def adjust_learning_rate(optimizer, iteration):
+    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = opt.lr / (1+iteration*1e-5)
+
+for iteration in range(1,opt.niter+1):
+    optimizer.zero_grad()
+    try:
+        content,_ = content_loader.next()
+    except IOError:
+        content,_ = content_loader.next()
+    except StopIteration:
+        content_loader = iter(content_loader_)
+        content,_ = content_loader.next()
+    except:
+        continue
+
+    try:
+        style,_ = style_loader.next()
+    except IOError:
+        style,_ = style_loader.next()
+    except StopIteration:
+        style_loader = iter(style_loader_)
+        style,_ = style_loader.next()
+    except:
+        continue
+
+    contentV.resize_(content.size()).copy_(content)
+    styleV.resize_(style.size()).copy_(style)
+
+    # forward
+    sF = vgg(styleV)
+    cF = vgg(contentV)
+
+    if(opt.layer == 'r41'):
+        feature,transmatrix = matrix(cF[opt.layer],sF[opt.layer])
+    else:
+        feature,transmatrix = matrix(cF,sF)
+    transfer = dec(feature)
+
+    sF_loss = vgg5(styleV)
+    cF_loss = vgg5(contentV)
+    tF = vgg5(transfer)
+    loss,styleLoss,contentLoss = criterion(tF,sF_loss,cF_loss)
+
+    # backward & optimization
+    loss.backward()
+    optimizer.step()
+    print('Iteration: [%d/%d] Loss: %.4f contentLoss: %.4f styleLoss: %.4f Learng Rate is %.6f'%
+         (opt.niter,iteration,loss,contentLoss,styleLoss,optimizer.param_groups[0]['lr']))
+
+    adjust_learning_rate(optimizer,iteration)
+
+    if((iteration) % opt.log_interval == 0):
+        transfer = transfer.clamp(0,1)
+        concat = torch.cat((content,style,transfer.cpu()),dim=0)
+        vutils.save_image(concat,'%s/%d.png'%(opt.outf,iteration),normalize=True,scale_each=True,nrow=opt.batchSize)
+
+    if(iteration > 0 and (iteration) % opt.save_interval == 0):
+        torch.save(matrix.state_dict(), '%s/%s.pth' % (opt.outf,opt.layer))

graph_networks/LinearStyleTransfer/TrainSPN.py

@@ -0,0 +1,141 @@
+from __future__ import print_function
+import os
+import argparse
+
+from libs.SPN import SPN
+from libs.Loader import Dataset
+from libs.models import encoder4
+from libs.models import decoder4
+from libs.utils import print_options
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+import torchvision.utils as vutils
+import torch.backends.cudnn as cudnn
+import torchvision.transforms as transforms
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--vgg_dir", default='models/vgg_r41.pth',
+                    help='pre-trained encoder path')
+parser.add_argument("--decoder_dir", default='models/dec_r41.pth',
+                    help='pre-trained decoder path')
+parser.add_argument("--contentPath", default="/home/xtli/DATA/MSCOCO/train2014/images/",
+                    help='path to MSCOCO dataset')
+parser.add_argument("--outf", default="trainingSPNOutput/",
+                    help='folder to output images and model checkpoints')
+parser.add_argument("--layer", default="r41",
+                    help='layers for content')
+parser.add_argument("--batchSize", type=int,default=8,
+                    help='batch size')
+parser.add_argument("--niter", type=int,default=100000,
+                    help='iterations to train the model')
+parser.add_argument('--loadSize', type=int, default=512,
+                    help='scale image size')
+parser.add_argument('--fineSize', type=int, default=256,
+                    help='crop image size')
+parser.add_argument("--lr", type=float, default=1e-3,
+                    help='learning rate')
+parser.add_argument("--log_interval", type=int, default=500,
+                    help='log interval')
+parser.add_argument("--save_interval", type=int, default=5000,
+                    help='checkpoint save interval')
+parser.add_argument("--spn_num", type=int, default=1,
+                    help='number of spn filters')
+
+################# PREPARATIONS #################
+opt = parser.parse_args()
+opt.cuda = torch.cuda.is_available()
+print_options(opt)
+
+
+os.makedirs(opt.outf, exist_ok = True)
+
+cudnn.benchmark = True
+
+################# DATA #################
+content_dataset = Dataset(opt.contentPath,opt.loadSize,opt.fineSize)
+content_loader_ = torch.utils.data.DataLoader(dataset=content_dataset,
+                                              batch_size = opt.batchSize,
+                                              shuffle = True,
+                                              num_workers = 4,
+                                              drop_last = True)
+content_loader = iter(content_loader_)
+
+################# MODEL #################
+spn = SPN(spn=opt.spn_num)
+if(opt.layer == 'r31'):
+    vgg = encoder3()
+    dec = decoder3()
+elif(opt.layer == 'r41'):
+    vgg = encoder4()
+    dec = decoder4()
+vgg.load_state_dict(torch.load(opt.vgg_dir))
+dec.load_state_dict(torch.load(opt.decoder_dir))
+
+for param in vgg.parameters():
+    param.requires_grad = False
+for param in dec.parameters():
+    param.requires_grad = False
+
+################# LOSS & OPTIMIZER #################
+criterion = nn.MSELoss(size_average=False)
+#optimizer_spn = optim.SGD(spn.parameters(), opt.lr)
+optimizer_spn = optim.Adam(spn.parameters(), opt.lr)
+
+################# GLOBAL VARIABLE #################
+contentV = torch.Tensor(opt.batchSize,3,opt.fineSize,opt.fineSize)
+
+################# GPU  #################
+if(opt.cuda):
+    vgg.cuda()
+    dec.cuda()
+    spn.cuda()
+    contentV = contentV.cuda()
+
+################# TRAINING #################
+def adjust_learning_rate(optimizer, iteration):
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = opt.lr / (1+iteration*1e-5)
+
+spn.train()
+for iteration in range(1,opt.niter+1):
+    optimizer_spn.zero_grad()
+    try:
+        content,_ = content_loader.next()
+    except IOError:
+        content,_ = content_loader.next()
+    except StopIteration:
+        content_loader = iter(content_loader_)
+        content,_ = content_loader.next()
+    except:
+        continue
+
+    contentV.resize_(content.size()).copy_(content)
+
+    # forward
+    cF = vgg(contentV)
+    transfer = dec(cF['r41'])
+
+
+    propagated = spn(transfer,contentV)
+    loss = criterion(propagated,contentV)
+
+    # backward & optimization
+    loss.backward()
+    #nn.utils.clip_grad_norm(spn.parameters(), 1000)
+    optimizer_spn.step()
+    print('Iteration: [%d/%d] Loss: %.4f Learng Rate is %.6f'
+         %(opt.niter,iteration,loss,optimizer_spn.param_groups[0]['lr']))
+
+    adjust_learning_rate(optimizer_spn,iteration)
+
+    if((iteration) % opt.log_interval == 0):
+        transfer = transfer.clamp(0,1)
+        propagated = propagated.clamp(0,1)
+        vutils.save_image(transfer,'%s/%d_transfer.png'%(opt.outf,iteration))
+        vutils.save_image(propagated,'%s/%d_propagated.png'%(opt.outf,iteration))
+
+    if(iteration > 0 and (iteration) % opt.save_interval == 0):
+        torch.save(spn.state_dict(), '%s/%s_spn.pth' % (opt.outf,opt.layer))

graph_networks/LinearStyleTransfer/libs/Criterion.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+
+class styleLoss(nn.Module):
+    def forward(self,input,target):
+        ib,ic,ih,iw = input.size()
+        iF = input.view(ib,ic,-1)
+        iMean = torch.mean(iF,dim=2)
+        iCov = GramMatrix()(input)
+
+        tb,tc,th,tw = target.size()
+        tF = target.view(tb,tc,-1)
+        tMean = torch.mean(tF,dim=2)
+        tCov = GramMatrix()(target)
+
+        loss = nn.MSELoss(size_average=False)(iMean,tMean) + nn.MSELoss(size_average=False)(iCov,tCov)
+        return loss/tb
+
+class GramMatrix(nn.Module):
+    def forward(self,input):
+        b, c, h, w = input.size()
+        f = input.view(b,c,h*w) # bxcx(hxw)
+        # torch.bmm(batch1, batch2, out=None)   #
+        # batch1: bxmxp, batch2: bxpxn -> bxmxn #
+        G = torch.bmm(f,f.transpose(1,2)) # f: bxcx(hxw), f.transpose: bx(hxw)xc -> bxcxc
+        return G.div_(c*h*w)
+
+class LossCriterion(nn.Module):
+    def __init__(self,style_layers,content_layers,style_weight,content_weight):
+        super(LossCriterion,self).__init__()
+
+        self.style_layers = style_layers
+        self.content_layers = content_layers
+        self.style_weight = style_weight
+        self.content_weight = content_weight
+
+        self.styleLosses = [styleLoss()] * len(style_layers)
+        self.contentLosses = [nn.MSELoss()] * len(content_layers)
+
+    def forward(self,tF,sF,cF):
+        # content loss
+        totalContentLoss = 0
+        for i,layer in enumerate(self.content_layers):
+            cf_i = cF[layer]
+            cf_i = cf_i.detach()
+            tf_i = tF[layer]
+            loss_i = self.contentLosses[i]
+            totalContentLoss += loss_i(tf_i,cf_i)
+        totalContentLoss = totalContentLoss * self.content_weight
+
+        # style loss
+        totalStyleLoss = 0
+        for i,layer in enumerate(self.style_layers):
+            sf_i = sF[layer]
+            sf_i = sf_i.detach()
+            tf_i = tF[layer]
+            loss_i = self.styleLosses[i]
+            totalStyleLoss += loss_i(tf_i,sf_i)
+        totalStyleLoss = totalStyleLoss * self.style_weight
+        loss = totalStyleLoss + totalContentLoss
+
+        return loss,totalStyleLoss,totalContentLoss

graph_networks/LinearStyleTransfer/libs/Loader.py

@@ -0,0 +1,44 @@
+import os
+from PIL import Image
+import torch.utils.data as data
+import torchvision.transforms as transforms
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg"])
+
+def default_loader(path):
+    return Image.open(path).convert('RGB')
+
+class Dataset(data.Dataset):
+    def __init__(self,dataPath,loadSize,fineSize,test=False,video=False):
+        super(Dataset,self).__init__()
+        self.dataPath = dataPath
+        self.image_list = [x for x in os.listdir(dataPath) if is_image_file(x)]
+        self.image_list = sorted(self.image_list)
+        if(video):
+            self.image_list = sorted(self.image_list)
+        if not test:
+            self.transform = transforms.Compose([
+            		         transforms.Resize(fineSize),
+            		         transforms.RandomCrop(fineSize),
+                             transforms.RandomHorizontalFlip(),
+            		         transforms.ToTensor()])
+        else:
+            self.transform = transforms.Compose([
+            		         transforms.Resize(fineSize),
+            		         transforms.ToTensor()])
+
+        self.test = test
+
+    def __getitem__(self,index):
+        dataPath = os.path.join(self.dataPath,self.image_list[index])
+
+        Img = default_loader(dataPath)
+        ImgA = self.transform(Img)
+
+        imgName = self.image_list[index]
+        imgName = imgName.split('.')[0]
+        return ImgA,imgName
+
+    def __len__(self):
+        return len(self.image_list)

graph_networks/LinearStyleTransfer/libs/LoaderPhotoReal.py

@@ -0,0 +1,162 @@
+from PIL import Image
+import torchvision.transforms as transforms
+import torchvision.utils as vutils
+import torch.utils.data as data
+from os import listdir
+from os.path import join
+import numpy as np
+import torch
+import os
+import torch.nn as nn
+from torch.autograd import Variable
+import numpy as np
+from libs.utils import whiten
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg"])
+
+def default_loader(path,fineSize):
+    img = Image.open(path).convert('RGB')
+    w,h = img.size
+    if(w < h):
+        neww = fineSize
+        newh = h * neww / w
+        newh = int(newh / 8) * 8
+    else:
+        newh = fineSize
+        neww = w * newh / h
+        neww = int(neww / 8) * 8
+    img = img.resize((neww,newh))
+    return img
+
+def MaskHelper(seg,color):
+    # green
+    mask = torch.Tensor()
+    if(color == 'green'):
+        mask = torch.lt(seg[0],0.1)
+        mask = torch.mul(mask,torch.gt(seg[1],1-0.1))
+        mask = torch.mul(mask,torch.lt(seg[2],0.1))
+    elif(color == 'black'):
+        mask = torch.lt(seg[0], 0.1)
+        mask = torch.mul(mask,torch.lt(seg[1], 0.1))
+        mask = torch.mul(mask,torch.lt(seg[2], 0.1))
+    elif(color == 'white'):
+        mask = torch.gt(seg[0], 1-0.1)
+        mask = torch.mul(mask,torch.gt(seg[1], 1-0.1))
+        mask = torch.mul(mask,torch.gt(seg[2], 1-0.1))
+    elif(color == 'red'):
+        mask = torch.gt(seg[0], 1-0.1)
+        mask = torch.mul(mask,torch.lt(seg[1], 0.1))
+        mask = torch.mul(mask,torch.lt(seg[2], 0.1))
+    elif(color == 'blue'):
+        mask = torch.lt(seg[0], 0.1)
+        mask = torch.mul(mask,torch.lt(seg[1], 0.1))
+        mask = torch.mul(mask,torch.gt(seg[2], 1-0.1))
+    elif(color == 'yellow'):
+        mask = torch.gt(seg[0], 1-0.1)
+        mask = torch.mul(mask,torch.gt(seg[1], 1-0.1))
+        mask = torch.mul(mask,torch.lt(seg[2], 0.1))
+    elif(color == 'grey'):
+        mask = torch.lt(seg[0], 0.1)
+        mask = torch.mul(mask,torch.lt(seg[1], 0.1))
+        mask = torch.mul(mask,torch.lt(seg[2], 0.1))
+    elif(color == 'lightblue'):
+        mask = torch.lt(seg[0], 0.1)
+        mask = torch.mul(mask,torch.gt(seg[1], 1-0.1))
+        mask = torch.mul(mask,torch.gt(seg[2], 1-0.1))
+    elif(color == 'purple'):
+        mask = torch.gt(seg[0], 1-0.1)
+        mask = torch.mul(mask,torch.lt(seg[1], 0.1))
+        mask = torch.mul(mask,torch.gt(seg[2], 1-0.1))
+    else:
+        print('MaskHelper(): color not recognized, color = ' + color)
+    return mask.float()
+
+def ExtractMask(Seg):
+    # Given segmentation for content and style, we get a list of segmentation for each color
+    '''
+    Test Code:
+        content_masks,style_masks = ExtractMask(contentSegImg,styleSegImg)
+        for i,mask in enumerate(content_masks):
+            vutils.save_image(mask,'samples/content_%d.png' % (i),normalize=True)
+        for i,mask in enumerate(style_masks):
+            vutils.save_image(mask,'samples/style_%d.png' % (i),normalize=True)
+    '''
+    color_codes = ['blue', 'green', 'black', 'white', 'red', 'yellow', 'grey', 'lightblue', 'purple']
+    masks = []
+    for color in color_codes:
+        mask = MaskHelper(Seg,color)
+        masks.append(mask)
+    return masks
+
+def calculate_size(h,w,fineSize):
+    if(h > w):
+        newh = fineSize
+        neww = int(w * 1.0 * newh / h)
+    else:
+        neww = fineSize
+        newh = int(h * 1.0 * neww / w)
+    newh = (newh // 8) * 8
+    neww = (neww // 8) * 8
+    return neww, newh
+
+class Dataset(data.Dataset):
+    def __init__(self,contentPath,stylePath,contentSegPath,styleSegPath,fineSize):
+        super(Dataset,self).__init__()
+        self.contentPath = contentPath
+        self.image_list = [x for x in listdir(contentPath) if is_image_file(x)]
+        self.stylePath = stylePath
+        self.contentSegPath = contentSegPath
+        self.styleSegPath = styleSegPath
+        self.fineSize = fineSize
+
+    def __getitem__(self,index):
+        contentImgPath = os.path.join(self.contentPath,self.image_list[index])
+        styleImgPath = os.path.join(self.stylePath,self.image_list[index])
+        contentImg = default_loader(contentImgPath,self.fineSize)
+        styleImg = default_loader(styleImgPath,self.fineSize)
+
+        try:
+            contentSegImgPath = os.path.join(self.contentSegPath,self.image_list[index])
+            contentSegImg = default_loader(contentSegImgPath,self.fineSize)
+        except :
+            print('no mask provided, fake a whole black one')
+            contentSegImg = Image.new('RGB', (contentImg.size))
+
+        try:
+            styleSegImgPath = os.path.join(self.styleSegPath,self.image_list[index])
+            styleSegImg = default_loader(styleSegImgPath,self.fineSize)
+        except :
+            print('no mask provided, fake a whole black one')
+            styleSegImg = Image.new('RGB', (styleImg.size))
+
+
+        hs, ws = styleImg.size
+        newhs, newws = calculate_size(hs,ws,self.fineSize)
+
+        transform = transforms.Compose([
+        		transforms.Resize((newhs, newws)),
+        		transforms.ToTensor()])
+        # Turning segmentation images into masks
+        styleSegImg = transform(styleSegImg)
+        styleImgArbi = transform(styleImg)
+
+        hc, wc = contentImg.size
+        newhc, newwc = calculate_size(hc,wc,self.fineSize)
+
+        transform = transforms.Compose([
+        		transforms.Resize((newhc, newwc)),
+        		transforms.ToTensor()])
+        contentSegImg = transform(contentSegImg)
+        contentImgArbi = transform(contentImg)
+
+        content_masks = ExtractMask(contentSegImg)
+        style_masks = ExtractMask(styleSegImg)
+
+        ImgW = whiten(contentImgArbi.view(3,-1).double())
+        ImgW = ImgW.view(contentImgArbi.size()).float()
+
+        return contentImgArbi.squeeze(0),styleImgArbi.squeeze(0),ImgW,content_masks,style_masks,self.image_list[index]
+
+    def __len__(self):
+        return len(self.image_list)

graph_networks/LinearStyleTransfer/libs/Matrix.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn as nn
+
+class CNN(nn.Module):
+    def __init__(self,layer,matrixSize=32):
+        super(CNN,self).__init__()
+        if(layer == 'r31'):
+            # 256x64x64
+            self.convs = nn.Sequential(nn.Conv2d(256,128,3,1,1),
+                                       nn.ReLU(inplace=True),
+                                       nn.Conv2d(128,64,3,1,1),
+                                       nn.ReLU(inplace=True),
+                                       nn.Conv2d(64,matrixSize,3,1,1))
+        elif(layer == 'r41'):
+            # 512x32x32
+            self.convs = nn.Sequential(nn.Conv2d(512,256,3,1,1),
+                                       nn.ReLU(inplace=True),
+                                       nn.Conv2d(256,128,3,1,1),
+                                       nn.ReLU(inplace=True),
+                                       nn.Conv2d(128,matrixSize,3,1,1))
+
+        # 32x8x8
+        self.fc = nn.Linear(matrixSize*matrixSize,matrixSize*matrixSize)
+        #self.fc = nn.Linear(32*64,256*256)
+
+    def forward(self,x):
+        out = self.convs(x)
+        # 32x8x8
+        b,c,h,w = out.size()
+        out = out.view(b,c,-1)
+        # 32x64
+        out = torch.bmm(out,out.transpose(1,2)).div(h*w)
+        # 32x32
+        out = out.view(out.size(0),-1)
+        return self.fc(out)
+
+class MulLayer(nn.Module):
+    def __init__(self,layer,matrixSize=32):
+        super(MulLayer,self).__init__()
+        self.snet = CNN(layer,matrixSize)
+        self.cnet = CNN(layer,matrixSize)
+        self.matrixSize = matrixSize
+
+        if(layer == 'r41'):
+            self.compress = nn.Conv2d(512,matrixSize,1,1,0)
+            self.unzip = nn.Conv2d(matrixSize,512,1,1,0)
+        elif(layer == 'r31'):
+            self.compress = nn.Conv2d(256,matrixSize,1,1,0)
+            self.unzip = nn.Conv2d(matrixSize,256,1,1,0)
+        self.transmatrix = None
+
+    def forward(self,cF,sF,trans=True):
+        cFBK = cF.clone()
+        cb,cc,ch,cw = cF.size()
+        cFF = cF.view(cb,cc,-1)
+        cMean = torch.mean(cFF,dim=2,keepdim=True)
+        cMean = cMean.unsqueeze(3)
+        cMean = cMean.expand_as(cF)
+        cF = cF - cMean
+
+        sb,sc,sh,sw = sF.size()
+        sFF = sF.view(sb,sc,-1)
+        sMean = torch.mean(sFF,dim=2,keepdim=True)
+        sMean = sMean.unsqueeze(3)
+        sMeanC = sMean.expand_as(cF)
+        sMeanS = sMean.expand_as(sF)
+        sF = sF - sMeanS
+
+
+        compress_content = self.compress(cF)
+        b,c,h,w = compress_content.size()
+        compress_content = compress_content.view(b,c,-1)
+
+        if(trans):
+            cMatrix = self.cnet(cF)
+            sMatrix = self.snet(sF)
+
+            sMatrix = sMatrix.view(sMatrix.size(0),self.matrixSize,self.matrixSize)
+            cMatrix = cMatrix.view(cMatrix.size(0),self.matrixSize,self.matrixSize)
+            transmatrix = torch.bmm(sMatrix,cMatrix)
+            print(cMatrix)
+            transfeature = torch.bmm(transmatrix,compress_content).view(b,c,h,w)
+            out = self.unzip(transfeature.view(b,c,h,w))
+            out = out + sMeanC
+            return out, transmatrix
+        else:
+            out = self.unzip(compress_content.view(b,c,h,w))
+            out = out + cMean
+            return out

graph_networks/LinearStyleTransfer/libs/MatrixTest.py

@@ -0,0 +1,154 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+import numpy as np
+import cv2
+from torch.autograd import Variable
+import torchvision.utils as vutils
+
+
+class CNN(nn.Module):
+    def __init__(self,layer,matrixSize=32):
+        super(CNN,self).__init__()
+        # 256x64x64
+        if(layer == 'r31'):
+            self.convs = nn.Sequential(nn.Conv2d(256,128,3,1,1),
+                                       nn.ReLU(inplace=True),
+                                       nn.Conv2d(128,64,3,1,1),
+                                       nn.ReLU(inplace=True),
+                                       nn.Conv2d(64,matrixSize,3,1,1))
+        elif(layer == 'r41'):
+            # 512x32x32
+            self.convs = nn.Sequential(nn.Conv2d(512,256,3,1,1),
+                                       nn.ReLU(inplace=True),
+                                       nn.Conv2d(256,128,3,1,1),
+                                       nn.ReLU(inplace=True),
+                                       nn.Conv2d(128,matrixSize,3,1,1))
+        self.fc = nn.Linear(32*32,32*32)
+
+    def forward(self,x,masks,style=False):
+        color_code_number = 9
+        xb,xc,xh,xw = x.size()
+        x = x.view(xc,-1)
+        feature_sub_mean = x.clone()
+        for i in range(color_code_number):
+            mask = masks[i].clone().squeeze(0)
+            mask = cv2.resize(mask.numpy(),(xw,xh),interpolation=cv2.INTER_NEAREST)
+            mask = torch.FloatTensor(mask)
+            mask = mask.long()
+            if(torch.sum(mask) >= 10):
+                mask = mask.view(-1)
+
+                # dilation here
+                """
+                kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(5,5))
+                mask = mask.cpu().numpy()
+                mask = cv2.dilate(mask.astype(np.float32), kernel)
+                mask = torch.from_numpy(mask)
+                mask = mask.squeeze()
+                """
+
+                fgmask = (mask>0).nonzero().squeeze(1)
+                fgmask = fgmask.cuda()
+                selectFeature = torch.index_select(x,1,fgmask) # 32x96
+                # subtract mean
+                f_mean = torch.mean(selectFeature,1)
+                f_mean = f_mean.unsqueeze(1).expand_as(selectFeature)
+                selectFeature = selectFeature - f_mean
+                feature_sub_mean.index_copy_(1,fgmask,selectFeature)
+
+        feature = self.convs(feature_sub_mean.view(xb,xc,xh,xw))
+        # 32x16x16
+        b,c,h,w = feature.size()
+        transMatrices = {}
+        feature = feature.view(c,-1)
+
+        for i in range(color_code_number):
+            mask = masks[i].clone().squeeze(0)
+            mask = cv2.resize(mask.numpy(),(w,h),interpolation=cv2.INTER_NEAREST)
+            mask = torch.FloatTensor(mask)
+            mask = mask.long()
+            if(torch.sum(mask) >= 10):
+                mask = mask.view(-1)
+                fgmask = Variable((mask==1).nonzero().squeeze(1))
+                fgmask = fgmask.cuda()
+                selectFeature = torch.index_select(feature,1,fgmask) # 32x96
+                tc,tN = selectFeature.size()
+
+                covMatrix = torch.mm(selectFeature,selectFeature.transpose(0,1)).div(tN)
+                transmatrix = self.fc(covMatrix.view(-1))
+                transMatrices[i] = transmatrix
+        return transMatrices,feature_sub_mean
+
+class MulLayer(nn.Module):
+    def __init__(self,layer,matrixSize=32):
+        super(MulLayer,self).__init__()
+        self.snet = CNN(layer)
+        self.cnet = CNN(layer)
+        self.matrixSize = matrixSize
+
+        if(layer == 'r41'):
+            self.compress = nn.Conv2d(512,matrixSize,1,1,0)
+            self.unzip = nn.Conv2d(matrixSize,512,1,1,0)
+        elif(layer == 'r31'):
+            self.compress = nn.Conv2d(256,matrixSize,1,1,0)
+            self.unzip = nn.Conv2d(matrixSize,256,1,1,0)
+
+    def forward(self,cF,sF,cmasks,smasks):
+
+        sb,sc,sh,sw = sF.size()
+
+        sMatrices,sF_sub_mean = self.snet(sF,smasks,style=True)
+        cMatrices,cF_sub_mean = self.cnet(cF,cmasks,style=False)
+
+        compress_content = self.compress(cF_sub_mean.view(cF.size()))
+        cb,cc,ch,cw = compress_content.size()
+        compress_content = compress_content.view(cc,-1)
+        transfeature = compress_content.clone()
+        color_code_number = 9
+        finalSMean = Variable(torch.zeros(cF.size()).cuda(0))
+        finalSMean = finalSMean.view(sc,-1)
+        for i in range(color_code_number):
+            cmask = cmasks[i].clone().squeeze(0)
+            smask = smasks[i].clone().squeeze(0)
+
+            cmask = cv2.resize(cmask.numpy(),(cw,ch),interpolation=cv2.INTER_NEAREST)
+            cmask = torch.FloatTensor(cmask)
+            cmask = cmask.long()
+            smask = cv2.resize(smask.numpy(),(sw,sh),interpolation=cv2.INTER_NEAREST)
+            smask = torch.FloatTensor(smask)
+            smask = smask.long()
+            if(torch.sum(cmask) >= 10 and torch.sum(smask) >= 10
+               and (i in sMatrices) and (i in cMatrices)):
+                cmask = cmask.view(-1)
+                fgcmask = Variable((cmask==1).nonzero().squeeze(1))
+                fgcmask = fgcmask.cuda()
+
+                smask = smask.view(-1)
+                fgsmask = Variable((smask==1).nonzero().squeeze(1))
+                fgsmask = fgsmask.cuda()
+
+                sFF = sF.view(sc,-1)
+                sFF_select = torch.index_select(sFF,1,fgsmask)
+                sMean = torch.mean(sFF_select,dim=1,keepdim=True)
+                sMean = sMean.view(1,sc,1,1)
+                sMean = sMean.expand_as(cF)
+
+                sMatrix = sMatrices[i]
+                cMatrix = cMatrices[i]
+
+                sMatrix = sMatrix.view(self.matrixSize,self.matrixSize)
+                cMatrix = cMatrix.view(self.matrixSize,self.matrixSize)
+
+                transmatrix = torch.mm(sMatrix,cMatrix) # (C*C)
+
+                compress_content_select = torch.index_select(compress_content,1,fgcmask)
+
+                transfeatureFG = torch.mm(transmatrix,compress_content_select)
+                transfeature.index_copy_(1,fgcmask,transfeatureFG)
+
+                sMean = sMean.contiguous()
+                sMean_select = torch.index_select(sMean.view(sc,-1),1,fgcmask)
+                finalSMean.index_copy_(1,fgcmask,sMean_select)
+        out = self.unzip(transfeature.view(cb,cc,ch,cw))
+        return out + finalSMean.view(out.size())

graph_networks/LinearStyleTransfer/libs/SPN.py

@@ -0,0 +1,156 @@
+import torch
+import torch.nn as nn
+from torchvision.models import vgg16
+from torch.autograd import Variable
+from collections import OrderedDict
+import torch.nn.functional as F
+import sys
+sys.path.append('../')
+from libs.pytorch_spn.modules.gaterecurrent2dnoind import GateRecurrent2dnoind
+
+class spn_block(nn.Module):
+    def __init__(self, horizontal, reverse):
+        super(spn_block, self).__init__()
+        self.propagator = GateRecurrent2dnoind(horizontal,reverse)
+
+    def forward(self,x,G1,G2,G3):
+        sum_abs = G1.abs() + G2.abs() + G3.abs()
+        sum_abs.data[sum_abs.data == 0] = 1e-6
+        mask_need_norm = sum_abs.ge(1)
+        mask_need_norm = mask_need_norm.float()
+        G1_norm = torch.div(G1, sum_abs)
+        G2_norm = torch.div(G2, sum_abs)
+        G3_norm = torch.div(G3, sum_abs)
+
+        G1 = torch.add(-mask_need_norm, 1) * G1 + mask_need_norm * G1_norm
+        G2 = torch.add(-mask_need_norm, 1) * G2 + mask_need_norm * G2_norm
+        G3 = torch.add(-mask_need_norm, 1) * G3 + mask_need_norm * G3_norm
+
+        return self.propagator(x,G1,G2,G3)
+
+class VGG(nn.Module):
+    def __init__(self,nf):
+        super(VGG,self).__init__()
+        self.conv1 = nn.Conv2d(3,nf,3,padding = 1)
+        # 256 x 256
+        self.pool1 = nn.MaxPool2d(kernel_size = 3, stride = 2,padding=1)
+        self.conv2 = nn.Conv2d(nf,nf*2,3,padding = 1)
+        # 128 x 128
+        self.pool2 = nn.MaxPool2d(kernel_size = 3, stride = 2,padding=1)
+        self.conv3 = nn.Conv2d(nf*2,nf*4,3,padding = 1)
+        # 64 x 64
+        self.pool3 = nn.MaxPool2d(kernel_size = 3, stride = 2,padding=1)
+        # 32 x 32
+        self.conv4 = nn.Conv2d(nf*4,nf*8,3,padding = 1)
+
+    def forward(self,x):
+        output = {}
+        output['conv1'] = self.conv1(x)
+        x = F.relu(output['conv1'])
+        x = self.pool1(x)
+        output['conv2'] = self.conv2(x)
+        # 128 x 128
+        x = F.relu(output['conv2'])
+        x = self.pool2(x)
+        output['conv3'] = self.conv3(x)
+        # 64 x 64
+        x = F.relu(output['conv3'])
+        output['pool3'] = self.pool3(x)
+        # 32 x 32
+        output['conv4'] = self.conv4(output['pool3'])
+        return output
+
+class Decoder(nn.Module):
+    def __init__(self,nf=32,spn=1):
+        super(Decoder,self).__init__()
+        # 32 x 32
+        self.layer0 = nn.Conv2d(nf*8,nf*4,1,1,0) # edge_conv5
+        self.layer1 = nn.Upsample(scale_factor=2,mode='bilinear')
+        self.layer2 = nn.Sequential(nn.Conv2d(nf*4,nf*4,3,1,1), # edge_conv8
+                                    nn.ELU(inplace=True))
+        # 64 x 64
+        self.layer3 = nn.Upsample(scale_factor=2,mode='bilinear')
+        self.layer4 = nn.Sequential(nn.Conv2d(nf*4,nf*2,3,1,1), # edge_conv8
+                                    nn.ELU(inplace=True))
+        # 128 x 128
+        self.layer5 = nn.Upsample(scale_factor=2,mode='bilinear')
+        self.layer6 = nn.Sequential(nn.Conv2d(nf*2,nf,3,1,1), # edge_conv8
+                                    nn.ELU(inplace=True))
+        if(spn == 1):
+            self.layer7 = nn.Conv2d(nf,nf*12,3,1,1)
+        else:
+            self.layer7 = nn.Conv2d(nf,nf*24,3,1,1)
+        self.spn = spn
+        # 256 x 256
+
+    def forward(self,encode_feature):
+        output = {}
+        output['0'] = self.layer0(encode_feature['conv4'])
+        output['1'] = self.layer1(output['0'])
+
+        output['2'] = self.layer2(output['1'])
+        output['2res'] = output['2'] + encode_feature['conv3']
+        # 64 x 64
+
+        output['3'] = self.layer3(output['2res'])
+        output['4'] = self.layer4(output['3'])
+        output['4res'] = output['4'] + encode_feature['conv2']
+        # 128 x 128
+
+        output['5'] = self.layer5(output['4res'])
+        output['6'] = self.layer6(output['5'])
+        output['6res'] = output['6'] + encode_feature['conv1']
+
+        output['7'] = self.layer7(output['6res'])
+
+        return output['7']
+
+
+class SPN(nn.Module):
+    def __init__(self,nf=32,spn=1):
+        super(SPN,self).__init__()
+        # conv for mask
+        self.mask_conv = nn.Conv2d(3,nf,3,1,1)
+
+        # guidance network
+        self.encoder = VGG(nf)
+        self.decoder = Decoder(nf,spn)
+
+        # spn blocks
+        self.left_right = spn_block(True,False)
+        self.right_left = spn_block(True,True)
+        self.top_down = spn_block(False, False)
+        self.down_top = spn_block(False,True)
+
+        # post upsample
+        self.post = nn.Conv2d(nf,3,3,1,1)
+        self.nf = nf
+
+    def forward(self,x,rgb):
+        # feature for mask
+        X = self.mask_conv(x)
+
+        # guidance
+        features = self.encoder(rgb)
+        guide = self.decoder(features)
+
+        G = torch.split(guide,self.nf,1)
+        out1 = self.left_right(X,G[0],G[1],G[2])
+        out2 = self.right_left(X,G[3],G[4],G[5])
+        out3 = self.top_down(X,G[6],G[7],G[8])
+        out4 = self.down_top(X,G[9],G[10],G[11])
+
+        out = torch.max(out1,out2)
+        out = torch.max(out,out3)
+        out = torch.max(out,out4)
+
+        return self.post(out)
+
+if __name__ == '__main__':
+    spn = SPN()
+    spn = spn.cuda()
+    for i in range(100):
+        x = Variable(torch.Tensor(1,3,256,256)).cuda()
+        rgb = Variable(torch.Tensor(1,3,256,256)).cuda()
+        output = spn(x,rgb)
+        print(output.size())

graph_networks/LinearStyleTransfer/libs/models.py

@@ -0,0 +1,662 @@
+import torch
+import torch.nn as nn
+
+class encoder3(nn.Module):
+    def __init__(self):
+        super(encoder3,self).__init__()
+        # vgg
+        # 224 x 224
+        self.conv1 = nn.Conv2d(3,3,1,1,0)
+        self.reflecPad1 = nn.ReflectionPad2d((1,1,1,1))
+        # 226 x 226
+
+        self.conv2 = nn.Conv2d(3,64,3,1,0)
+        self.relu2 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad3 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv3 = nn.Conv2d(64,64,3,1,0)
+        self.relu3 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.maxPool = nn.MaxPool2d(kernel_size=2,stride=2,return_indices = True)
+        # 112 x 112
+
+        self.reflecPad4 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv4 = nn.Conv2d(64,128,3,1,0)
+        self.relu4 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad5 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv5 = nn.Conv2d(128,128,3,1,0)
+        self.relu5 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.maxPool2 = nn.MaxPool2d(kernel_size=2,stride=2,return_indices = True)
+        # 56 x 56
+
+        self.reflecPad6 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv6 = nn.Conv2d(128,256,3,1,0)
+        self.relu6 = nn.ReLU(inplace=True)
+        # 56 x 56
+    def forward(self,x):
+        out = self.conv1(x)
+        out = self.reflecPad1(out)
+        out = self.conv2(out)
+        out = self.relu2(out)
+        out = self.reflecPad3(out)
+        out = self.conv3(out)
+        pool1 = self.relu3(out)
+        out,pool_idx = self.maxPool(pool1)
+        out = self.reflecPad4(out)
+        out = self.conv4(out)
+        out = self.relu4(out)
+        out = self.reflecPad5(out)
+        out = self.conv5(out)
+        pool2 = self.relu5(out)
+        out,pool_idx2 = self.maxPool2(pool2)
+        out = self.reflecPad6(out)
+        out = self.conv6(out)
+        out = self.relu6(out)
+        return out
+
+class decoder3(nn.Module):
+    def __init__(self):
+        super(decoder3,self).__init__()
+        # decoder
+        self.reflecPad7 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv7 = nn.Conv2d(256,128,3,1,0)
+        self.relu7 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
+        # 112 x 112
+
+        self.reflecPad8 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv8 = nn.Conv2d(128,128,3,1,0)
+        self.relu8 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad9 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv9 = nn.Conv2d(128,64,3,1,0)
+        self.relu9 = nn.ReLU(inplace=True)
+
+        self.unpool2 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 224 x 224
+
+        self.reflecPad10 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv10 = nn.Conv2d(64,64,3,1,0)
+        self.relu10 = nn.ReLU(inplace=True)
+
+        self.reflecPad11 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv11 = nn.Conv2d(64,3,3,1,0)
+
+    def forward(self,x):
+        output = {}
+        out = self.reflecPad7(x)
+        out = self.conv7(out)
+        out = self.relu7(out)
+        out = self.unpool(out)
+        out = self.reflecPad8(out)
+        out = self.conv8(out)
+        out = self.relu8(out)
+        out = self.reflecPad9(out)
+        out = self.conv9(out)
+        out_relu9 = self.relu9(out)
+        out = self.unpool2(out_relu9)
+        out = self.reflecPad10(out)
+        out = self.conv10(out)
+        out = self.relu10(out)
+        out = self.reflecPad11(out)
+        out = self.conv11(out)
+        return out
+
+class encoder4(nn.Module):
+    def __init__(self):
+        super(encoder4,self).__init__()
+        # vgg
+        # 224 x 224
+        self.conv1 = nn.Conv2d(3,3,1,1,0)
+        self.reflecPad1 = nn.ReflectionPad2d((1,1,1,1))
+        # 226 x 226
+
+        self.conv2 = nn.Conv2d(3,64,3,1,0)
+        self.relu2 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad3 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv3 = nn.Conv2d(64,64,3,1,0)
+        self.relu3 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.maxPool = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 112 x 112
+
+        self.reflecPad4 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv4 = nn.Conv2d(64,128,3,1,0)
+        self.relu4 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad5 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv5 = nn.Conv2d(128,128,3,1,0)
+        self.relu5 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.maxPool2 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 56 x 56
+
+        self.reflecPad6 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv6 = nn.Conv2d(128,256,3,1,0)
+        self.relu6 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad7 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv7 = nn.Conv2d(256,256,3,1,0)
+        self.relu7 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad8 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv8 = nn.Conv2d(256,256,3,1,0)
+        self.relu8 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad9 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv9 = nn.Conv2d(256,256,3,1,0)
+        self.relu9 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.maxPool3 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 28 x 28
+
+        self.reflecPad10 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv10 = nn.Conv2d(256,512,3,1,0)
+        self.relu10 = nn.ReLU(inplace=True)
+        # 28 x 28
+    def forward(self,x,sF=None,matrix11=None,matrix21=None,matrix31=None):
+        output = {}
+        out = self.conv1(x)
+        out = self.reflecPad1(out)
+        out = self.conv2(out)
+        output['r11'] = self.relu2(out)
+        out = self.reflecPad7(output['r11'])
+
+        out = self.conv3(out)
+        output['r12'] = self.relu3(out)
+
+        output['p1'] = self.maxPool(output['r12'])
+        out = self.reflecPad4(output['p1'])
+        out = self.conv4(out)
+        output['r21'] = self.relu4(out)
+        out = self.reflecPad7(output['r21'])
+
+        out = self.conv5(out)
+        output['r22'] = self.relu5(out)
+
+        output['p2'] = self.maxPool2(output['r22'])
+        out = self.reflecPad6(output['p2'])
+        out = self.conv6(out)
+        output['r31'] = self.relu6(out)
+        if(matrix31 is not None):
+            feature3,transmatrix3 = matrix31(output['r31'],sF['r31'])
+            out = self.reflecPad7(feature3)
+        else:
+            out = self.reflecPad7(output['r31'])
+        out = self.conv7(out)
+        output['r32'] = self.relu7(out)
+
+        out = self.reflecPad8(output['r32'])
+        out = self.conv8(out)
+        output['r33'] = self.relu8(out)
+
+        out = self.reflecPad9(output['r33'])
+        out = self.conv9(out)
+        output['r34'] = self.relu9(out)
+
+        output['p3'] = self.maxPool3(output['r34'])
+        out = self.reflecPad10(output['p3'])
+        out = self.conv10(out)
+        output['r41'] = self.relu10(out)
+
+        return output
+
+class decoder4(nn.Module):
+    def __init__(self):
+        super(decoder4,self).__init__()
+        # decoder
+        self.reflecPad11 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv11 = nn.Conv2d(512,256,3,1,0)
+        self.relu11 = nn.ReLU(inplace=True)
+        # 28 x 28
+
+        self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
+        # 56 x 56
+
+        self.reflecPad12 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv12 = nn.Conv2d(256,256,3,1,0)
+        self.relu12 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad13 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv13 = nn.Conv2d(256,256,3,1,0)
+        self.relu13 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad14 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv14 = nn.Conv2d(256,256,3,1,0)
+        self.relu14 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad15 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv15 = nn.Conv2d(256,128,3,1,0)
+        self.relu15 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.unpool2 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 112 x 112
+
+        self.reflecPad16 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv16 = nn.Conv2d(128,128,3,1,0)
+        self.relu16 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad17 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv17 = nn.Conv2d(128,64,3,1,0)
+        self.relu17 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.unpool3 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 224 x 224
+
+        self.reflecPad18 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv18 = nn.Conv2d(64,64,3,1,0)
+        self.relu18 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad19 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv19 = nn.Conv2d(64,3,3,1,0)
+
+    def forward(self,x):
+        # decoder
+        out = self.reflecPad11(x)
+        out = self.conv11(out)
+        out = self.relu11(out)
+        out = self.unpool(out)
+        out = self.reflecPad12(out)
+        out = self.conv12(out)
+
+        out = self.relu12(out)
+        out = self.reflecPad13(out)
+        out = self.conv13(out)
+        out = self.relu13(out)
+        out = self.reflecPad14(out)
+        out = self.conv14(out)
+        out = self.relu14(out)
+        out = self.reflecPad15(out)
+        out = self.conv15(out)
+        out = self.relu15(out)
+        out = self.unpool2(out)
+        out = self.reflecPad16(out)
+        out = self.conv16(out)
+        out = self.relu16(out)
+        out = self.reflecPad17(out)
+        out = self.conv17(out)
+        out = self.relu17(out)
+        out = self.unpool3(out)
+        out = self.reflecPad18(out)
+        out = self.conv18(out)
+        out = self.relu18(out)
+        out = self.reflecPad19(out)
+        out = self.conv19(out)
+        return out
+
+class decoder4(nn.Module):
+    def __init__(self):
+        super(decoder4,self).__init__()
+        # decoder
+        self.reflecPad11 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv11 = nn.Conv2d(512,256,3,1,0)
+        self.relu11 = nn.ReLU(inplace=True)
+        # 28 x 28
+
+        self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
+        # 56 x 56
+
+        self.reflecPad12 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv12 = nn.Conv2d(256,256,3,1,0)
+        self.relu12 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad13 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv13 = nn.Conv2d(256,256,3,1,0)
+        self.relu13 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad14 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv14 = nn.Conv2d(256,256,3,1,0)
+        self.relu14 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad15 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv15 = nn.Conv2d(256,128,3,1,0)
+        self.relu15 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.unpool2 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 112 x 112
+
+        self.reflecPad16 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv16 = nn.Conv2d(128,128,3,1,0)
+        self.relu16 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad17 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv17 = nn.Conv2d(128,64,3,1,0)
+        self.relu17 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.unpool3 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 224 x 224
+
+        self.reflecPad18 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv18 = nn.Conv2d(64,64,3,1,0)
+        self.relu18 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad19 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv19 = nn.Conv2d(64,3,3,1,0)
+
+    def forward(self,x):
+        # decoder
+        out = self.reflecPad11(x)
+        out = self.conv11(out)
+        out = self.relu11(out)
+        out = self.unpool(out)
+        out = self.reflecPad12(out)
+        out = self.conv12(out)
+
+        out = self.relu12(out)
+        out = self.reflecPad13(out)
+        out = self.conv13(out)
+        out = self.relu13(out)
+        out = self.reflecPad14(out)
+        out = self.conv14(out)
+        out = self.relu14(out)
+        out = self.reflecPad15(out)
+        out = self.conv15(out)
+        out = self.relu15(out)
+        out = self.unpool2(out)
+        out = self.reflecPad16(out)
+        out = self.conv16(out)
+        out = self.relu16(out)
+        out = self.reflecPad17(out)
+        out = self.conv17(out)
+        out = self.relu17(out)
+        out = self.unpool3(out)
+        out = self.reflecPad18(out)
+        out = self.conv18(out)
+        out = self.relu18(out)
+        out = self.reflecPad19(out)
+        out = self.conv19(out)
+        return out
+
+class encoder5(nn.Module):
+    def __init__(self):
+        super(encoder5,self).__init__()
+        # vgg
+        # 224 x 224
+        self.conv1 = nn.Conv2d(3,3,1,1,0)
+        self.reflecPad1 = nn.ReflectionPad2d((1,1,1,1))
+        # 226 x 226
+
+        self.conv2 = nn.Conv2d(3,64,3,1,0)
+        self.relu2 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad3 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv3 = nn.Conv2d(64,64,3,1,0)
+        self.relu3 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.maxPool = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 112 x 112
+
+        self.reflecPad4 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv4 = nn.Conv2d(64,128,3,1,0)
+        self.relu4 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad5 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv5 = nn.Conv2d(128,128,3,1,0)
+        self.relu5 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.maxPool2 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 56 x 56
+
+        self.reflecPad6 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv6 = nn.Conv2d(128,256,3,1,0)
+        self.relu6 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad7 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv7 = nn.Conv2d(256,256,3,1,0)
+        self.relu7 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad8 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv8 = nn.Conv2d(256,256,3,1,0)
+        self.relu8 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad9 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv9 = nn.Conv2d(256,256,3,1,0)
+        self.relu9 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.maxPool3 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 28 x 28
+
+        self.reflecPad10 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv10 = nn.Conv2d(256,512,3,1,0)
+        self.relu10 = nn.ReLU(inplace=True)
+
+        self.reflecPad11 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv11 = nn.Conv2d(512,512,3,1,0)
+        self.relu11 = nn.ReLU(inplace=True)
+
+        self.reflecPad12 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv12 = nn.Conv2d(512,512,3,1,0)
+        self.relu12 = nn.ReLU(inplace=True)
+
+        self.reflecPad13 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv13 = nn.Conv2d(512,512,3,1,0)
+        self.relu13 = nn.ReLU(inplace=True)
+
+        self.maxPool4 = nn.MaxPool2d(kernel_size=2,stride=2)
+        self.reflecPad14 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv14 = nn.Conv2d(512,512,3,1,0)
+        self.relu14 = nn.ReLU(inplace=True)
+
+    def forward(self,x,sF=None,contentV256=None,styleV256=None,matrix11=None,matrix21=None,matrix31=None):
+        output = {}
+        out = self.conv1(x)
+        out = self.reflecPad1(out)
+        out = self.conv2(out)
+        output['r11'] = self.relu2(out)
+        out = self.reflecPad7(output['r11'])
+
+        #out = self.reflecPad3(output['r11'])
+        out = self.conv3(out)
+        output['r12'] = self.relu3(out)
+
+        output['p1'] = self.maxPool(output['r12'])
+        out = self.reflecPad4(output['p1'])
+        out = self.conv4(out)
+        output['r21'] = self.relu4(out)
+        out = self.reflecPad7(output['r21'])
+
+        #out = self.reflecPad5(output['r21'])
+        out = self.conv5(out)
+        output['r22'] = self.relu5(out)
+
+        output['p2'] = self.maxPool2(output['r22'])
+        out = self.reflecPad6(output['p2'])
+        out = self.conv6(out)
+        output['r31'] = self.relu6(out)
+        if(styleV256 is not None):
+            feature = matrix31(output['r31'],sF['r31'],contentV256,styleV256)
+            out = self.reflecPad7(feature)
+        else:
+            out = self.reflecPad7(output['r31'])
+        out = self.conv7(out)
+        output['r32'] = self.relu7(out)
+
+        out = self.reflecPad8(output['r32'])
+        out = self.conv8(out)
+        output['r33'] = self.relu8(out)
+
+        out = self.reflecPad9(output['r33'])
+        out = self.conv9(out)
+        output['r34'] = self.relu9(out)
+
+        output['p3'] = self.maxPool3(output['r34'])
+        out = self.reflecPad10(output['p3'])
+        out = self.conv10(out)
+        output['r41'] = self.relu10(out)
+
+        out = self.reflecPad11(output['r41'])
+        out = self.conv11(out)
+        output['r42'] = self.relu11(out)
+
+        out = self.reflecPad12(output['r42'])
+        out = self.conv12(out)
+        output['r43'] = self.relu12(out)
+
+        out = self.reflecPad13(output['r43'])
+        out = self.conv13(out)
+        output['r44'] = self.relu13(out)
+
+        output['p4'] = self.maxPool4(output['r44'])
+
+        out = self.reflecPad14(output['p4'])
+        out = self.conv14(out)
+        output['r51'] = self.relu14(out)
+        return output
+
+class decoder5(nn.Module):
+    def __init__(self):
+        super(decoder5,self).__init__()
+
+        # decoder
+        self.reflecPad15 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv15 = nn.Conv2d(512,512,3,1,0)
+        self.relu15 = nn.ReLU(inplace=True)
+
+        self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
+        # 28 x 28
+
+        self.reflecPad16 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv16 = nn.Conv2d(512,512,3,1,0)
+        self.relu16 = nn.ReLU(inplace=True)
+        # 28 x 28
+
+        self.reflecPad17 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv17 = nn.Conv2d(512,512,3,1,0)
+        self.relu17 = nn.ReLU(inplace=True)
+        # 28 x 28
+
+        self.reflecPad18 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv18 = nn.Conv2d(512,512,3,1,0)
+        self.relu18 = nn.ReLU(inplace=True)
+        # 28 x 28
+
+        self.reflecPad19 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv19 = nn.Conv2d(512,256,3,1,0)
+        self.relu19 = nn.ReLU(inplace=True)
+        # 28 x 28
+
+        self.unpool2 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 56 x 56
+
+        self.reflecPad20 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv20 = nn.Conv2d(256,256,3,1,0)
+        self.relu20 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad21 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv21 = nn.Conv2d(256,256,3,1,0)
+        self.relu21 = nn.ReLU(inplace=True)
+
+        self.reflecPad22 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv22 = nn.Conv2d(256,256,3,1,0)
+        self.relu22 = nn.ReLU(inplace=True)
+
+        self.reflecPad23 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv23 = nn.Conv2d(256,128,3,1,0)
+        self.relu23 = nn.ReLU(inplace=True)
+
+        self.unpool3 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 112 X 112
+
+        self.reflecPad24 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv24 = nn.Conv2d(128,128,3,1,0)
+        self.relu24 = nn.ReLU(inplace=True)
+
+        self.reflecPad25 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv25 = nn.Conv2d(128,64,3,1,0)
+        self.relu25 = nn.ReLU(inplace=True)
+
+        self.unpool4 = nn.UpsamplingNearest2d(scale_factor=2)
+
+        self.reflecPad26 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv26 = nn.Conv2d(64,64,3,1,0)
+        self.relu26 = nn.ReLU(inplace=True)
+
+        self.reflecPad27 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv27 = nn.Conv2d(64,3,3,1,0)
+
+    def forward(self,x):
+        # decoder
+        out = self.reflecPad15(x)
+        out = self.conv15(out)
+        out = self.relu15(out)
+        out = self.unpool(out)
+        out = self.reflecPad16(out)
+        out = self.conv16(out)
+        out = self.relu16(out)
+        out = self.reflecPad17(out)
+        out = self.conv17(out)
+        out = self.relu17(out)
+        out = self.reflecPad18(out)
+        out = self.conv18(out)
+        out = self.relu18(out)
+        out = self.reflecPad19(out)
+        out = self.conv19(out)
+        out = self.relu19(out)
+        out = self.unpool2(out)
+        out = self.reflecPad20(out)
+        out = self.conv20(out)
+        out = self.relu20(out)
+        out = self.reflecPad21(out)
+        out = self.conv21(out)
+        out = self.relu21(out)
+        out = self.reflecPad22(out)
+        out = self.conv22(out)
+        out = self.relu22(out)
+        out = self.reflecPad23(out)
+        out = self.conv23(out)
+        out = self.relu23(out)
+        out = self.unpool3(out)
+        out = self.reflecPad24(out)
+        out = self.conv24(out)
+        out = self.relu24(out)
+        out = self.reflecPad25(out)
+        out = self.conv25(out)
+        out = self.relu25(out)
+        out = self.unpool4(out)
+        out = self.reflecPad26(out)
+        out = self.conv26(out)
+        out = self.relu26(out)
+        out = self.reflecPad27(out)
+        out = self.conv27(out)
+        return out

graph_networks/LinearStyleTransfer/libs/pytorch_spn/README.md

@@ -0,0 +1,12 @@
+# pytorch_spn
+To build, install [pytorch](https://github.com/pytorch) and run:
+
+$ sh make.sh
+
+See left_right_demo.py for usage:
+
+$ mv left_right_demo.py ../
+
+$ python left_right_demo.py
+
+The original codes (caffe) and models will be relesed [HERE](https://github.com/Liusifei/caffe-spn.git).

graph_networks/LinearStyleTransfer/libs/pytorch_spn/_ext/gaterecurrent2dnoind/__init__.py

@@ -0,0 +1,15 @@
+
+from torch.utils.ffi import _wrap_function
+from ._gaterecurrent2dnoind import lib as _lib, ffi as _ffi
+
+__all__ = []
+def _import_symbols(locals):
+    for symbol in dir(_lib):
+        fn = getattr(_lib, symbol)
+        if callable(fn):
+            locals[symbol] = _wrap_function(fn, _ffi)
+        else:
+            locals[symbol] = fn
+        __all__.append(symbol)
+
+_import_symbols(locals())

graph_networks/LinearStyleTransfer/libs/pytorch_spn/build.py

@@ -0,0 +1,34 @@
+import os
+import torch
+from torch.utils.ffi import create_extension
+
+this_file = os.path.dirname(__file__)
+
+sources = []
+headers = []
+defines = []
+with_cuda = False
+
+if torch.cuda.is_available():
+    print('Including CUDA code.')
+    sources += ['src/gaterecurrent2dnoind_cuda.c']
+    headers += ['src/gaterecurrent2dnoind_cuda.h']
+    defines += [('WITH_CUDA', None)]
+    with_cuda = True
+
+this_file = os.path.dirname(os.path.realpath(__file__))
+extra_objects = ['src/cuda/gaterecurrent2dnoind_kernel.cu.o']
+extra_objects = [os.path.join(this_file, fname) for fname in extra_objects]
+
+ffi = create_extension(
+    '_ext.gaterecurrent2dnoind',
+    headers=headers,
+    sources=sources,
+    define_macros=defines,
+    relative_to=__file__,
+    with_cuda=with_cuda,
+    extra_objects=extra_objects
+)
+
+if __name__ == '__main__':
+    ffi.build()

graph_networks/LinearStyleTransfer/libs/pytorch_spn/functions/gaterecurrent2dnoind.py

@@ -0,0 +1,47 @@
+import torch
+from torch.autograd import Function
+from .._ext import gaterecurrent2dnoind as gaterecurrent2d
+
+class GateRecurrent2dnoindFunction(Function):
+    def __init__(self, horizontal_, reverse_):
+        self.horizontal = horizontal_
+        self.reverse = reverse_
+
+    def forward(self, X, G1, G2, G3):
+        num, channels, height, width = X.size()
+        output = torch.zeros(num, channels, height, width)
+
+        if not X.is_cuda:
+            print("cpu version is not ready at this time")
+            return 0
+        else:
+            output = output.cuda()
+            gaterecurrent2d.gaterecurrent2dnoind_forward_cuda(self.horizontal,self.reverse, X, G1, G2, G3, output)
+
+            self.X = X
+            self.G1 = G1
+            self.G2 = G2
+            self.G3 = G3
+            self.output = output
+            self.hiddensize = X.size()
+            return output
+
+    def backward(self, grad_output):
+        assert(self.hiddensize is not None and grad_output.is_cuda)
+        num, channels, height, width = self.hiddensize
+
+        grad_X = torch.zeros(num, channels, height, width).cuda()
+        grad_G1 = torch.zeros(num, channels, height, width).cuda()
+        grad_G2 = torch.zeros(num, channels, height, width).cuda()
+        grad_G3 = torch.zeros(num, channels, height, width).cuda()
+
+        gaterecurrent2d.gaterecurrent2dnoind_backward_cuda(self.horizontal, self.reverse, self.output, grad_output, self.X, self.G1, self.G2, self.G3, grad_X, grad_G1, grad_G2, grad_G3)
+
+        del self.hiddensize
+        del self.G1
+        del self.G2
+        del self.G3
+        del self.output
+        del self.X
+
+        return grad_X, grad_G1, grad_G2, grad_G3

graph_networks/LinearStyleTransfer/libs/pytorch_spn/left_right_demo.py

@@ -0,0 +1,46 @@
+"""
+An example of left->right propagation
+
+Other direction settings:
+left->right: Propagator = GateRecurrent2dnoind(True,False)
+right->left: Propagator = GateRecurrent2dnoind(True,True)
+top->bottom: Propagator = GateRecurrent2dnoind(False,False)
+bottom->top: Propagator = GateRecurrent2dnoind(False,True)
+
+X: any signal/feature map to be filtered
+G1~G3: three coefficient maps (e.g., left-top, left-center, left-bottom)
+
+Note:
+1. G1~G3 constitute the affinity, they can be a bounch of output maps coming from any CNN, with the input of any useful known information (e.g., RGB images)
+2. for any pixel i, |G1(i)| + |G2(i)| + |G3(i)| <= 1 is a sufficent condition for model stability (see paper)
+"""
+import torch
+from torch.autograd import Variable
+from pytorch_spn.modules.gaterecurrent2dnoind import GateRecurrent2dnoind
+
+Propagator = GateRecurrent2dnoind(True,False)
+
+X = Variable(torch.randn(1,3,10,10))
+G1 = Variable(torch.randn(1,3,10,10))
+G2 = Variable(torch.randn(1,3,10,10))
+G3 = Variable(torch.randn(1,3,10,10))
+
+sum_abs = G1.abs() + G2.abs() + G3.abs()
+mask_need_norm = sum_abs.ge(1)
+mask_need_norm = mask_need_norm.float()
+G1_norm = torch.div(G1, sum_abs)
+G2_norm = torch.div(G2, sum_abs)
+G3_norm = torch.div(G3, sum_abs)
+
+G1 = torch.add(-mask_need_norm, 1) * G1 + mask_need_norm * G1_norm
+G2 = torch.add(-mask_need_norm, 1) * G2 + mask_need_norm * G2_norm
+G3 = torch.add(-mask_need_norm, 1) * G3 + mask_need_norm * G3_norm
+
+X = X.cuda()
+G1 = G1.cuda()
+G2 = G2.cuda()
+G3 = G3.cuda()
+
+output = Propagator.forward(X,G1,G2,G3)
+print(X)
+print(output)

graph_networks/LinearStyleTransfer/libs/pytorch_spn/make.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+CUDA_PATH=/usr/local/cuda/
+
+cd src/cuda/
+echo "Compiling gaterecurrent2dnoind layer kernels by nvcc..."
+nvcc -c -o gaterecurrent2dnoind_kernel.cu.o gaterecurrent2dnoind_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52
+cd ../../
+python build.py

graph_networks/LinearStyleTransfer/libs/pytorch_spn/modules/__init__.py

@@ -0,0 +1 @@
+

graph_networks/LinearStyleTransfer/libs/pytorch_spn/modules/gaterecurrent2dnoind.py

@@ -0,0 +1,12 @@
+import torch.nn as nn
+from ..functions.gaterecurrent2dnoind import GateRecurrent2dnoindFunction
+
+class GateRecurrent2dnoind(nn.Module):
+    """docstring for ."""
+    def __init__(self, horizontal_, reverse_):
+        super(GateRecurrent2dnoind, self).__init__()
+        self.horizontal = horizontal_
+        self.reverse = reverse_
+
+    def forward(self, X, G1, G2, G3):
+        return GateRecurrent2dnoindFunction(self.horizontal, self.reverse)(X, G1, G2, G3)

graph_networks/LinearStyleTransfer/libs/pytorch_spn/src/cuda/gaterecurrent2dnoind_kernel.cu

@@ -0,0 +1,697 @@
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <stdio.h>
+#include <math.h>
+#include <float.h>
+#include "gaterecurrent2dnoind_kernel.h"
+
+#define CUDA_1D_KERNEL_LOOP(i, n)                            \
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \
+       i += blockDim.x * gridDim.x)
+
+__device__ void get_gate_idx_sf(int h1, int w1, int h2, int w2, int * out, int horizontal, int reverse)
+{
+	if(horizontal && ! reverse) // left -> right
+	{
+		if(w1>w2)
+		{
+			out[0]=h1;
+			out[1]=w1;
+		}
+		else
+		{
+			out[0]=h2;
+			out[1]=w2;
+		}
+	}
+	if(horizontal && reverse)  // right -> left
+	{
+		if(w1<w2)
+		{
+			out[0]=h1;
+			out[1]=w1;
+		}
+		else
+		{
+			out[0]=h2;
+			out[1]=w2;
+		}
+	}
+	if(!horizontal && !reverse)  // top  -> bottom
+	{
+		if(h1>h2)
+		{
+			out[0]=h1;
+			out[1]=w1;
+		}
+		else
+		{
+			out[0]=h2;
+			out[1]=w2;
+		}
+	}
+	if(!horizontal && reverse)  // bottom -> top
+	{
+		if(h1<h2)
+		{
+			out[0]=h1;
+			out[1]=w1;
+		}
+		else
+		{
+			out[0]=h2;
+			out[1]=w2;
+		}
+	}
+
+}
+
+__device__ float get_data_sf(float * data, int num, int channels,int height, int width,int n,int c,int h,int w)
+{
+	if(h<0 || h >=height)
+		return 0;
+	if(w<0 || w >= width)
+		return 0;
+
+	return data[n*channels*height*width + c * height*width + h * width + w];
+}
+
+__device__ void set_data_sf(float * data, int num, int channels,int height, int width,int n,int c,int h,int w, float v)
+{
+	if(h<0 || h >=height)
+		return ;
+	if(w<0 || w >= width)
+		return ;
+
+	data[n*channels*height*width + c * height*width + h * width + w]=v;
+}
+
+__device__ float get_gate_sf(float * data, int num, int channels,int height, int width,int n,int c,int h1,int w1,int h2,int w2,int horizontal,int reverse)
+{
+	if(h1<0 || h1 >=height)
+		return 0;
+	if(w1<0 || w1 >= width)
+		return 0;
+	if(h2<0 || h2 >=height)
+		return 0;
+	if(w2<0 || w2 >= width)
+		return 0;
+	int idx[2];
+
+	get_gate_idx_sf(h1,w1,h2,w2, idx,horizontal, reverse);
+
+	int h = idx[0];
+	int w = idx[1];
+
+	return data[n*channels*height*width + c * height*width + h * width + w];
+}
+
+__device__ void set_gate_sf(float * data, int num, int channels,int height, int width,int n,int c,int h1,int w1,int h2,int w2,int horizontal,int reverse, float v)
+{
+	if(h1<0 || h1 >=height)
+		return ;
+	if(w1<0 || w1 >= width)
+		return ;
+	if(h2<0 || h2 >=height)
+		return ;
+	if(w2<0 || w2 >= width)
+		return ;
+	int idx[2];
+
+	get_gate_idx_sf(h1,w1,h2,w2, idx,horizontal, reverse);
+
+	int h = idx[0];
+	int w = idx[1];
+
+	data[n*channels*height*width + c * height*width + h * width + w]=v;
+}
+
+// we do not use set_gate_add_sf(...) in the caffe implimentation
+// avoid using atomicAdd
+
+__global__ void forward_one_col_left_right( int count, int T, int num,int channels, int height,  int width, float* X,  float* G1,  float* G2, float* G3, float* H, int horizontal, int reverse) {
+  CUDA_1D_KERNEL_LOOP(index, count) {
+
+  	int hc_count = height * channels;
+
+  	int n,c,h,w;
+  	int temp=index;
+  	w = T;
+  	n = temp / hc_count;
+  	temp = temp % hc_count;
+  	c = temp / height;
+  	temp = temp % height;
+  	h = temp;
+
+
+  	float x_data = get_data_sf(X,num,channels,height,width,n,c,h,w);
+
+  	float g_data_1 = get_gate_sf(G1,num,channels,height,width,n,c,h,w,h-1,w-1,horizontal,reverse);
+  	float h_minus1_data_1 = get_data_sf(H,num,channels,height,width,n,c,h-1,w-1);
+  	float h1_minus1 = g_data_1 * h_minus1_data_1;
+
+  	float g_data_2 = get_gate_sf(G2,num,channels,height,width,n,c,h,w,h,w-1,horizontal,reverse);
+  	float h_minus1_data_2 = get_data_sf(H,num,channels,height,width,n,c,h,w-1);
+  	float h2_minus1 = g_data_2 * h_minus1_data_2;
+
+  	float g_data_3 = get_gate_sf(G3,num,channels,height,width,n,c,h,w,h+1,w-1,horizontal,reverse);
+  	float h_minus1_data_3 = get_data_sf(H,num,channels,height,width,n,c,h+1,w-1);
+  	float h3_minus1 = g_data_3 * h_minus1_data_3;
+
+  	float h_hype = h1_minus1 + h2_minus1 + h3_minus1;
+  	float x_hype = (1 - g_data_1 - g_data_2 - g_data_3) * x_data;
+
+  	float h_data = x_hype + h_hype;
+
+  	set_data_sf(H,num,channels,height,width,n,c,h,w,h_data);
+
+  }
+}
+
+__global__ void forward_one_col_right_left( int count, int T, int num,int channels, int height,  int width, float* X,  float* G1,  float* G2, float* G3, float* H,int horizontal,int reverse) {
+  CUDA_1D_KERNEL_LOOP(index, count) {
+
+  	int hc_count = height * channels;
+  	int n,c,h,w;
+  	int temp=index;
+  	w = T;
+  	n = temp / hc_count;
+  	temp = temp % hc_count;
+  	c = temp / height;
+  	temp = temp % height;
+  	h = temp;
+
+  	float x_data = get_data_sf(X,num,channels,height,width,n,c,h,w);
+
+  	float g_data_1 = get_gate_sf(G1,num,channels,height,width,n,c,h,w,h-1,w+1,horizontal,reverse);
+  	float h_minus1_data_1 = get_data_sf(H,num,channels,height,width,n,c,h-1,w+1);
+  	float h1_minus1 = g_data_1 * h_minus1_data_1;
+
+  	float g_data_2 = get_gate_sf(G2,num,channels,height,width,n,c,h,w,h,w+1,horizontal,reverse);
+  	float h_minus1_data_2 = get_data_sf(H,num,channels,height,width,n,c,h,w+1);
+  	float h2_minus1 = g_data_2 * h_minus1_data_2;
+
+  	float g_data_3 = get_gate_sf(G3,num,channels,height,width,n,c,h,w,h+1,w+1,horizontal,reverse);
+  	float h_minus1_data_3 = get_data_sf(H,num,channels,height,width,n,c,h+1,w+1);
+  	float h3_minus1 = g_data_3 * h_minus1_data_3;
+
+  	float h_hype = h1_minus1 + h2_minus1 + h3_minus1;
+  	float x_hype = (1 - g_data_1 - g_data_2 - g_data_3) * x_data;
+
+  	float h_data = x_hype + h_hype;
+
+  	set_data_sf(H,num,channels,height,width,n,c,h,w,h_data);
+
+  }
+}
+
+__global__ void forward_one_row_top_bottom( int count, int T, int num,int channels, int height,  int width, float* X,  float* G1,  float* G2, float* G3, float* H,int horizontal,int reverse) {
+  CUDA_1D_KERNEL_LOOP(index, count) {
+
+  	int wc_count = width * channels;
+
+  	int n,c,h,w;
+  	int temp=index;
+  	h = T;
+  	n = temp / wc_count;
+  	temp = temp % wc_count;
+  	c = temp / width;
+  	temp = temp % width;
+  	w = temp;
+
+
+  	float x_data = get_data_sf(X,num,channels,height,width,n,c,h,w);
+
+
+  	float g_data_1 = get_gate_sf(G1,num,channels,height,width,n,c,h,w,h-1,w-1,horizontal,reverse);
+  	float h_minus1_data_1 = get_data_sf(H,num,channels,height,width,n,c,h-1,w-1);
+  	float h1_minus1 = g_data_1 * h_minus1_data_1;
+
+  	float g_data_2 = get_gate_sf(G2,num,channels,height,width,n,c,h,w,h-1,w,horizontal,reverse);
+  	float h_minus1_data_2 = get_data_sf(H,num,channels,height,width,n,c,h-1,w);
+  	float h2_minus1 = g_data_2 * h_minus1_data_2;
+
+  	float g_data_3 = get_gate_sf(G3,num,channels,height,width,n,c,h,w,h-1,w+1,horizontal,reverse);
+  	float h_minus1_data_3 = get_data_sf(H,num,channels,height,width,n,c,h-1,w+1);
+  	float h3_minus1 = g_data_3 * h_minus1_data_3;
+
+  	float h_hype = h1_minus1 + h2_minus1 + h3_minus1;
+  	float x_hype = (1 - g_data_1 - g_data_2 - g_data_3) * x_data;
+
+  	float h_data = x_hype + h_hype;
+
+  	set_data_sf(H,num,channels,height,width,n,c,h,w,h_data);
+
+  }
+}
+
+__global__ void forward_one_row_bottom_top( int count, int T, int num,int channels, int height,  int width, float* X,  float* G1,  float* G2, float* G3, float* H,int horizontal,int reverse) {
+  CUDA_1D_KERNEL_LOOP(index, count) {
+
+  	int wc_count = width * channels;
+
+  	int n,c,h,w;
+  	int temp=index;
+  	h = T;
+  	n = temp / wc_count;
+  	temp = temp % wc_count;
+  	c = temp / width;
+  	temp = temp % width;
+  	w = temp;
+
+
+  	float x_data = get_data_sf(X,num,channels,height,width,n,c,h,w);
+
+
+  	float g_data_1 = get_gate_sf(G1,num,channels,height,width,n,c,h,w,h+1,w-1,horizontal,reverse);
+  	float h_minus1_data_1 = get_data_sf(H,num,channels,height,width,n,c,h+1,w-1);
+  	float h1_minus1 = g_data_1 * h_minus1_data_1;
+
+
+  	float g_data_2 = get_gate_sf(G2,num,channels,height,width,n,c,h,w,h+1,w,horizontal,reverse);
+  	float h_minus1_data_2 = get_data_sf(H,num,channels,height,width,n,c,h+1,w);
+  	float h2_minus1 = g_data_2 * h_minus1_data_2;
+
+  	float g_data_3 = get_gate_sf(G3,num,channels,height,width,n,c,h,w,h+1,w+1,horizontal,reverse);
+  	float h_minus1_data_3 = get_data_sf(H,num,channels,height,width,n,c,h+1,w+1);
+  	float h3_minus1 = g_data_3 * h_minus1_data_3;
+
+  	float h_hype = h1_minus1 + h2_minus1 + h3_minus1;
+  	float x_hype = (1 - g_data_1 - g_data_2 - g_data_3) * x_data;
+
+  	float h_data = x_hype + h_hype;
+
+  	set_data_sf(H,num,channels,height,width,n,c,h,w,h_data);
+
+  }
+}
+
+
+__global__ void backward_one_col_left_right( int count, int T, int num,int channels, int height,  int width, float* X,  float* G1,  float* G2, float* G3,  float* H, float * X_diff, float * G1_diff,float* G2_diff,float * G3_diff, float * Hdiff,int horizontal,int reverse) {
+  CUDA_1D_KERNEL_LOOP(index, count) {
+
+  	int hc_count = height * channels;
+
+  	int n,c,h,w;
+  	int temp=index;
+
+  	w = T;
+  	n = temp / hc_count;
+  	temp = temp % hc_count;
+  	c = temp / height;
+  	temp = temp % height;
+  	h = temp;
+
+  	float x_data = get_data_sf(X,num,channels,height,width,n,c,h,w);
+
+  	//h(t)_diff = top(t)_diff
+  	float h_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h,w);
+
+  	//h(t)_diff += h(t+1)_diff * g(t+1) if t<T
+  	float add1_h3_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h-1,w+1);
+  	float add1_g3_data = get_gate_sf(G3,num,channels,height,width,n,c,h,w,h-1,w+1,horizontal,reverse);
+
+  	float add1_h2_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h,w+1);
+  	float add1_g2_data = get_gate_sf(G2,num,channels,height,width,n,c,h,w,h,w+1,horizontal,reverse);
+
+  	float add1_h1_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h+1,w+1);
+  	float add1_g1_data = get_gate_sf(G1,num,channels,height,width,n,c,h,w,h+1,w+1,horizontal,reverse);
+
+  	h_diff = h_diff + add1_h3_diff * add1_g3_data + add1_h2_diff * add1_g2_data + add1_h1_diff * add1_g1_data;
+
+
+  	//Hdiff[n*channels*height*width + c*height*width + h*width + w]=0;
+  	set_data_sf(Hdiff,num,channels,height,width,n,c,h,w,h_diff);
+
+
+  	//x(t)_diff=(1-sum(g_date))*h(t)_diff
+    float g1_data =  get_gate_sf(G1,num,channels,height,width,n,c,h,w,h-1,w-1,horizontal,reverse);
+  	float g2_data =  get_gate_sf(G2,num,channels,height,width,n,c,h,w,h,w-1,horizontal,reverse);
+  	float g3_data =  get_gate_sf(G3,num,channels,height,width,n,c,h,w,h+1,w-1,horizontal,reverse);
+
+  	float x_diff = (1- g1_data -g2_data -g3_data) * h_diff;
+  	set_data_sf(X_diff,num,channels,height,width,n,c,h,w,x_diff);
+
+
+  	// g_diff = h_diff * (h_data(t-1) - x_data)
+  	float h1_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h-1,w-1);
+  	float g1_diff = h_diff * (h1_minus1_data - x_data);
+  	set_gate_sf(G1_diff,num,channels,height,width,n,c,h,w,h-1,w-1,horizontal,reverse,g1_diff);
+
+  	float h2_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h,w-1);
+  	float g2_diff = h_diff * (h2_minus1_data - x_data);
+  	set_gate_sf(G2_diff,num,channels,height,width,n,c,h,w,h,w-1,horizontal,reverse,g2_diff);
+
+  	float h3_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h+1,w-1);
+  	float g3_diff = h_diff * (h3_minus1_data - x_data);
+  	set_gate_sf(G3_diff,num,channels,height,width,n,c,h,w,h+1,w-1,horizontal,reverse,g3_diff);
+
+  }
+}
+
+
+__global__ void backward_one_col_right_left( int count, int T, int num,int channels, int height,  int width, float* X, float* G1,  float* G2, float* G3,  float* H, float * X_diff, float * G1_diff,float* G2_diff,float * G3_diff, float * Hdiff,int horizontal,int reverse) {
+  CUDA_1D_KERNEL_LOOP(index, count) {
+
+  	int hc_count = height * channels;
+
+  	int n,c,h,w;
+  	int temp=index;
+
+
+  	w = T;
+  	n = temp / hc_count;
+  	temp = temp % hc_count;
+  	c = temp / height;
+  	temp = temp % height;
+  	h = temp;
+
+
+  	float x_data = get_data_sf(X,num,channels,height,width,n,c,h,w);
+
+  	//h(t)_diff = top(t)_diff
+  	float h_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h,w);
+
+  	///h(t)_diff += h(t+1)_diff * g(t+1) if t<T
+  	float add1_h3_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h-1,w-1);
+  	float add1_g3_data = get_gate_sf(G3,num,channels,height,width,n,c,h,w,h-1,w-1,horizontal,reverse);
+
+  	float add1_h2_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h,w-1);
+  	float add1_g2_data = get_gate_sf(G2,num,channels,height,width,n,c,h,w,h,w-1,horizontal,reverse);
+
+  	float add1_h1_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h+1,w-1);
+  	float add1_g1_data = get_gate_sf(G1,num,channels,height,width,n,c,h,w,h+1,w-1,horizontal,reverse);
+
+  	h_diff = h_diff + add1_h3_diff * add1_g3_data + add1_h2_diff * add1_g2_data + add1_h1_diff * add1_g1_data;
+
+
+  	set_data_sf(Hdiff,num,channels,height,width,n,c,h,w,h_diff);
+
+    float g1_data =  get_gate_sf(G1,num,channels,height,width,n,c,h,w,h-1,w+1,horizontal,reverse);
+  	float g2_data =  get_gate_sf(G2,num,channels,height,width,n,c,h,w,h,w+1,horizontal,reverse);
+  	float g3_data =  get_gate_sf(G3,num,channels,height,width,n,c,h,w,h+1,w+1,horizontal,reverse);
+  	float x_diff = (1- g1_data -g2_data -g3_data) * h_diff;
+  	set_data_sf(X_diff,num,channels,height,width,n,c,h,w,x_diff);
+
+      // g_diff = h_diff * (h_data(t-1) - x_data)
+  	float h1_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h-1,w+1);
+  	float g1_diff = h_diff * (h1_minus1_data - x_data);
+  	set_gate_sf(G1_diff,num,channels,height,width,n,c,h,w,h-1,w+1,horizontal,reverse,g1_diff);
+
+
+  	float h2_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h,w+1);
+  	float g2_diff = h_diff * (h2_minus1_data - x_data);
+  	set_gate_sf(G2_diff,num,channels,height,width,n,c,h,w,h,w+1,horizontal,reverse,g2_diff);
+
+  	float h3_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h+1,w+1);
+  	float g3_diff = h_diff * (h3_minus1_data - x_data);
+  	set_gate_sf(G3_diff,num,channels,height,width,n,c,h,w,h+1,w+1,horizontal,reverse,g3_diff);
+
+
+  }
+}
+
+__global__ void backward_one_row_top_bottom( int count, int T, int num,int channels, int height,  int width, float* X, float* G1,  float* G2, float* G3,  float* H, float * X_diff, float * G1_diff,float* G2_diff,float * G3_diff, float * Hdiff,int horizontal,int reverse) {
+  CUDA_1D_KERNEL_LOOP(index, count) {
+
+
+  	int wc_count = width * channels;
+
+  	int n,c,h,w;
+  	int temp=index;
+  	h = T;
+  	n = temp / wc_count;
+  	temp = temp % wc_count;
+  	c = temp / width;
+  	temp = temp % width;
+  	w = temp;
+
+  	float x_data = get_data_sf(X,num,channels,height,width,n,c,h,w);
+
+  	float h_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h,w);
+
+  	//h(t)_diff += h(t+1)_diff * g(t+1) if t<T
+  	float add1_h3_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h+1,w-1);
+  	float add1_g3_data = get_gate_sf(G3,num,channels,height,width,n,c,h,w,h+1,w-1,horizontal,reverse);
+
+  	float add1_h2_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h+1,w);
+  	float add1_g2_data = get_gate_sf(G2,num,channels,height,width,n,c,h,w,h+1,w,horizontal,reverse);
+
+  	float add1_h1_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h+1,w+1);
+  	float add1_g1_data = get_gate_sf(G1,num,channels,height,width,n,c,h,w,h+1,w+1,horizontal,reverse);
+
+  	h_diff = h_diff + add1_h3_diff * add1_g3_data + add1_h2_diff * add1_g2_data + add1_h1_diff * add1_g1_data;
+
+
+  	set_data_sf(Hdiff,num,channels,height,width,n,c,h,w,h_diff);
+
+
+  	//x(t)_diff=(1-g(t))*h(t)_diff
+  	float g1_data =  get_gate_sf(G1,num,channels,height,width,n,c,h,w,h-1,w-1,horizontal,reverse);
+  	float g2_data =  get_gate_sf(G2,num,channels,height,width,n,c,h,w,h-1,w,horizontal,reverse);
+  	float g3_data =  get_gate_sf(G3,num,channels,height,width,n,c,h,w,h-1,w+1,horizontal,reverse);
+  	float x_diff = (1- g1_data -g2_data -g3_data) * h_diff;
+  	set_data_sf(X_diff,num,channels,height,width,n,c,h,w,x_diff);
+
+
+
+  	// g_diff = h_diff * (h_data(t-1) - x_data)
+  	float h1_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h-1,w-1);
+  	float g1_diff = h_diff * (h1_minus1_data - x_data);
+  	set_gate_sf(G1_diff,num,channels,height,width,n,c,h,w,h-1,w-1,horizontal,reverse,g1_diff);
+
+  	float h2_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h-1,w);
+  	float g2_diff = h_diff * (h2_minus1_data - x_data);
+  	set_gate_sf(G2_diff,num,channels,height,width,n,c,h,w,h-1,w,horizontal,reverse,g2_diff);
+
+  	float h3_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h-1,w+1);
+  	float g3_diff = h_diff * (h3_minus1_data - x_data);
+  	set_gate_sf(G3_diff,num,channels,height,width,n,c,h,w,h-1,w+1,horizontal,reverse,g3_diff);
+
+  }
+}
+
+__global__ void backward_one_row_bottom_top( int count, int T, int num,int channels, int height,  int width,  float* X,  float* G1,  float* G2, float* G3,  float* H, float * X_diff, float * G1_diff,float* G2_diff,float * G3_diff, float * Hdiff,int horizontal,int reverse) {
+  CUDA_1D_KERNEL_LOOP(index, count) {
+
+  	int wc_count = width * channels;
+
+  	int n,c,h,w;
+  	int temp=index;
+  	h = T;
+  	n = temp / wc_count;
+  	temp = temp % wc_count;
+  	c = temp / width;
+  	temp = temp % width;
+  	w = temp;
+
+  	float x_data = get_data_sf(X,num,channels,height,width,n,c,h,w);
+
+  	//h(t)_diff = top(t)_diff
+  	float h_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h,w);
+
+  	//h(t)_diff += h(t+1)_diff * g(t+1) if t<T
+  	float add1_h3_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h-1,w-1);
+  	float add1_g3_data = get_gate_sf(G3,num,channels,height,width,n,c,h,w,h-1,w-1,horizontal,reverse);
+
+  	float add1_h2_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h-1,w);
+  	float add1_g2_data = get_gate_sf(G2,num,channels,height,width,n,c,h,w,h-1,w,horizontal,reverse);
+
+  	float add1_h1_diff = get_data_sf(Hdiff,num,channels,height,width,n,c,h-1,w+1);
+  	float add1_g1_data = get_gate_sf(G1,num,channels,height,width,n,c,h,w,h-1,w+1,horizontal,reverse);
+
+  	h_diff = h_diff + add1_h3_diff * add1_g3_data + add1_h2_diff * add1_g2_data + add1_h1_diff * add1_g1_data;
+
+
+  	set_data_sf(Hdiff,num,channels,height,width,n,c,h,w,h_diff);
+
+
+  	//x(t)_diff=(1-g(t))*h(t)_diff
+  	float g1_data =  get_gate_sf(G1,num,channels,height,width,n,c,h,w,h+1,w-1,horizontal,reverse);
+  	float g2_data =  get_gate_sf(G2,num,channels,height,width,n,c,h,w,h+1,w,horizontal,reverse);
+  	float g3_data =  get_gate_sf(G3,num,channels,height,width,n,c,h,w,h+1,w+1,horizontal,reverse);
+   	float x_diff = (1- g1_data -g2_data -g3_data) * h_diff;
+  	set_data_sf(X_diff,num,channels,height,width,n,c,h,w,x_diff);
+
+
+  	// g_diff = h_diff * (h_data(t-1) - x_data)
+  	float h1_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h+1,w-1);
+  	float g1_diff = h_diff * (h1_minus1_data - x_data);
+  	set_gate_sf(G1_diff,num,channels,height,width,n,c,h,w,h+1,w-1,horizontal,reverse,g1_diff);
+
+  	//float g2_diff = h_diff * g2_idx * x_data * -1;
+  	float h2_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h+1,w);
+  	float g2_diff = h_diff * (h2_minus1_data - x_data);
+  	set_gate_sf(G2_diff,num,channels,height,width,n,c,h,w,h+1,w,horizontal,reverse,g2_diff);
+
+  	//float g3_diff = h_diff * g3_idx * x_data * -1;
+  	float h3_minus1_data = get_data_sf(H,num,channels,height,width,n,c,h+1,w+1);
+  	float g3_diff = h_diff * (h3_minus1_data - x_data);
+  	set_gate_sf(G3_diff,num,channels,height,width,n,c,h,w,h+1,w+1,horizontal,reverse,g3_diff);
+
+
+  }
+}
+
+
+int Forward_left_right(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3, float * H, int horizontal_, int reverse_, cudaStream_t stream)
+{
+   int count = height_ * channels_ * num_;
+   int kThreadsPerBlock = 1024;
+  cudaError_t err;
+
+  for(int t=0; t<width_; t++) {
+    	forward_one_col_left_right<<<(count + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>(count, t, num_, channels_, height_, width_, X, G1, G2, G3, H, horizontal_, reverse_);
+
+      err = cudaGetLastError();
+      if(cudaSuccess != err)
+      {
+          fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) );
+          exit( -1 );
+      }
+  }
+  return 1;
+}
+
+int Forward_right_left(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3, float * H, int horizontal_, int reverse_, cudaStream_t stream)
+{
+   int count = height_ * channels_ * num_;
+   int kThreadsPerBlock = 1024;
+  cudaError_t err;
+
+  for(int t = width_ - 1; t >= 0; t--) {
+    	forward_one_col_right_left<<<(count + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>(count, t, num_, channels_, height_, width_, X, G1, G2, G3, H, horizontal_, reverse_);
+
+      err = cudaGetLastError();
+      if(cudaSuccess != err)
+      {
+          fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) );
+          exit( -1 );
+      }
+  }
+  return 1;
+}
+
+int Forward_top_bottom(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3, float * H, int horizontal_, int reverse_, cudaStream_t stream)
+{
+   int count = width_ * channels_ * num_;
+   int kThreadsPerBlock = 1024;
+  cudaError_t err;
+
+  for(int t=0; t< height_; t++) {
+    	forward_one_row_top_bottom<<<(count + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>(count, t, num_, channels_, height_, width_, X, G1, G2, G3, H, horizontal_, reverse_);
+
+      err = cudaGetLastError();
+      if(cudaSuccess != err)
+      {
+          fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) );
+          exit( -1 );
+      }
+  }
+  return 1;
+}
+
+int Forward_bottom_top(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3, float * H, int horizontal_, int reverse_, cudaStream_t stream)
+{
+   int count = width_ * channels_ * num_;
+   int kThreadsPerBlock = 1024;
+  cudaError_t err;
+
+  for(int t = height_-1; t >= 0; t--) {
+    	forward_one_row_bottom_top<<<(count + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>(count, t, num_, channels_, height_, width_, X, G1, G2, G3, H, horizontal_, reverse_);
+
+      err = cudaGetLastError();
+      if(cudaSuccess != err)
+      {
+          fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) );
+          exit( -1 );
+      }
+  }
+  return 1;
+}
+
+int Backward_left_right(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3,  float * H, float * X_diff, float * G1_diff, float * G2_diff, float * G3_diff, float * H_diff, int horizontal_, int reverse_, cudaStream_t stream)
+{
+   int count =  height_ * channels_ * num_;
+   int kThreadsPerBlock = 1024;
+  cudaError_t err;
+
+  for(int t = width_ -1; t>=0; t--)
+  {
+    backward_one_col_left_right<<<(count + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>(count, t, num_, channels_, height_, width_, X, G1, G2, G3, H, X_diff, G1_diff, G2_diff, G3_diff, H_diff, horizontal_, reverse_);
+
+    err = cudaGetLastError();
+    if(cudaSuccess != err)
+    {
+        fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) );
+        exit( -1 );
+    }
+  }
+  return 1;
+}
+
+int Backward_right_left(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3,  float * H, float * X_diff, float * G1_diff, float * G2_diff, float * G3_diff, float * H_diff, int horizontal_, int reverse_, cudaStream_t stream)
+{
+   int count =  height_ * channels_ * num_;
+   int kThreadsPerBlock = 1024;
+  cudaError_t err;
+
+  for(int t = 0; t<width_; t++)
+  {
+    backward_one_col_right_left<<<(count + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>(count, t, num_, channels_, height_, width_, X, G1, G2, G3, H, X_diff, G1_diff, G2_diff, G3_diff, H_diff, horizontal_, reverse_);
+
+    err = cudaGetLastError();
+    if(cudaSuccess != err)
+    {
+        fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) );
+        exit( -1 );
+    }
+  }
+  return 1;
+}
+
+int Backward_top_bottom(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3,  float * H, float * X_diff, float * G1_diff, float * G2_diff, float * G3_diff, float * H_diff, int horizontal_, int reverse_, cudaStream_t stream)
+{
+   int count =  width_ * channels_ * num_;
+   int kThreadsPerBlock = 1024;
+  cudaError_t err;
+
+  for(int t = height_-1; t>=0; t--)
+  {
+    backward_one_row_top_bottom<<<(count + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>(count, t, num_, channels_, height_, width_, X, G1, G2, G3, H, X_diff, G1_diff, G2_diff, G3_diff, H_diff, horizontal_, reverse_);
+
+    err = cudaGetLastError();
+    if(cudaSuccess != err)
+    {
+        fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) );
+        exit( -1 );
+    }
+  }
+  return 1;
+}
+
+int Backward_bottom_top(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3,  float * H, float * X_diff, float * G1_diff, float * G2_diff, float * G3_diff, float * H_diff, int horizontal_, int reverse_, cudaStream_t stream)
+{
+   int count =  width_ * channels_ * num_;
+   int kThreadsPerBlock = 1024;
+  cudaError_t err;
+
+  for(int t = 0; t<height_; t++)
+  {
+    backward_one_row_bottom_top<<<(count + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>(count, t, num_, channels_, height_, width_, X, G1, G2, G3, H, X_diff, G1_diff, G2_diff, G3_diff, H_diff, horizontal_, reverse_);
+
+    err = cudaGetLastError();
+    if(cudaSuccess != err)
+    {
+        fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) );
+        exit( -1 );
+    }
+  }
+  return 1;
+}
+
+#ifdef __cplusplus
+}
+#endif

graph_networks/LinearStyleTransfer/libs/pytorch_spn/src/cuda/gaterecurrent2dnoind_kernel.h

@@ -0,0 +1,28 @@
+#ifndef _GATERECURRENT2DNOIND_KERNEL
+#define _GATERECURRENT2DNOIND_KERNEL
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+int Forward_left_right(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3, float * H, int horizontal_, int reverse_, cudaStream_t stream);
+
+int Forward_right_left(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3, float * H, int horizontal_, int reverse_, cudaStream_t stream);
+
+int Forward_top_bottom(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3, float * H, int horizontal_, int reverse_, cudaStream_t stream);
+
+int Forward_bottom_top(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3, float * H, int horizontal_, int reverse_, cudaStream_t stream);
+
+int Backward_left_right(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3,  float * H, float * X_diff, float * G1_diff, float * G2_diff, float * G3_diff, float * H_diff, int horizontal_, int reverse_, cudaStream_t stream);
+
+int Backward_right_left(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3,  float * H, float * X_diff, float * G1_diff, float * G2_diff, float * G3_diff, float * H_diff, int horizontal_, int reverse_, cudaStream_t stream);
+
+int Backward_top_bottom(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3,  float * H, float * X_diff, float * G1_diff, float * G2_diff, float * G3_diff, float * H_diff, int horizontal_, int reverse_, cudaStream_t stream);
+
+int Backward_bottom_top(int num_, int channels_, int height_, int width_,  float * X,  float * G1,  float * G2,  float * G3,  float * H, float * X_diff, float * G1_diff, float * G2_diff, float * G3_diff, float * H_diff, int horizontal_, int reverse_, cudaStream_t stream);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

graph_networks/LinearStyleTransfer/libs/pytorch_spn/src/gaterecurrent2dnoind_cuda.c

@@ -0,0 +1,91 @@
+// gaterecurrent2dnoind_cuda.c
+#include <THC/THC.h>
+#include <math.h>
+#include "gaterecurrent2dnoind_cuda.h"
+#include "cuda/gaterecurrent2dnoind_kernel.h"
+
+// typedef bool boolean;
+
+// this symbol will be resolved automatically from PyTorch libs
+extern THCState *state;
+
+int gaterecurrent2dnoind_forward_cuda(int horizontal_, int reverse_, THCudaTensor * X, THCudaTensor * G1, THCudaTensor * G2, THCudaTensor * G3, THCudaTensor * output)
+{
+	// Grab the input tensor to flat
+	float * X_data = THCudaTensor_data(state, X);
+	float * G1_data = THCudaTensor_data(state, G1);
+	float * G2_data = THCudaTensor_data(state, G2);
+	float * G3_data = THCudaTensor_data(state, G3);
+	float * H_data = THCudaTensor_data(state, output);
+
+	// dimensions
+	int num_ = THCudaTensor_size(state, X, 0);
+	int channels_ = THCudaTensor_size(state, X, 1);
+	int height_ = THCudaTensor_size(state, X, 2);
+	int width_ = THCudaTensor_size(state, X, 3);
+
+	cudaStream_t stream = THCState_getCurrentStream(state);
+
+	if(horizontal_ && !reverse_) // left to right
+	{
+		//const int count = height_ * channels_ * num_;
+		Forward_left_right(num_, channels_, height_, width_, X_data, G1_data, G2_data, G3_data, H_data, horizontal_, reverse_, stream);
+	}
+	else if(horizontal_ && reverse_) // right to left
+	{
+		Forward_right_left(num_, channels_, height_, width_, X_data, G1_data, G2_data, G3_data, H_data, horizontal_, reverse_, stream);
+	}
+	else if(!horizontal_ && !reverse_) // top to bottom
+	{
+		Forward_top_bottom(num_, channels_, height_, width_, X_data, G1_data, G2_data, G3_data, H_data, horizontal_, reverse_, stream);
+	}
+	else
+	{
+		Forward_bottom_top(num_, channels_, height_, width_, X_data, G1_data, G2_data, G3_data, H_data, horizontal_, reverse_, stream);
+	}
+
+	return 1;
+}
+
+int gaterecurrent2dnoind_backward_cuda(int horizontal_, int reverse_, THCudaTensor* top, THCudaTensor* top_grad, THCudaTensor * X, THCudaTensor * G1, THCudaTensor * G2, THCudaTensor * G3, THCudaTensor * X_grad, THCudaTensor * G1_grad, THCudaTensor * G2_grad, THCudaTensor * G3_grad)
+{
+	//Grab the input tensor to flat
+	float * X_data = THCudaTensor_data(state, X);
+	float * G1_data = THCudaTensor_data(state, G1);
+	float * G2_data = THCudaTensor_data(state, G2);
+	float * G3_data = THCudaTensor_data(state, G3);
+	float * H_data = THCudaTensor_data(state, top);
+
+	float * H_diff = THCudaTensor_data(state, top_grad);
+
+	float * X_diff = THCudaTensor_data(state, X_grad);
+	float * G1_diff = THCudaTensor_data(state, G1_grad);
+	float * G2_diff = THCudaTensor_data(state, G2_grad);
+	float * G3_diff = THCudaTensor_data(state, G3_grad);
+
+	// dimensions
+	int num_ = THCudaTensor_size(state, X, 0);
+	int channels_ = THCudaTensor_size(state, X, 1);
+	int height_ = THCudaTensor_size(state, X, 2);
+	int width_ = THCudaTensor_size(state, X, 3);
+
+	cudaStream_t stream = THCState_getCurrentStream(state);
+
+	if(horizontal_ && ! reverse_) //left to right
+	{
+		Backward_left_right(num_, channels_, height_, width_, X_data, G1_data, G2_data, G3_data, H_data, X_diff, G1_diff, G2_diff, G3_diff, H_diff, horizontal_, reverse_, stream);
+	}
+	else if(horizontal_ &&  reverse_) //right to left
+	{
+		Backward_right_left(num_, channels_, height_, width_, X_data, G1_data, G2_data, G3_data, H_data, X_diff, G1_diff, G2_diff, G3_diff, H_diff, horizontal_, reverse_, stream);
+	}
+	else if(!horizontal_ &&  !reverse_) //top to bottom
+	{
+		Backward_top_bottom(num_, channels_, height_, width_, X_data, G1_data, G2_data, G3_data, H_data, X_diff, G1_diff, G2_diff, G3_diff, H_diff, horizontal_, reverse_, stream);
+	}
+	else {
+		Backward_bottom_top(num_, channels_, height_, width_, X_data, G1_data, G2_data, G3_data, H_data, X_diff, G1_diff, G2_diff, G3_diff, H_diff, horizontal_, reverse_, stream);
+	}
+
+	return 1;
+}

graph_networks/LinearStyleTransfer/libs/pytorch_spn/src/gaterecurrent2dnoind_cuda.h

@@ -0,0 +1,6 @@
+
+// #include <stdbool.h>
+// gaterecurrent2dnoind_cuda.h
+int gaterecurrent2dnoind_forward_cuda(int horizontal_, int reverse_, THCudaTensor * X, THCudaTensor * G1, THCudaTensor * G2, THCudaTensor * G3, THCudaTensor * output);
+
+int gaterecurrent2dnoind_backward_cuda(int horizontal_, int reverse_, THCudaTensor* top, THCudaTensor* top_grad, THCudaTensor * X, THCudaTensor * G1, THCudaTensor * G2, THCudaTensor * G3, THCudaTensor * X_diff, THCudaTensor * G1_diff, THCudaTensor * G2_diff, THCudaTensor * G3_diff);

graph_networks/LinearStyleTransfer/libs/smooth_filter.py

@@ -0,0 +1,407 @@
+"""
+Code cc from https://github.com/LouieYang/deep-photo-styletransfer-tf/blob/master/smooth_local_affine.py
+"""
+src = '''
+	#include "/usr/local/cuda/include/math_functions.h"
+	#define TB 256
+	#define EPS 1e-7
+
+	__device__ bool InverseMat4x4(double m_in[4][4], double inv_out[4][4]) {
+		double m[16], inv[16];
+		for (int i = 0; i < 4; i++) {
+			for (int j = 0; j < 4; j++) {
+				m[i * 4 + j] = m_in[i][j];
+			}
+		}
+
+	    inv[0] = m[5]  * m[10] * m[15] -
+	             m[5]  * m[11] * m[14] -
+	             m[9]  * m[6]  * m[15] +
+	             m[9]  * m[7]  * m[14] +
+	             m[13] * m[6]  * m[11] -
+	             m[13] * m[7]  * m[10];
+
+	    inv[4] = -m[4]  * m[10] * m[15] +
+	              m[4]  * m[11] * m[14] +
+	              m[8]  * m[6]  * m[15] -
+	              m[8]  * m[7]  * m[14] -
+	              m[12] * m[6]  * m[11] +
+	              m[12] * m[7]  * m[10];
+
+	    inv[8] = m[4]  * m[9] * m[15] -
+	             m[4]  * m[11] * m[13] -
+	             m[8]  * m[5] * m[15] +
+	             m[8]  * m[7] * m[13] +
+	             m[12] * m[5] * m[11] -
+	             m[12] * m[7] * m[9];
+
+	    inv[12] = -m[4]  * m[9] * m[14] +
+	               m[4]  * m[10] * m[13] +
+	               m[8]  * m[5] * m[14] -
+	               m[8]  * m[6] * m[13] -
+	               m[12] * m[5] * m[10] +
+	               m[12] * m[6] * m[9];
+
+	    inv[1] = -m[1]  * m[10] * m[15] +
+	              m[1]  * m[11] * m[14] +
+	              m[9]  * m[2] * m[15] -
+	              m[9]  * m[3] * m[14] -
+	              m[13] * m[2] * m[11] +
+	              m[13] * m[3] * m[10];
+
+	    inv[5] = m[0]  * m[10] * m[15] -
+	             m[0]  * m[11] * m[14] -
+	             m[8]  * m[2] * m[15] +
+	             m[8]  * m[3] * m[14] +
+	             m[12] * m[2] * m[11] -
+	             m[12] * m[3] * m[10];
+
+	    inv[9] = -m[0]  * m[9] * m[15] +
+	              m[0]  * m[11] * m[13] +
+	              m[8]  * m[1] * m[15] -
+	              m[8]  * m[3] * m[13] -
+	              m[12] * m[1] * m[11] +
+	              m[12] * m[3] * m[9];
+
+	    inv[13] = m[0]  * m[9] * m[14] -
+	              m[0]  * m[10] * m[13] -
+	              m[8]  * m[1] * m[14] +
+	              m[8]  * m[2] * m[13] +
+	              m[12] * m[1] * m[10] -
+	              m[12] * m[2] * m[9];
+
+	    inv[2] = m[1]  * m[6] * m[15] -
+	             m[1]  * m[7] * m[14] -
+	             m[5]  * m[2] * m[15] +
+	             m[5]  * m[3] * m[14] +
+	             m[13] * m[2] * m[7] -
+	             m[13] * m[3] * m[6];
+
+	    inv[6] = -m[0]  * m[6] * m[15] +
+	              m[0]  * m[7] * m[14] +
+	              m[4]  * m[2] * m[15] -
+	              m[4]  * m[3] * m[14] -
+	              m[12] * m[2] * m[7] +
+	              m[12] * m[3] * m[6];
+
+	    inv[10] = m[0]  * m[5] * m[15] -
+	              m[0]  * m[7] * m[13] -
+	              m[4]  * m[1] * m[15] +
+	              m[4]  * m[3] * m[13] +
+	              m[12] * m[1] * m[7] -
+	              m[12] * m[3] * m[5];
+
+	    inv[14] = -m[0]  * m[5] * m[14] +
+	               m[0]  * m[6] * m[13] +
+	               m[4]  * m[1] * m[14] -
+	               m[4]  * m[2] * m[13] -
+	               m[12] * m[1] * m[6] +
+	               m[12] * m[2] * m[5];
+
+	    inv[3] = -m[1] * m[6] * m[11] +
+	              m[1] * m[7] * m[10] +
+	              m[5] * m[2] * m[11] -
+	              m[5] * m[3] * m[10] -
+	              m[9] * m[2] * m[7] +
+	              m[9] * m[3] * m[6];
+
+	    inv[7] = m[0] * m[6] * m[11] -
+	             m[0] * m[7] * m[10] -
+	             m[4] * m[2] * m[11] +
+	             m[4] * m[3] * m[10] +
+	             m[8] * m[2] * m[7] -
+	             m[8] * m[3] * m[6];
+
+	    inv[11] = -m[0] * m[5] * m[11] +
+	               m[0] * m[7] * m[9] +
+	               m[4] * m[1] * m[11] -
+	               m[4] * m[3] * m[9] -
+	               m[8] * m[1] * m[7] +
+	               m[8] * m[3] * m[5];
+
+	    inv[15] = m[0] * m[5] * m[10] -
+	              m[0] * m[6] * m[9] -
+	              m[4] * m[1] * m[10] +
+	              m[4] * m[2] * m[9] +
+	              m[8] * m[1] * m[6] -
+	              m[8] * m[2] * m[5];
+
+	    double det = m[0] * inv[0] + m[1] * inv[4] + m[2] * inv[8] + m[3] * inv[12];
+
+	    if (abs(det) < 1e-9) {
+	        return false;
+	    }
+
+
+	    det = 1.0 / det;
+
+	    for (int i = 0; i < 4; i++) {
+	    	for (int j = 0; j < 4; j++) {
+	    		inv_out[i][j] = inv[i * 4 + j] * det;
+	    	}
+	    }
+
+	    return true;
+	}
+
+  extern "C"
+	__global__ void best_local_affine_kernel(
+		float *output, float *input, float *affine_model,
+		int h, int w, float epsilon, int kernel_radius
+	)
+	{
+		int size = h * w;
+		int id = blockIdx.x * blockDim.x + threadIdx.x;
+
+		if (id < size) {
+			int x = id % w, y = id / w;
+
+			double Mt_M[4][4] = {}; // 4x4
+			double invMt_M[4][4] = {};
+			double Mt_S[3][4] = {}; // RGB -> 1x4
+			double A[3][4] = {};
+			for (int i = 0; i < 4; i++)
+				for (int j = 0; j < 4; j++) {
+					Mt_M[i][j] = 0, invMt_M[i][j] = 0;
+					if (i != 3) {
+						Mt_S[i][j] = 0, A[i][j] = 0;
+						if (i == j)
+				    		Mt_M[i][j] = 1e-3;
+				    }
+				}
+
+			for (int dy = -kernel_radius; dy <= kernel_radius; dy++) {
+				for (int dx = -kernel_radius; dx <= kernel_radius; dx++) {
+
+					int xx = x + dx, yy = y + dy;
+					int id2 = yy * w + xx;
+
+					if (0 <= xx && xx < w && 0 <= yy && yy < h) {
+
+						Mt_M[0][0] += input[id2 + 2*size] * input[id2 + 2*size];
+						Mt_M[0][1] += input[id2 + 2*size] * input[id2 + size];
+						Mt_M[0][2] += input[id2 + 2*size] * input[id2];
+						Mt_M[0][3] += input[id2 + 2*size];
+
+						Mt_M[1][0] += input[id2 + size] * input[id2 + 2*size];
+						Mt_M[1][1] += input[id2 + size] * input[id2 + size];
+						Mt_M[1][2] += input[id2 + size] * input[id2];
+						Mt_M[1][3] += input[id2 + size];
+
+						Mt_M[2][0] += input[id2] * input[id2 + 2*size];
+						Mt_M[2][1] += input[id2] * input[id2 + size];
+						Mt_M[2][2] += input[id2] * input[id2];
+						Mt_M[2][3] += input[id2];
+
+						Mt_M[3][0] += input[id2 + 2*size];
+						Mt_M[3][1] += input[id2 + size];
+						Mt_M[3][2] += input[id2];
+						Mt_M[3][3] += 1;
+
+						Mt_S[0][0] += input[id2 + 2*size] * output[id2 + 2*size];
+						Mt_S[0][1] += input[id2 + size] * output[id2 + 2*size];
+						Mt_S[0][2] += input[id2] * output[id2 + 2*size];
+						Mt_S[0][3] += output[id2 + 2*size];
+
+						Mt_S[1][0] += input[id2 + 2*size] * output[id2 + size];
+						Mt_S[1][1] += input[id2 + size] * output[id2 + size];
+						Mt_S[1][2] += input[id2] * output[id2 + size];
+						Mt_S[1][3] += output[id2 + size];
+
+						Mt_S[2][0] += input[id2 + 2*size] * output[id2];
+						Mt_S[2][1] += input[id2 + size] * output[id2];
+						Mt_S[2][2] += input[id2] * output[id2];
+						Mt_S[2][3] += output[id2];
+					}
+				}
+			}
+
+			bool success = InverseMat4x4(Mt_M, invMt_M);
+
+			for (int i = 0; i < 3; i++) {
+				for (int j = 0; j < 4; j++) {
+					for (int k = 0; k < 4; k++) {
+						A[i][j] += invMt_M[j][k] * Mt_S[i][k];
+					}
+				}
+			}
+
+			for (int i = 0; i < 3; i++) {
+				for (int j = 0; j < 4; j++) {
+					int affine_id = i * 4 + j;
+					affine_model[12 * id + affine_id] = A[i][j];
+				}
+			}
+		}
+		return ;
+	}
+
+  extern "C"
+	__global__ void bilateral_smooth_kernel(
+		float *affine_model, float *filtered_affine_model, float *guide,
+		int h, int w, int kernel_radius, float sigma1, float sigma2
+	)
+	{
+		int id = blockIdx.x * blockDim.x + threadIdx.x;
+		int size = h * w;
+		if (id < size) {
+			int x = id % w;
+			int y = id / w;
+
+			double sum_affine[12] = {};
+			double sum_weight = 0;
+			for (int dx = -kernel_radius; dx <= kernel_radius; dx++) {
+				for (int dy = -kernel_radius; dy <= kernel_radius; dy++) {
+					int yy = y + dy, xx = x + dx;
+					int id2 = yy * w + xx;
+					if (0 <= xx && xx < w && 0 <= yy && yy < h) {
+						float color_diff1 = guide[yy*w + xx] - guide[y*w + x];
+						float color_diff2 = guide[yy*w + xx + size] - guide[y*w + x + size];
+						float color_diff3 = guide[yy*w + xx + 2*size] - guide[y*w + x + 2*size];
+						float color_diff_sqr =
+							(color_diff1*color_diff1 + color_diff2*color_diff2 + color_diff3*color_diff3) / 3;
+
+						float v1 = exp(-(dx * dx + dy * dy) / (2 * sigma1 * sigma1));
+						float v2 = exp(-(color_diff_sqr) / (2 * sigma2 * sigma2));
+						float weight = v1 * v2;
+
+						for (int i = 0; i < 3; i++) {
+							for (int j = 0; j < 4; j++) {
+								int affine_id = i * 4 + j;
+								sum_affine[affine_id] += weight * affine_model[id2*12 + affine_id];
+							}
+						}
+						sum_weight += weight;
+					}
+				}
+			}
+
+			for (int i = 0; i < 3; i++) {
+				for (int j = 0; j < 4; j++) {
+					int affine_id = i * 4 + j;
+					filtered_affine_model[id*12 + affine_id] = sum_affine[affine_id] / sum_weight;
+				}
+			}
+		}
+		return ;
+	}
+
+
+  extern "C"
+	__global__ void reconstruction_best_kernel(
+		float *input, float *filtered_affine_model, float *filtered_best_output,
+		int h, int w
+	)
+	{
+		int id = blockIdx.x * blockDim.x + threadIdx.x;
+		int size = h * w;
+		if (id < size) {
+			double out1 =
+				input[id + 2*size] * filtered_affine_model[id*12 + 0] + // A[0][0] +
+				input[id + size]   * filtered_affine_model[id*12 + 1] + // A[0][1] +
+				input[id]          * filtered_affine_model[id*12 + 2] + // A[0][2] +
+									 filtered_affine_model[id*12 + 3]; //A[0][3];
+			double out2 =
+				input[id + 2*size] * filtered_affine_model[id*12 + 4] + //A[1][0] +
+				input[id + size]   * filtered_affine_model[id*12 + 5] + //A[1][1] +
+				input[id]          * filtered_affine_model[id*12 + 6] + //A[1][2] +
+									 filtered_affine_model[id*12 + 7]; //A[1][3];
+			double out3 =
+				input[id + 2*size] * filtered_affine_model[id*12 + 8] + //A[2][0] +
+				input[id + size]   * filtered_affine_model[id*12 + 9] + //A[2][1] +
+				input[id]          * filtered_affine_model[id*12 + 10] + //A[2][2] +
+									 filtered_affine_model[id*12 + 11]; // A[2][3];
+
+			filtered_best_output[id] = out1;
+			filtered_best_output[id + size] = out2;
+			filtered_best_output[id + 2*size] = out3;
+		}
+		return ;
+	}
+	'''
+
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+from cupy.cuda import function
+from pynvrtc.compiler import Program
+from collections import namedtuple
+
+
+def smooth_local_affine(output_cpu, input_cpu, epsilon, patch, h, w, f_r, f_e):
+    # program = Program(src.encode('utf-8'), 'best_local_affine_kernel.cu'.encode('utf-8'))
+    # ptx = program.compile(['-I/usr/local/cuda/include'.encode('utf-8')])
+    program = Program(src, 'best_local_affine_kernel.cu')
+    ptx = program.compile(['-I/usr/local/cuda/include'])
+    m = function.Module()
+    m.load(bytes(ptx.encode()))
+
+    _reconstruction_best_kernel = m.get_function('reconstruction_best_kernel')
+    _bilateral_smooth_kernel = m.get_function('bilateral_smooth_kernel')
+    _best_local_affine_kernel = m.get_function('best_local_affine_kernel')
+    Stream = namedtuple('Stream', ['ptr'])
+    s = Stream(ptr=torch.cuda.current_stream().cuda_stream)
+
+    filter_radius = f_r
+    sigma1 = filter_radius / 3
+    sigma2 = f_e
+    radius = (patch - 1) / 2
+
+    filtered_best_output = torch.zeros(np.shape(input_cpu)).cuda()
+    affine_model =    torch.zeros((h * w, 12)).cuda()
+    filtered_affine_model =torch.zeros((h * w, 12)).cuda()
+
+    input_ = torch.from_numpy(input_cpu).cuda()
+    output_ = torch.from_numpy(output_cpu).cuda()
+    _best_local_affine_kernel(
+        grid=(int((h * w) / 256 + 1), 1),
+        block=(256, 1, 1),
+        args=[output_.data_ptr(), input_.data_ptr(), affine_model.data_ptr(),
+             np.int32(h), np.int32(w), np.float32(epsilon), np.int32(radius)], stream=s
+     )
+
+    _bilateral_smooth_kernel(
+        grid=(int((h * w) / 256 + 1), 1),
+        block=(256, 1, 1),
+        args=[affine_model.data_ptr(), filtered_affine_model.data_ptr(), input_.data_ptr(), np.int32(h), np.int32(w), np.int32(f_r), np.float32(sigma1), np.float32(sigma2)], stream=s
+    )
+
+    _reconstruction_best_kernel(
+        grid=(int((h * w) / 256 + 1), 1),
+        block=(256, 1, 1),
+        args=[input_.data_ptr(), filtered_affine_model.data_ptr(), filtered_best_output.data_ptr(),
+        np.int32(h), np.int32(w)], stream=s
+    )
+    numpy_filtered_best_output = filtered_best_output.cpu().numpy()
+    return numpy_filtered_best_output
+
+
+def smooth_filter(initImg, contentImg, f_radius=15,f_edge=1e-1):
+	'''
+	:param initImg: intermediate output. Either image path or PIL Image
+	:param contentImg: content image output. Either path or PIL Image
+	:return: stylized output image. PIL Image
+	'''
+	if type(initImg) == str:
+	    initImg = Image.open(initImg).convert("RGB")
+	best_image_bgr = np.array(initImg, dtype=np.float32)
+	bW, bH, bC = best_image_bgr.shape
+	best_image_bgr = best_image_bgr[:, :, ::-1]
+	best_image_bgr = best_image_bgr.transpose((2, 0, 1))
+
+	if type(contentImg) == str:
+		contentImg = Image.open(contentImg).convert("RGB")
+		content_input = contentImg.resize((bH,bW))
+	else:
+		content_input = cv2.resize(contentImg,(bH,bW))
+	content_input = np.array(content_input, dtype=np.float32)
+	content_input = content_input[:, :, ::-1]
+	content_input = content_input.transpose((2, 0, 1))
+	input_ = np.ascontiguousarray(content_input, dtype=np.float32) / 255.
+	_, H, W = np.shape(input_)
+	output_ = np.ascontiguousarray(best_image_bgr, dtype=np.float32) / 255.
+	best_ = smooth_local_affine(output_, input_, 1e-7, 3, H, W, f_radius, f_edge)
+	best_ = best_.transpose(1, 2, 0)
+	result = Image.fromarray(np.uint8(np.clip(best_ * 255., 0, 255.)))
+	return result

graph_networks/LinearStyleTransfer/libs/utils.py

@@ -0,0 +1,92 @@
+from __future__ import division
+import os
+import cv2
+import time
+import torch
+import scipy.misc
+import numpy as np
+import scipy.sparse
+from PIL import Image
+import scipy.sparse.linalg
+from cv2.ximgproc import jointBilateralFilter
+from numpy.lib.stride_tricks import as_strided
+
+def whiten(cF):
+    cFSize = cF.size()
+    c_mean = torch.mean(cF,1) # c x (h x w)
+    c_mean = c_mean.unsqueeze(1).expand_as(cF)
+    cF = cF - c_mean
+
+    contentConv = torch.mm(cF,cF.t()).div(cFSize[1]-1) + torch.eye(cFSize[0]).double()
+    c_u,c_e,c_v = torch.svd(contentConv,some=False)
+
+    k_c = cFSize[0]
+    for i in range(cFSize[0]):
+        if c_e[i] < 0.00001:
+            k_c = i
+            break
+
+    c_d = (c_e[0:k_c]).pow(-0.5)
+    step1 = torch.mm(c_v[:,0:k_c],torch.diag(c_d))
+    step2 = torch.mm(step1,(c_v[:,0:k_c].t()))
+    whiten_cF = torch.mm(step2,cF)
+    return whiten_cF
+
+def numpy2cv2(cont,style,prop,width,height):
+    cont = cont.transpose((1,2,0))
+    cont = cont[...,::-1]
+    cont = cont * 255
+    cont = cv2.resize(cont,(width,height))
+    #cv2.resize(iimg,(width,height))
+    style = style.transpose((1,2,0))
+    style = style[...,::-1]
+    style = style * 255
+    style = cv2.resize(style,(width,height))
+
+    prop = prop.transpose((1,2,0))
+    prop = prop[...,::-1]
+    prop = prop * 255
+    prop = cv2.resize(prop,(width,height))
+
+    #return np.concatenate((cont,np.concatenate((style,prop),axis=1)),axis=1)
+    return prop,cont
+
+def makeVideo(content,style,props,outf):
+    print('Stack transferred frames back to video...')
+    layers,height,width = content[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*'MJPG')
+    video = cv2.VideoWriter(os.path.join(outf,'transfer.avi'),fourcc,10.0,(width,height))
+    ori_video = cv2.VideoWriter(os.path.join(outf,'content.avi'),fourcc,10.0,(width,height))
+    for j in range(len(content)):
+        prop,cont = numpy2cv2(content[j],style,props[j],width,height)
+        cv2.imwrite('prop.png',prop)
+        cv2.imwrite('content.png',cont)
+        # TODO: this is ugly, fix this
+        imgj = cv2.imread('prop.png')
+        imgc = cv2.imread('content.png')
+
+        video.write(imgj)
+        ori_video.write(imgc)
+        # RGB or BRG, yuks
+    video.release()
+    ori_video.release()
+    os.remove('prop.png')
+    os.remove('content.png')
+    print('Transferred video saved at %s.'%outf)
+
+def print_options(opt):
+    message = ''
+    message += '----------------- Options ---------------\n'
+    for k, v in sorted(vars(opt).items()):
+        comment = ''
+        message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
+    message += '----------------- End -------------------'
+    print(message)
+
+    # save to the disk
+    expr_dir = os.path.join(opt.outf)
+    os.makedirs(expr_dir,exist_ok=True)
+    file_name = os.path.join(expr_dir, 'opt.txt')
+    with open(file_name, 'wt') as opt_file:
+        opt_file.write(message)
+        opt_file.write('\n')

graph_networks/LinearStyleTransfer/models/dec_r31.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ccc3bbc97a15e1d002d0b13523543c518dda2d0346c8f4d39c1d381d8490f68
+size 2221888

graph_networks/LinearStyleTransfer/models/dec_r41.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6858f96d3d0882fa3b40652a0315928219086e4bcb0e3efbe43bd04ea631911
+size 14023509

graph_networks/LinearStyleTransfer/models/r31.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8bb75be684b331a105a8fad266343556067e4eec888249f7abb693a66b5ad7e3
+size 11564438