Upload LDR to HDR application files

app.py

@@ -0,0 +1,339 @@
+
+
+
+# import gradio as gr
+# import torch
+# import os
+# import tempfile
+# import shutil
+# from PIL import Image
+# import numpy as np
+# from pathlib import Path
+# import sys
+# import copy
+
+# # --- Import logic from your project ---
+# from options.test_options import TestOptions
+# from data import create_dataset
+# from models import create_model
+# try:
+#     from best_ldr import compute_metrics_for_images, score_records
+# except ImportError:
+#     raise ImportError("Could not import from best_ldr.py. Make sure the file is in the same directory as app.py.")
+
+# print("--- Initializing LDR-to-HDR Model (this may take a moment) ---")
+
+# # --- Global Setup: Load the CycleGAN model once when the app starts ---
+
+# # We need to satisfy the parser's requirement for a dataroot at startup
+# if '--dataroot' not in sys.argv:
+#     sys.argv.extend(['--dataroot', './dummy_dataroot_for_init'])
+
+# # Load the base options
+# opt = TestOptions().parse()
+
+# # Manually override settings for our model
+# opt.name = 'ldr2hdr_cyclegan_728'
+# opt.model = 'test'
+# opt.netG = 'resnet_9blocks'
+# opt.norm = 'instance'
+# opt.no_dropout = True
+# opt.checkpoints_dir = './checkpoints'
+# opt.gpu_ids = [0] if torch.cuda.is_available() else []
+# opt.device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu')
+
+# # Create the model using these options
+# model = create_model(opt)
+# model.setup(opt)
+# model.eval()
+
+# print("--- Model Loaded Successfully ---")
+
+
+# # --- Helper Function for Inference ---
+
+# def run_inference(model, image_path, process_options):
+#     """
+#     A reusable function to run the model with specific preprocessing options.
+#     """
+#     # Deep copy the base options to avoid modifying the global state
+#     local_opt = copy.deepcopy(opt)
+
+#     # Apply the specific settings for this run
+#     for key, value in process_options.items():
+#         setattr(local_opt, key, value)
+
+#     with tempfile.TemporaryDirectory() as temp_dir:
+#         shutil.copy(image_path, temp_dir)
+#         local_opt.dataroot = temp_dir
+#         local_opt.num_test = 1
+#         dataset = create_dataset(local_opt)
+
+#         for i, data in enumerate(dataset):
+#             model.set_input(data)
+#             model.test()
+#             visuals = model.get_current_visuals()
+            
+#             for label, image_tensor in visuals.items():
+#                 if label == 'fake':
+#                     image_numpy = (np.transpose(image_tensor.cpu().float().numpy()[0], (1, 2, 0)) + 1) / 2.0 * 255.0
+#                     return Image.fromarray(image_numpy.astype(np.uint8))
+
+# # --- The Main Gradio Processing Function ---
+
+# def process_images_and_display(list_of_temp_files):
+#     """
+#     The main workflow: select best LDR, then run two inference modes.
+#     """
+#     if not list_of_temp_files:
+#         raise gr.Error("Please upload your bracketed LDR images.")
+#     if len(list_of_temp_files) < 2:
+#         gr.Warning("For best results, upload at least 2 bracketed LDR images.")
+
+#     uploaded_filepaths = [Path(f.name) for f in list_of_temp_files]
+    
+#     try:
+#         # --- Step 1: Select the Best LDR ---
+#         print(f"Analyzing {len(uploaded_filepaths)} uploaded images...")
+#         weights = {"clipped": 0.35, "coverage": 0.25, "exposure": 0.15, "sharpness": 0.15, "noise": 0.10}
+#         records = compute_metrics_for_images(uploaded_filepaths, resize_max=1024)
+#         scored_records = score_records(records, weights)
+#         if not scored_records:
+#             raise gr.Error("Could not read or score any of the uploaded images.")
+            
+#         best_ldr_record = scored_records[0]
+#         best_ldr_path = best_ldr_record['path']
+#         print(f"Best LDR selected: {os.path.basename(best_ldr_path)} (Score: {best_ldr_record['score']:.4f})")
+#         chosen_ldr_image = Image.open(best_ldr_path).convert("RGB")
+
+#         # --- Step 2: Run Inference in Both Modes ---
+        
+#         # Mode A: High-Quality Crop (at model's native resolution)
+#         print("Running Mode A: High-Quality Crop...")
+#         crop_options = {
+#             'preprocess': 'resize_and_crop',
+#             'load_size': 728,
+#             'crop_size': 728
+#         }
+#         hdr_cropped = run_inference(model, best_ldr_path, crop_options)
+#         print("Mode A successful.")
+        
+#         # Mode B: Full Image (at a higher resolution)
+#         print("Running Mode B: Full Image (High-Res Scaled)...")
+#         scale_options = {
+#             'preprocess': 'scale_width',
+#             'load_size': 1024, # <-- THIS IS THE CHANGE FOR HIGHER RESOLUTION
+#             'crop_size': 728 # This value is ignored by scale_width but needs to be present
+#         }
+#         hdr_scaled = run_inference(model, best_ldr_path, scale_options)
+#         print("Mode B successful.")
+        
+#         # Return all the images to update the UI
+#         return uploaded_filepaths, chosen_ldr_image, hdr_cropped, hdr_scaled
+
+#     except Exception as e:
+#         print(f"An error occurred: {e}")
+#         raise gr.Error(f"An error occurred during processing: {e}")
+
+# # --- Create and Launch the Gradio Interface ---
+
+# with gr.Blocks(theme=gr.themes.Monochrome(), css="footer {display: none !important}") as demo:
+#     gr.Markdown("# LDR Bracketing to HDR Converter")
+#     gr.Markdown("Upload a set of bracketed LDR images. The app will automatically select the best one and convert it to HDR using two different methods for comparison.")
+    
+#     with gr.Row():
+#         with gr.Column(scale=1, min_width=300):
+#             input_files = gr.Files(
+#                 label="Upload Bracketed LDR Images",
+#                 file_types=["image"]
+#             )
+#             process_button = gr.Button("Process Images", variant="primary")
+            
+#             with gr.Accordion("See Your Uploads", open=False):
+#                  input_gallery = gr.Gallery(label="Uploaded LDR Bracket", show_label=False, columns=3, height="auto")
+
+#         with gr.Column(scale=2):
+#             gr.Markdown("## Results")
+#             with gr.Row():
+#                 chosen_ldr_display = gr.Image(label="Best LDR Chosen by Algorithm", type="pil", interactive=False)
+#             with gr.Row():
+#                 output_cropped = gr.Image(label="Result 1: High-Quality Crop (728x728)", type="pil", interactive=False)
+#                 output_scaled = gr.Image(label="Result 2: Full Image (Scaled to 1024px Width)", type="pil", interactive=False)
+
+#     process_button.click(
+#         fn=process_images_and_display,
+#         inputs=input_files,
+#         outputs=[input_gallery, chosen_ldr_display, output_cropped, output_scaled]
+#     )
+
+# print("--- Launching Gradio App ---")
+# demo.launch(share=True)
+
+
+
+
+
+
+
+
+
+import gradio as gr
+import torch
+import os
+import tempfile
+import shutil
+from PIL import Image
+import numpy as np
+from pathlib import Path
+import sys
+import copy
+
+# --- Import logic from your project ---
+from options.test_options import TestOptions
+from data import create_dataset
+from models import create_model
+try:
+    from best_ldr import compute_metrics_for_images, score_records
+except ImportError:
+    raise ImportError("Could not import from best_ldr.py. Make sure the file is in the same directory as app.py.")
+
+print("--- Initializing LDR-to-HDR Model (this may take a moment) ---")
+
+# --- Global Setup: Load the CycleGAN model once when the app starts ---
+
+# We need to satisfy the parser's requirement for a dataroot at startup
+if '--dataroot' not in sys.argv:
+    sys.argv.extend(['--dataroot', './dummy_dataroot_for_init'])
+
+# Load the base options
+opt = TestOptions().parse()
+
+# Manually override settings for our model
+opt.name = 'ldr2hdr_cyclegan_728'
+opt.model = 'test'
+opt.netG = 'resnet_9blocks'
+opt.norm = 'instance'
+opt.no_dropout = True
+opt.checkpoints_dir = './checkpoints'
+opt.gpu_ids = [0] if torch.cuda.is_available() else []
+opt.device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu')
+
+# Create the model using these options
+model = create_model(opt)
+model.setup(opt)
+model.eval()
+
+print("--- Model Loaded Successfully ---")
+
+
+# --- The Main Gradio Processing Function ---
+
+def process_images_to_hdr(list_of_temp_files):
+    """
+    The main workflow: select best LDR, run inference, and return results for the UI.
+    """
+    if not list_of_temp_files:
+        raise gr.Error("Please upload your bracketed LDR images.")
+    if len(list_of_temp_files) < 2:
+        gr.Warning("For best results, upload at least 2 bracketed LDR images.")
+
+    uploaded_filepaths = [Path(f.name) for f in list_of_temp_files]
+    
+    try:
+        # --- Step 1: Select the Best LDR ---
+        print(f"Analyzing {len(uploaded_filepaths)} uploaded images...")
+        weights = {"clipped": 0.35, "coverage": 0.25, "exposure": 0.15, "sharpness": 0.15, "noise": 0.10}
+        records = compute_metrics_for_images(uploaded_filepaths, resize_max=1024)
+        scored_records = score_records(records, weights)
+        if not scored_records:
+            raise gr.Error("Could not read or score any of the uploaded images.")
+            
+        best_ldr_record = scored_records[0]
+        best_ldr_path = best_ldr_record['path']
+        print(f"Best LDR selected: {os.path.basename(best_ldr_path)} (Score: {best_ldr_record['score']:.4f})")
+
+        # --- Step 2: Run Inference ---
+        print("Running Full Image (High-Res Scaled) Inference...")
+        
+        # We only need the one set of options now
+        inference_options = {
+            'preprocess': 'scale_width',
+            'load_size': 1024, # Generate the high-resolution, full image
+            'crop_size': 728   # This value is ignored but required by the parser
+        }
+        
+        # Deep copy the base options to avoid modifying the global state
+        local_opt = copy.deepcopy(opt)
+        for key, value in inference_options.items():
+            setattr(local_opt, key, value)
+
+        # Run the model
+        with tempfile.TemporaryDirectory() as temp_dir:
+            shutil.copy(best_ldr_path, temp_dir)
+            local_opt.dataroot = temp_dir
+            local_opt.num_test = 1
+            dataset = create_dataset(local_opt)
+
+            for i, data in enumerate(dataset):
+                model.set_input(data)
+                model.test()
+                visuals = model.get_current_visuals()
+                
+                for label, image_tensor in visuals.items():
+                    if label == 'fake':
+                        image_numpy = (np.transpose(image_tensor.cpu().float().numpy()[0], (1, 2, 0)) + 1) / 2.0 * 255.0
+                        final_hdr_image = Image.fromarray(image_numpy.astype(np.uint8))
+                        print("Conversion to HDR successful.")
+                        # Return the gallery of inputs and the single final HDR image
+                        return uploaded_filepaths, final_hdr_image
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        raise gr.Error(f"An error occurred during processing: {e}")
+
+# --- Create and Launch the Gradio Interface ---
+
+with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}") as demo:
+    gr.Markdown(
+        """
+        # LDR Bracketing to HDR Converter
+        Upload a set of bracketed LDR images. The app will automatically select the best one and convert it to a vibrant, full-resolution HDR image.
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=1, min_width=350):
+            # --- INPUT ---
+            input_files = gr.Files(
+                label="Upload Bracketed LDR Images",
+                file_types=["image"]
+            )
+            process_button = gr.Button("Process Images", variant="primary")
+            with gr.Accordion("See Your Uploaded Images", open=True):
+                 input_gallery = gr.Gallery(label="Uploaded Images", show_label=False, columns=[2, 3], height="auto")
+
+        with gr.Column(scale=2):
+            # --- OUTPUT ---
+            gr.Markdown("## Generated HDR Result")
+            output_image = gr.Image(label="Final HDR Image", type="pil", interactive=False, show_download_button=True)
+
+    process_button.click(
+        fn=process_images_to_hdr,
+        inputs=input_files,
+        outputs=[input_gallery, output_image]
+    )
+    
+    # gr.Markdown("### Examples")
+    # gr.Examples(
+    #     examples=[
+    #         [
+    #             "../pix2pix_dataset/testA/077A2406.jpg", 
+    #             "../pix2pix_dataset/testA/077A4049.jpg", 
+    #             "../pix2pix_dataset/testA/077A4073.jpg"
+    #         ]
+    #     ],
+    #     inputs=input_files
+    # )
+
+print("--- Launching Gradio App ---")
+demo.launch(share=True)

best_ldr.py

@@ -0,0 +1,327 @@
+#!/usr/bin/env python3
+"""
+make_cyclegan_dataset.py
+
+Create paired datasets (setA, setB) for CycleGAN training from your dataset.
+
+What it does:
+ - Walks the dataset root (e.g. ../jpeg_stage1Just0)
+ - Finds scene directories that contain both a `source/` subfolder with >= min_images
+   and an `output/` subfolder with at least one image.
+ - For each scene: selects the best LDR from `source/` (using metrics: clipped, coverage,
+   exposure centering, sharpness, noise), copies that chosen source image into outdir/setA/,
+   copies the scene's output image into outdir/setB/ but renames it to the chosen source filename.
+ - Writes CSV and JSON reports with metric breakdowns.
+
+Usage:
+  python make_cyclegan_dataset.py --root ../jpeg_stage1Just0 --outdir ./cyclegan_data
+
+Dependencies:
+  pip install opencv-python pillow numpy
+
+Author: ChatGPT (opinionated: default weights favor low clipping and good coverage)
+"""
+
+import argparse
+import os
+from pathlib import Path
+import json
+import csv
+import shutil
+from math import fabs
+
+import numpy as np
+import cv2
+from PIL import Image, ExifTags
+
+IMG_EXTS = {".jpg", ".jpeg", ".png", ".tif", ".tiff", ".bmp", ".webp"}
+
+# ------------------ Image / metric helpers ------------------
+
+def is_image_file(p: Path):
+    return p.suffix.lower() in IMG_EXTS and p.is_file()
+
+def list_images(folder: Path):
+    if not folder.exists():
+        return []
+    return sorted([p for p in folder.iterdir() if is_image_file(p)])
+
+def read_image_gray(path: Path, resize_max=None):
+    """Read color then convert to grayscale float32 [0,1]. Uses cv2.imdecode to handle weird filenames."""
+    arr = np.fromfile(str(path), dtype=np.uint8)
+    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+    if img is None:
+        raise IOError(f"Failed to read image {path}")
+    if resize_max:
+        h, w = img.shape[:2]
+        scale = resize_max / max(h, w) if max(h, w) > resize_max else 1.0
+        if scale != 1.0:
+            img = cv2.resize(img, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_AREA)
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
+    return gray
+
+def clipped_ratio(gray):
+    total = gray.size
+    high = np.count_nonzero(gray >= 0.992)
+    low = np.count_nonzero(gray <= 0.008)
+    return float(high + low) / float(total)
+
+def histogram_coverage(gray, bins=256, min_frac=0.001):
+    hist, _ = np.histogram((gray * 255).astype(np.uint8), bins=bins, range=(0,255))
+    threshold = max(1, int(min_frac * gray.size))
+    covered = np.count_nonzero(hist >= threshold)
+    return float(covered) / float(bins)
+
+def exposure_distance(gray):
+    return float(abs(float(np.mean(gray)) - 0.5))
+
+def sharpness_metric(gray):
+    lap = cv2.Laplacian((gray * 255).astype(np.uint8), cv2.CV_64F)
+    return float(np.var(lap))
+
+def noise_estimate(gray):
+    blur = cv2.GaussianBlur(gray, (3,3), 0)
+    hf = gray - blur
+    return float(np.std(hf))
+
+def minmax_normalize(vals, eps=1e-8):
+    arr = np.array(vals, dtype=np.float64)
+    mn = float(arr.min())
+    mx = float(arr.max())
+    if mx - mn < eps:
+        # all equal -> zeros
+        return np.zeros_like(arr)
+    return (arr - mn) / (mx - mn)
+
+# ------------------ Selection & scene processing ------------------
+
+def compute_metrics_for_images(image_paths, resize_max):
+    records = []
+    for p in image_paths:
+        try:
+            g = read_image_gray(p, resize_max=resize_max)
+        except Exception as e:
+            print(f"  WARNING: cannot read {p}: {e}")
+            continue
+        rec = {
+            "path": str(p),
+            "name": p.name,
+            "clipped": clipped_ratio(g),
+            "coverage": histogram_coverage(g),
+            "exposure_dist": exposure_distance(g),
+            "sharpness": sharpness_metric(g),
+            "noise": noise_estimate(g)
+        }
+        records.append(rec)
+    return records
+
+def score_records(records, weights):
+    if not records:
+        return []
+    clipped_vals = [r["clipped"] for r in records]
+    cov_vals = [r["coverage"] for r in records]
+    exp_vals = [r["exposure_dist"] for r in records]
+    sharp_vals = [r["sharpness"] for r in records]
+    noise_vals = [r["noise"] for r in records]
+
+    clipped_n = minmax_normalize(clipped_vals)
+    cov_n = minmax_normalize(cov_vals)
+    exp_n = minmax_normalize(exp_vals)
+    sharp_n = minmax_normalize(sharp_vals)
+    noise_n = minmax_normalize(noise_vals)
+
+    scored = []
+    for i, r in enumerate(records):
+        score = 0.0
+        score += weights["clipped"] * (1.0 - float(clipped_n[i]))   # less clipping -> better
+        score += weights["coverage"] * float(cov_n[i])              # more coverage -> better
+        score += weights["exposure"] * (1.0 - float(exp_n[i]))      # closer to mid gray -> better
+        score += weights["sharpness"] * float(sharp_n[i])           # sharper -> better
+        score += weights["noise"] * (1.0 - float(noise_n[i]))       # less noise -> better
+
+        rec = dict(r)
+        rec.update({
+            "clipped_n": float(clipped_n[i]),
+            "coverage_n": float(cov_n[i]),
+            "exposure_n": float(exp_n[i]),
+            "sharpness_n": float(sharp_n[i]),
+            "noise_n": float(noise_n[i]),
+            "score": float(score)
+        })
+        scored.append(rec)
+    scored_sorted = sorted(scored, key=lambda x: x["score"], reverse=True)
+    return scored_sorted
+
+def find_output_image(output_folder: Path):
+    imgs = list_images(output_folder)
+    if not imgs:
+        return None
+    # Prefer file with same name as parent folder (if present), else pick largest file
+    parent_name = output_folder.parent.name
+    for p in imgs:
+        if p.stem == parent_name:
+            return p
+    # otherwise pick largest by file size (likely the main image)
+    imgs_sorted = sorted(imgs, key=lambda x: x.stat().st_size, reverse=True)
+    return imgs_sorted[0]
+
+# ------------------ Main procedure ------------------
+
+def make_dataset(root: Path, outdir: Path, min_images: int,
+                 resize_max: int, weights: dict, copy_method="copy"):
+    scenes_found = 0
+    results = []
+    setA = outdir / "setA"
+    setB = outdir / "setB"
+    os.makedirs(setA, exist_ok=True)
+    os.makedirs(setB, exist_ok=True)
+
+    # Walk the tree and find directories that contain both source/ and output/
+    for dirpath, dirnames, filenames in os.walk(root):
+        d = Path(dirpath)
+        src_dir = d / "source"
+        out_dir = d / "output"
+        if not src_dir.exists() or not out_dir.exists():
+            continue
+        src_imgs = list_images(src_dir)
+        if len(src_imgs) < min_images:
+            # skip small scenes
+            continue
+
+        scenes_found += 1
+        print(f"[{scenes_found}] Scene: {d}  ({len(src_imgs)} source images)")
+
+        # compute metrics and choose best
+        records = compute_metrics_for_images(src_imgs, resize_max=resize_max)
+        if not records:
+            print("  No readable source images, skipping.")
+            continue
+        scored = score_records(records, weights)
+        chosen = scored[0]
+        chosen_path = Path(chosen["path"])
+        chosen_name = chosen_path.name  # used for setA filename (and setB target name)
+
+        # find output image for this scene
+        out_img = find_output_image(out_dir)
+        if out_img is None:
+            print(f"  WARNING: no output image found in {out_dir}; skipping copying pair.")
+            out_img_path = None
+        else:
+            out_img_path = out_img
+
+        # destination paths
+        destA = setA / chosen_name
+        destB = setB / chosen_name
+
+        # copy or symlink
+        try:
+            if copy_method == "symlink":
+                if destA.exists():
+                    destA.unlink()
+                os.symlink(os.path.abspath(chosen_path), destA)
+            else:
+                shutil.copy2(chosen_path, destA)
+        except Exception as e:
+            print(f"  ERROR copying source -> {destA}: {e}")
+
+        if out_img_path is not None:
+            try:
+                if copy_method == "symlink":
+                    if destB.exists():
+                        destB.unlink()
+                    os.symlink(os.path.abspath(out_img_path), destB)
+                else:
+                    shutil.copy2(out_img_path, destB)
+            except Exception as e:
+                print(f"  ERROR copying output -> {destB}: {e}")
+
+        # record result
+        result = {
+            "scene_dir": str(d),
+            "source_dir": str(src_dir),
+            "output_dir": str(out_dir),
+            "chosen_source_path": str(chosen_path),
+            "chosen_source_name": chosen_name,
+            "chosen_score": chosen["score"],
+            "metrics": {
+                "clipped": chosen["clipped"],
+                "coverage": chosen["coverage"],
+                "exposure_dist": chosen["exposure_dist"],
+                "sharpness": chosen["sharpness"],
+                "noise": chosen["noise"],
+                "clipped_n": chosen["clipped_n"],
+                "coverage_n": chosen["coverage_n"],
+                "exposure_n": chosen["exposure_n"],
+                "sharpness_n": chosen["sharpness_n"],
+                "noise_n": chosen["noise_n"],
+            },
+            "output_image_used": str(out_img_path) if out_img_path is not None else None,
+            "destA": str(destA),
+            "destB": str(destB) if out_img_path is not None else None
+        }
+        results.append(result)
+
+        # print top 3 for quick audit
+        print("  Top candidates:")
+        for c in scored[:3]:
+            print(f"    {c['score']:.4f}  clipped={c['clipped']:.4f} cov={c['coverage']:.4f} expd={c['exposure_dist']:.4f} sharp={c['sharpness']:.1f} noise={c['noise']:.5f}  -> {Path(c['path']).name}")
+
+    # write reports
+    out_csv = outdir / "paired_selection.csv"
+    out_json = outdir / "paired_selection.json"
+    with open(out_json, "w", encoding="utf-8") as jf:
+        json.dump(results, jf, indent=2)
+    with open(out_csv, "w", newline="", encoding="utf-8") as cf:
+        writer = csv.writer(cf)
+        header = ["scene_dir", "source_dir", "output_dir", "chosen_source_name", "chosen_source_path",
+                  "chosen_score", "output_image_used", "destA", "destB"]
+        writer.writerow(header)
+        for r in results:
+            writer.writerow([r.get(h, "") for h in header])
+
+    print(f"\nDone. Scenes processed: {scenes_found}")
+    print(f"Paired data saved to:\n  {setA}\n  {setB}")
+    print(f"Reports: {out_csv} , {out_json}")
+    return results
+
+# ------------------ CLI ------------------
+
+def parse_weights(s):
+    parts = [float(x.strip()) for x in s.split(",")]
+    if len(parts) != 5:
+        raise argparse.ArgumentTypeError("weights must be 5 comma-separated numbers")
+    ssum = sum(parts)
+    if ssum == 0:
+        raise argparse.ArgumentTypeError("weights sum must be > 0")
+    return [p / ssum for p in parts]
+
+def main():
+    ap = argparse.ArgumentParser(description="Make paired CycleGAN dataset from your LDR/HDR scene layout.")
+    ap.add_argument("--root", "-r", required=True, help="Root of dataset (e.g. ../jpeg_stage1Just0)")
+    ap.add_argument("--outdir", "-o", default="./cyclegan_data", help="Output folder for paired dataset")
+    ap.add_argument("--min_images", type=int, default=2, help="Minimum images in source/ to consider scene")
+    ap.add_argument("--resize_max", type=int, default=1024, help="Resize longest side for metric calc (speeds up)")
+    ap.add_argument("--weights", type=parse_weights, default="0.35,0.25,0.15,0.15,0.10",
+                    help="5 weights: clipped,coverage,exposure,sharpness,noise  (will be normalized)")
+    ap.add_argument("--copy_method", choices=["copy", "symlink"], default="copy",
+                    help="copy files or create symlinks (symlink saves disk space)")
+    args = ap.parse_args()
+
+    root = Path(args.root).expanduser().resolve()
+    outdir = Path(args.outdir).expanduser().resolve()
+    w = args.weights if isinstance(args.weights, list) else args.weights  # parse_weights returns list
+    weights = {
+        "clipped": w[0],
+        "coverage": w[1],
+        "exposure": w[2],
+        "sharpness": w[3],
+        "noise": w[4]
+    }
+    print("Using weights:", weights)
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    make_dataset(root, outdir, min_images=args.min_images,
+                 resize_max=args.resize_max, weights=weights, copy_method=args.copy_method)
+
+if __name__ == "__main__":
+    main()

checkpoints/ldr2hdr_cyclegan_728/latest_net_G.pth

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

data/__init__.py

@@ -0,0 +1,107 @@
+"""This package includes all the modules related to data loading and preprocessing
+
+ To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.
+ You need to implement four functions:
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point from data loader.
+    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
+
+Now you can use the dataset class by specifying flag '--dataset_mode dummy'.
+See our template dataset class 'template_dataset.py' for more details.
+"""
+
+import importlib
+import torch.utils.data
+from torch.utils.data.distributed import DistributedSampler
+import torch.distributed as dist
+import os
+from data.base_dataset import BaseDataset
+
+
+def find_dataset_using_name(dataset_name):
+    """Import the module "data/[dataset_name]_dataset.py".
+
+    In the file, the class called DatasetNameDataset() will
+    be instantiated. It has to be a subclass of BaseDataset,
+    and it is case-insensitive.
+    """
+    dataset_filename = "data." + dataset_name + "_dataset"
+    datasetlib = importlib.import_module(dataset_filename)
+
+    dataset = None
+    target_dataset_name = dataset_name.replace("_", "") + "dataset"
+    for name, cls in datasetlib.__dict__.items():
+        if name.lower() == target_dataset_name.lower() and issubclass(cls, BaseDataset):
+            dataset = cls
+
+    if dataset is None:
+        raise NotImplementedError("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name))
+
+    return dataset
+
+
+def get_option_setter(dataset_name):
+    """Return the static method <modify_commandline_options> of the dataset class."""
+    dataset_class = find_dataset_using_name(dataset_name)
+    return dataset_class.modify_commandline_options
+
+
+def create_dataset(opt):
+    """Create a dataset given the option.
+
+    This function wraps the class CustomDatasetDataLoader.
+        This is the main interface between this package and 'train.py'/'test.py'
+
+    Example:
+        >>> from data import create_dataset
+        >>> dataset = create_dataset(opt)
+    """
+    data_loader = CustomDatasetDataLoader(opt)
+    dataset = data_loader.load_data()
+    return dataset
+
+
+class CustomDatasetDataLoader:
+    """Wrapper class of Dataset class that performs multi-threaded data loading"""
+
+    def __init__(self, opt):
+        """Initialize this class
+
+        Step 1: create a dataset instance given the name [dataset_mode]
+        Step 2: create a multi-threaded data loader.
+        """
+        self.opt = opt
+        dataset_class = find_dataset_using_name(opt.dataset_mode)
+        self.dataset = dataset_class(opt)
+        print("dataset [%s] was created" % type(self.dataset).__name__)
+
+        # Use DistributedSampler for DDP training
+        if "LOCAL_RANK" in os.environ:
+            print(f'create DDP sampler on rank {int(os.environ["LOCAL_RANK"])}')
+            self.sampler = DistributedSampler(self.dataset, shuffle=not opt.serial_batches)
+            shuffle = False  # DistributedSampler handles shuffling
+        else:
+            self.sampler = None
+            shuffle = not opt.serial_batches
+
+        self.dataloader = torch.utils.data.DataLoader(self.dataset, batch_size=opt.batch_size, shuffle=shuffle, sampler=self.sampler, num_workers=int(opt.num_threads))
+
+    def load_data(self):
+        return self
+
+    def __len__(self):
+        """Return the number of data in the dataset"""
+        return min(len(self.dataset), self.opt.max_dataset_size)
+
+    def __iter__(self):
+        """Return a batch of data"""
+        for i, data in enumerate(self.dataloader):
+            if i * self.opt.batch_size >= self.opt.max_dataset_size:
+                break
+            yield data
+
+    def set_epoch(self, epoch):
+        """Set epoch for DistributedSampler to ensure proper shuffling"""
+        if self.sampler is not None:
+            self.sampler.set_epoch(epoch)

data/aligned_dataset.py

@@ -0,0 +1,60 @@
+import os
+from data.base_dataset import BaseDataset, get_params, get_transform
+from data.image_folder import make_dataset
+from PIL import Image
+
+
+class AlignedDataset(BaseDataset):
+    """A dataset class for paired image dataset.
+
+    It assumes that the directory '/path/to/data/train' contains image pairs in the form of {A,B}.
+    During test time, you need to prepare a directory '/path/to/data/test'.
+    """
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseDataset.__init__(self, opt)
+        self.dir_AB = os.path.join(opt.dataroot, opt.phase)  # get the image directory
+        self.AB_paths = sorted(make_dataset(self.dir_AB, opt.max_dataset_size))  # get image paths
+        assert self.opt.load_size >= self.opt.crop_size  # crop_size should be smaller than the size of loaded image
+        self.input_nc = self.opt.output_nc if self.opt.direction == "BtoA" else self.opt.input_nc
+        self.output_nc = self.opt.input_nc if self.opt.direction == "BtoA" else self.opt.output_nc
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns a dictionary that contains A, B, A_paths and B_paths
+            A (tensor) - - an image in the input domain
+            B (tensor) - - its corresponding image in the target domain
+            A_paths (str) - - image paths
+            B_paths (str) - - image paths (same as A_paths)
+        """
+        # read a image given a random integer index
+        AB_path = self.AB_paths[index]
+        AB = Image.open(AB_path).convert("RGB")
+        # split AB image into A and B
+        w, h = AB.size
+        w2 = int(w / 2)
+        A = AB.crop((0, 0, w2, h))
+        B = AB.crop((w2, 0, w, h))
+
+        # apply the same transform to both A and B
+        transform_params = get_params(self.opt, A.size)
+        A_transform = get_transform(self.opt, transform_params, grayscale=(self.input_nc == 1))
+        B_transform = get_transform(self.opt, transform_params, grayscale=(self.output_nc == 1))
+
+        A = A_transform(A)
+        B = B_transform(B)
+
+        return {"A": A, "B": B, "A_paths": AB_path, "B_paths": AB_path}
+
+    def __len__(self):
+        """Return the total number of images in the dataset."""
+        return len(self.AB_paths)

data/base_dataset.py

@@ -0,0 +1,167 @@
+"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.
+
+It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.
+"""
+
+import random
+import numpy as np
+import torch.utils.data as data
+from PIL import Image
+import torchvision.transforms as transforms
+from abc import ABC, abstractmethod
+
+
+class BaseDataset(data.Dataset, ABC):
+    """This class is an abstract base class (ABC) for datasets.
+
+    To create a subclass, you need to implement the following four functions:
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point.
+    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
+    """
+
+    def __init__(self, opt):
+        """Initialize the class; save the options in the class
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        self.opt = opt
+        self.root = opt.dataroot
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        return parser
+
+    @abstractmethod
+    def __len__(self):
+        """Return the total number of images in the dataset."""
+        return 0
+
+    @abstractmethod
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns:
+            a dictionary of data with their names. It ususally contains the data itself and its metadata information.
+        """
+        pass
+
+
+def get_params(opt, size):
+    w, h = size
+    new_h = h
+    new_w = w
+    if opt.preprocess == "resize_and_crop":
+        new_h = new_w = opt.load_size
+    elif opt.preprocess == "scale_width_and_crop":
+        new_w = opt.load_size
+        new_h = opt.load_size * h // w
+
+    x = random.randint(0, np.maximum(0, new_w - opt.crop_size))
+    y = random.randint(0, np.maximum(0, new_h - opt.crop_size))
+
+    flip = random.random() > 0.5
+
+    return {"crop_pos": (x, y), "flip": flip}
+
+
+def get_transform(opt, params=None, grayscale=False, method=transforms.InterpolationMode.BICUBIC, convert=True):
+    transform_list = []
+    if grayscale:
+        transform_list.append(transforms.Grayscale(1))
+    if "resize" in opt.preprocess:
+        osize = [opt.load_size, opt.load_size]
+        transform_list.append(transforms.Resize(osize, method))
+    elif "scale_width" in opt.preprocess:
+        transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, opt.crop_size, method)))
+
+    if "crop" in opt.preprocess:
+        if params is None:
+            transform_list.append(transforms.RandomCrop(opt.crop_size))
+        else:
+            transform_list.append(transforms.Lambda(lambda img: __crop(img, params["crop_pos"], opt.crop_size)))
+
+    if opt.preprocess == "none":
+        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base=4, method=method)))
+
+    if not opt.no_flip:
+        if params is None:
+            transform_list.append(transforms.RandomHorizontalFlip())
+        elif params["flip"]:
+            transform_list.append(transforms.Lambda(lambda img: __flip(img, params["flip"])))
+
+    if convert:
+        transform_list += [transforms.ToTensor()]
+        if grayscale:
+            transform_list += [transforms.Normalize((0.5,), (0.5,))]
+        else:
+            transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+    return transforms.Compose(transform_list)
+
+
+def __transforms2pil_resize(method):
+    mapper = {
+        transforms.InterpolationMode.BILINEAR: Image.BILINEAR,
+        transforms.InterpolationMode.BICUBIC: Image.BICUBIC,
+        transforms.InterpolationMode.NEAREST: Image.NEAREST,
+        transforms.InterpolationMode.LANCZOS: Image.LANCZOS,
+    }
+    return mapper[method]
+
+
+def __make_power_2(img, base, method=transforms.InterpolationMode.BICUBIC):
+    method = __transforms2pil_resize(method)
+    ow, oh = img.size
+    h = int(round(oh / base) * base)
+    w = int(round(ow / base) * base)
+    if h == oh and w == ow:
+        return img
+
+    __print_size_warning(ow, oh, w, h)
+    return img.resize((w, h), method)
+
+
+def __scale_width(img, target_size, crop_size, method=transforms.InterpolationMode.BICUBIC):
+    method = __transforms2pil_resize(method)
+    ow, oh = img.size
+    if ow == target_size and oh >= crop_size:
+        return img
+    w = target_size
+    h = int(max(target_size * oh / ow, crop_size))
+    return img.resize((w, h), method)
+
+
+def __crop(img, pos, size):
+    ow, oh = img.size
+    x1, y1 = pos
+    tw = th = size
+    if ow > tw or oh > th:
+        return img.crop((x1, y1, x1 + tw, y1 + th))
+    return img
+
+
+def __flip(img, flip):
+    if flip:
+        return img.transpose(Image.FLIP_LEFT_RIGHT)
+    return img
+
+
+def __print_size_warning(ow, oh, w, h):
+    """Print warning information about image size(only print once)"""
+    if not hasattr(__print_size_warning, "has_printed"):
+        print("The image size needs to be a multiple of 4. " "The loaded image size was (%d, %d), so it was adjusted to " "(%d, %d). This adjustment will be done to all images " "whose sizes are not multiples of 4" % (ow, oh, w, h))
+        __print_size_warning.has_printed = True

data/colorization_dataset.py

@@ -0,0 +1,69 @@
+import os
+from data.base_dataset import BaseDataset, get_transform
+from data.image_folder import make_dataset
+from skimage import color  # require skimage
+from PIL import Image
+import numpy as np
+import torchvision.transforms as transforms
+
+
+class ColorizationDataset(BaseDataset):
+    """This dataset class can load a set of natural images in RGB, and convert RGB format into (L, ab) pairs in Lab color space.
+
+    This dataset is required by pix2pix-based colorization model ('--model colorization')
+    """
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+
+        By default, the number of channels for input image  is 1 (L) and
+        the number of channels for output image is 2 (ab). The direction is from A to B
+        """
+        parser.set_defaults(input_nc=1, output_nc=2, direction="AtoB")
+        return parser
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseDataset.__init__(self, opt)
+        self.dir = os.path.join(opt.dataroot, opt.phase)
+        self.AB_paths = sorted(make_dataset(self.dir, opt.max_dataset_size))
+        assert opt.input_nc == 1 and opt.output_nc == 2 and opt.direction == "AtoB"
+        self.transform = get_transform(self.opt, convert=False)
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns a dictionary that contains A, B, A_paths and B_paths
+            A (tensor) - - the L channel of an image
+            B (tensor) - - the ab channels of the same image
+            A_paths (str) - - image paths
+            B_paths (str) - - image paths (same as A_paths)
+        """
+        path = self.AB_paths[index]
+        im = Image.open(path).convert("RGB")
+        im = self.transform(im)
+        im = np.array(im)
+        lab = color.rgb2lab(im).astype(np.float32)
+        lab_t = transforms.ToTensor()(lab)
+        A = lab_t[[0], ...] / 50.0 - 1.0
+        B = lab_t[[1, 2], ...] / 110.0
+        return {"A": A, "B": B, "A_paths": path, "B_paths": path}
+
+    def __len__(self):
+        """Return the total number of images in the dataset."""
+        return len(self.AB_paths)

data/image_folder.py

@@ -0,0 +1,72 @@
+"""A modified image folder class
+
+We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)
+so that this class can load images from both current directory and its subdirectories.
+"""
+
+import torch.utils.data as data
+from pathlib import Path
+from PIL import Image
+
+IMG_EXTENSIONS = [
+    ".jpg",
+    ".JPG",
+    ".jpeg",
+    ".JPEG",
+    ".png",
+    ".PNG",
+    ".ppm",
+    ".PPM",
+    ".bmp",
+    ".BMP",
+    ".tif",
+    ".TIF",
+    ".tiff",
+    ".TIFF",
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir, max_dataset_size=float("inf")):
+    images = []
+    dir_path = Path(dir)
+    assert dir_path.is_dir(), f"{dir} is not a valid directory"
+
+    for path in sorted(dir_path.rglob("*")):
+        if path.is_file() and is_image_file(path.name):
+            images.append(str(path))
+    return images[: min(max_dataset_size, len(images))]
+
+
+def default_loader(path):
+    return Image.open(path).convert("RGB")
+
+
+class ImageFolder(data.Dataset):
+
+    def __init__(self, root, transform=None, return_paths=False, loader=default_loader):
+        imgs = make_dataset(root)
+        if len(imgs) == 0:
+            raise (RuntimeError("Found 0 images in: " + root + "\n" "Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
+
+        self.root = root
+        self.imgs = imgs
+        self.transform = transform
+        self.return_paths = return_paths
+        self.loader = loader
+
+    def __getitem__(self, index):
+        path = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.return_paths:
+            return img, path
+        else:
+            return img
+
+    def __len__(self):
+        return len(self.imgs)

data/single_dataset.py

@@ -0,0 +1,40 @@
+from data.base_dataset import BaseDataset, get_transform
+from data.image_folder import make_dataset
+from PIL import Image
+
+
+class SingleDataset(BaseDataset):
+    """This dataset class can load a set of images specified by the path --dataroot /path/to/data.
+
+    It can be used for generating CycleGAN results only for one side with the model option '-model test'.
+    """
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseDataset.__init__(self, opt)
+        self.A_paths = sorted(make_dataset(opt.dataroot, opt.max_dataset_size))
+        input_nc = self.opt.output_nc if self.opt.direction == "BtoA" else self.opt.input_nc
+        self.transform = get_transform(opt, grayscale=(input_nc == 1))
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns a dictionary that contains A and A_paths
+            A(tensor) - - an image in one domain
+            A_paths(str) - - the path of the image
+        """
+        A_path = self.A_paths[index]
+        A_img = Image.open(A_path).convert("RGB")
+        A = self.transform(A_img)
+        return {"A": A, "A_paths": A_path}
+
+    def __len__(self):
+        """Return the total number of images in the dataset."""
+        return len(self.A_paths)

data/template_dataset.py

@@ -0,0 +1,78 @@
+"""Dataset class template
+
+This module provides a template for users to implement custom datasets.
+You can specify '--dataset_mode template' to use this dataset.
+The class name should be consistent with both the filename and its dataset_mode option.
+The filename should be <dataset_mode>_dataset.py
+The class name should be <Dataset_mode>Dataset.py
+You need to implement the following functions:
+    -- <modify_commandline_options>:　Add dataset-specific options and rewrite default values for existing options.
+    -- <__init__>: Initialize this dataset class.
+    -- <__getitem__>: Return a data point and its metadata information.
+    -- <__len__>: Return the number of images.
+"""
+
+from data.base_dataset import BaseDataset, get_transform
+
+# from data.image_folder import make_dataset
+# from PIL import Image
+
+
+class TemplateDataset(BaseDataset):
+    """A template dataset class for you to implement custom datasets."""
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        parser.add_argument("--new_dataset_option", type=float, default=1.0, help="new dataset option")
+        parser.set_defaults(max_dataset_size=10, new_dataset_option=2.0)  # specify dataset-specific default values
+        return parser
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+
+        A few things can be done here.
+        - save the options (have been done in BaseDataset)
+        - get image paths and meta information of the dataset.
+        - define the image transformation.
+        """
+        # save the option and dataset root
+        BaseDataset.__init__(self, opt)
+        # get the image paths of your dataset;
+        self.image_paths = []  # You can call sorted(make_dataset(self.root, opt.max_dataset_size)) to get all the image paths under the directory self.root
+        # define the default transform function. You can use <base_dataset.get_transform>; You can also define your custom transform function
+        self.transform = get_transform(opt)
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index -- a random integer for data indexing
+
+        Returns:
+            a dictionary of data with their names. It usually contains the data itself and its metadata information.
+
+        Step 1: get a random image path: e.g., path = self.image_paths[index]
+        Step 2: load your data from the disk: e.g., image = Image.open(path).convert('RGB').
+        Step 3: convert your data to a PyTorch tensor. You can use helpder functions such as self.transform. e.g., data = self.transform(image)
+        Step 4: return a data point as a dictionary.
+        """
+        path = "temp"  # needs to be a string
+        data_A = None  # needs to be a tensor
+        data_B = None  # needs to be a tensor
+        return {"data_A": data_A, "data_B": data_B, "path": path}
+
+    def __len__(self):
+        """Return the total number of images."""
+        return len(self.image_paths)

data/unaligned_dataset.py

@@ -0,0 +1,71 @@
+import os
+from data.base_dataset import BaseDataset, get_transform
+from data.image_folder import make_dataset
+from PIL import Image
+import random
+
+
+class UnalignedDataset(BaseDataset):
+    """
+    This dataset class can load unaligned/unpaired datasets.
+
+    It requires two directories to host training images from domain A '/path/to/data/trainA'
+    and from domain B '/path/to/data/trainB' respectively.
+    You can train the model with the dataset flag '--dataroot /path/to/data'.
+    Similarly, you need to prepare two directories:
+    '/path/to/data/testA' and '/path/to/data/testB' during test time.
+    """
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseDataset.__init__(self, opt)
+        self.dir_A = os.path.join(opt.dataroot, opt.phase + "A")  # create a path '/path/to/data/trainA'
+        self.dir_B = os.path.join(opt.dataroot, opt.phase + "B")  # create a path '/path/to/data/trainB'
+
+        self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size))  # load images from '/path/to/data/trainA'
+        self.B_paths = sorted(make_dataset(self.dir_B, opt.max_dataset_size))  # load images from '/path/to/data/trainB'
+        self.A_size = len(self.A_paths)  # get the size of dataset A
+        self.B_size = len(self.B_paths)  # get the size of dataset B
+        btoA = self.opt.direction == "BtoA"
+        input_nc = self.opt.output_nc if btoA else self.opt.input_nc  # get the number of channels of input image
+        output_nc = self.opt.input_nc if btoA else self.opt.output_nc  # get the number of channels of output image
+        self.transform_A = get_transform(self.opt, grayscale=(input_nc == 1))
+        self.transform_B = get_transform(self.opt, grayscale=(output_nc == 1))
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index (int)      -- a random integer for data indexing
+
+        Returns a dictionary that contains A, B, A_paths and B_paths
+            A (tensor)       -- an image in the input domain
+            B (tensor)       -- its corresponding image in the target domain
+            A_paths (str)    -- image paths
+            B_paths (str)    -- image paths
+        """
+        A_path = self.A_paths[index % self.A_size]  # make sure index is within then range
+        if self.opt.serial_batches:  # make sure index is within then range
+            index_B = index % self.B_size
+        else:  # randomize the index for domain B to avoid fixed pairs.
+            index_B = random.randint(0, self.B_size - 1)
+        B_path = self.B_paths[index_B]
+        A_img = Image.open(A_path).convert("RGB")
+        B_img = Image.open(B_path).convert("RGB")
+        # apply image transformation
+        A = self.transform_A(A_img)
+        B = self.transform_B(B_img)
+
+        return {"A": A, "B": B, "A_paths": A_path, "B_paths": B_path}
+
+    def __len__(self):
+        """Return the total number of images in the dataset.
+
+        As we have two datasets with potentially different number of images,
+        we take a maximum of
+        """
+        return max(self.A_size, self.B_size)

models/__init__.py

@@ -0,0 +1,58 @@
+"""This package contains modules related to objective functions, optimizations, and network architectures.
+
+To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
+You need to implement the following five functions:
+    -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
+    -- <set_input>:                     unpack data from dataset and apply preprocessing.
+    -- <forward>:                       produce intermediate results.
+    -- <optimize_parameters>:           calculate loss, gradients, and update network weights.
+    -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
+
+In the function <__init__>, you need to define four lists:
+    -- self.loss_names (str list):          specify the training losses that you want to plot and save.
+    -- self.model_names (str list):         define networks used in our training.
+    -- self.visual_names (str list):        specify the images that you want to display and save.
+    -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.
+
+Now you can use the model class by specifying flag '--model dummy'.
+See our template model class 'template_model.py' for more details.
+"""
+
+import importlib
+from models.base_model import BaseModel
+
+
+def find_model_using_name(model_name: str):
+    """Import the module "models/[model_name]_model.py".
+
+    In the file, the class called DatasetNameModel() will
+    be instantiated. It has to be a subclass of BaseModel,
+    and it is case-insensitive.
+    """
+    model_filename = "models." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+    model = None
+    target_model_name = model_name.replace("_", "") + "model"
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() and issubclass(cls, BaseModel):
+            model = cls
+
+    if model is None:
+        print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
+        exit(0)
+
+    return model
+
+
+def get_option_setter(model_name: str):
+    """Return the static method <modify_commandline_options> of the model class."""
+    model_class = find_model_using_name(model_name)
+    return model_class.modify_commandline_options
+
+
+def create_model(opt):
+    """Create a model given the option."""
+    model = find_model_using_name(opt.model)
+    instance = model(opt)
+    print(f"model [{type(instance).__name__}] was created")
+    return instance

models/base_model.py

@@ -0,0 +1,306 @@
+import os
+import torch
+import torch.distributed as dist
+from pathlib import Path
+from collections import OrderedDict
+from abc import ABC, abstractmethod
+from . import networks
+
+
+class BaseModel(ABC):
+    """This class is an abstract base class (ABC) for models.
+    To create a subclass, you need to implement the following five functions:
+        -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
+        -- <set_input>:                     unpack data from dataset and apply preprocessing.
+        -- <forward>:                       produce intermediate results.
+        -- <optimize_parameters>:           calculate losses, gradients, and update network weights.
+        -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
+    """
+
+    def __init__(self, opt):
+        """Initialize the BaseModel class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+
+        When creating your custom class, you need to implement your own initialization.
+        In this function, you should first call <BaseModel.__init__(self, opt)>
+        Then, you need to define four lists:
+            -- self.loss_names (str list):          specify the training losses that you want to plot and save.
+            -- self.model_names (str list):         define networks used in our training.
+            -- self.visual_names (str list):        specify the images that you want to display and save.
+            -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
+        """
+        self.opt = opt
+        self.isTrain = opt.isTrain
+        self.save_dir = Path(opt.checkpoints_dir) / opt.name  # save all the checkpoints to save_dir
+        self.device = opt.device
+        # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
+        if opt.preprocess != "scale_width":
+            torch.backends.cudnn.benchmark = True
+        self.loss_names = []
+        self.model_names = []
+        self.visual_names = []
+        self.optimizers = []
+        self.image_paths = []
+        self.metric = 0  # used for learning rate policy 'plateau'
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new model-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        return parser
+
+    @abstractmethod
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input (dict): includes the data itself and its metadata information.
+        """
+        pass
+
+    @abstractmethod
+    def forward(self):
+        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
+        pass
+
+    @abstractmethod
+    def optimize_parameters(self):
+        """Calculate losses, gradients, and update network weights; called in every training iteration"""
+        pass
+
+    def setup(self, opt):
+        """Load and print networks; create schedulers
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        # Initialize all networks and load if needed
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, "net" + name)
+                net = networks.init_net(net, opt.init_type, opt.init_gain)
+
+                # Load networks if needed
+                if not self.isTrain or opt.continue_train:
+                    load_suffix = f"iter_{opt.load_iter}" if opt.load_iter > 0 else opt.epoch
+                    load_filename = f"{load_suffix}_net_{name}.pth"
+                    load_path = self.save_dir / load_filename
+
+                    if isinstance(net, torch.nn.parallel.DistributedDataParallel):
+                        net = net.module
+                    print(f"loading the model from {load_path}")
+
+                    state_dict = torch.load(load_path, map_location=str(self.device), weights_only=True)
+
+                    if hasattr(state_dict, "_metadata"):
+                        del state_dict._metadata
+
+                    # patch InstanceNorm checkpoints
+                    for key in list(state_dict.keys()):
+                        self.__patch_instance_norm_state_dict(state_dict, net, key.split("."))
+                    net.load_state_dict(state_dict)
+
+                # Move network to device
+                net.to(self.device)
+
+                # Wrap networks with DDP after loading
+                if dist.is_initialized():
+                    # Check if using syncbatch normalization for DDP
+                    if self.opt.norm == "syncbatch":
+                        raise ValueError(f"For distributed training, opt.norm must be 'syncbatch' or 'inst', but got '{self.opt.norm}'. " "Please set --norm syncbatch for multi-GPU training.")
+
+                    net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[self.device.index])
+                    # Sync all processes after DDP wrapping
+                    dist.barrier()
+
+                setattr(self, "net" + name, net)
+
+        self.print_networks(opt.verbose)
+
+        if self.isTrain:
+            self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
+
+    def eval(self):
+        """Make models eval mode during test time"""
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, "net" + name)
+                net.eval()
+
+    def test(self):
+        """Forward function used in test time.
+
+        This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
+        It also calls <compute_visuals> to produce additional visualization results
+        """
+        with torch.no_grad():
+            self.forward()
+            self.compute_visuals()
+
+    def compute_visuals(self):
+        """Calculate additional output images for visdom and HTML visualization"""
+        pass
+
+    def get_image_paths(self):
+        """Return image paths that are used to load current data"""
+        return self.image_paths
+
+    def update_learning_rate(self):
+        """Update learning rates for all the networks; called at the end of every epoch"""
+        old_lr = self.optimizers[0].param_groups[0]["lr"]
+        for scheduler in self.schedulers:
+            if self.opt.lr_policy == "plateau":
+                scheduler.step(self.metric)
+            else:
+                scheduler.step()
+
+        lr = self.optimizers[0].param_groups[0]["lr"]
+        print(f"learning rate {old_lr:.7f} -> {lr:.7f}")
+
+    def get_current_visuals(self):
+        """Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""
+        visual_ret = OrderedDict()
+        for name in self.visual_names:
+            if isinstance(name, str):
+                visual_ret[name] = getattr(self, name)
+        return visual_ret
+
+    def get_current_losses(self):
+        """Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""
+        errors_ret = OrderedDict()
+        for name in self.loss_names:
+            if isinstance(name, str):
+                errors_ret[name] = float(getattr(self, "loss_" + name))  # float(...) works for both scalar tensor and float number
+        return errors_ret
+
+    def save_networks(self, epoch):
+        """Save all the networks to the disk, unwrapping them first."""
+
+        # Only allow the main process (rank 0) to save the checkpoint
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            for name in self.model_names:
+                if isinstance(name, str):
+                    save_filename = f"{epoch}_net_{name}.pth"
+                    save_path = self.save_dir / save_filename
+                    net = getattr(self, "net" + name)
+
+                    # 1. First, unwrap from DDP if it exists
+                    if hasattr(net, "module"):
+                        model_to_save = net.module
+                    else:
+                        model_to_save = net
+
+                    # 2. Second, unwrap from torch.compile if it exists
+                    if hasattr(model_to_save, "_orig_mod"):
+                        model_to_save = model_to_save._orig_mod
+
+                    # 3. Save the final, clean state_dict
+                    torch.save(model_to_save.state_dict(), save_path)
+
+    def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
+        """Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
+        key = keys[i]
+        if i + 1 == len(keys):  # at the end, pointing to a parameter/buffer
+            if module.__class__.__name__.startswith("InstanceNorm") and (key == "running_mean" or key == "running_var"):
+                if getattr(module, key) is None:
+                    state_dict.pop(".".join(keys))
+            if module.__class__.__name__.startswith("InstanceNorm") and (key == "num_batches_tracked"):
+                state_dict.pop(".".join(keys))
+        else:
+            self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
+
+    def load_networks(self, epoch):
+        """Load all networks from the disk for DDP."""
+
+        for name in self.model_names:
+            if isinstance(name, str):
+                load_filename = f"{epoch}_net_{name}.pth"
+                load_path = self.save_dir / load_filename
+                net = getattr(self, "net" + name)
+
+                if isinstance(net, torch.nn.parallel.DistributedDataParallel):
+                    net = net.module
+                print(f"loading the model from {load_path}")
+
+                state_dict = torch.load(load_path, map_location=str(self.device), weights_only=True)
+
+                if hasattr(state_dict, "_metadata"):
+                    del state_dict._metadata
+
+                # patch InstanceNorm checkpoints
+                for key in list(state_dict.keys()):
+                    self.__patch_instance_norm_state_dict(state_dict, net, key.split("."))
+                net.load_state_dict(state_dict)
+
+        # Add a barrier to sync all processes before continuing
+        if dist.is_initialized():
+            dist.barrier()
+
+    def print_networks(self, verbose):
+        """Print the total number of parameters in the network and (if verbose) network architecture
+
+        Parameters:
+            verbose (bool) -- if verbose: print the network architecture
+        """
+        print("---------- Networks initialized -------------")
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, "net" + name)
+                num_params = 0
+                for param in net.parameters():
+                    num_params += param.numel()
+                if verbose:
+                    print(net)
+                print(f"[Network {name}] Total number of parameters : {num_params / 1e6:.3f} M")
+        print("-----------------------------------------------")
+
+    def set_requires_grad(self, nets, requires_grad=False):
+        """Set requies_grad=Fasle for all the networks to avoid unnecessary computations
+        Parameters:
+            nets (network list)   -- a list of networks
+            requires_grad (bool)  -- whether the networks require gradients or not
+        """
+        if not isinstance(nets, list):
+            nets = [nets]
+        for net in nets:
+            if net is not None:
+                for param in net.parameters():
+                    param.requires_grad = requires_grad
+
+    def init_networks(self, init_type="normal", init_gain=0.02):
+        """Initialize all networks: 1. move to device; 2. initialize weights
+
+        Parameters:
+            init_type (str) -- initialization method: normal | xavier | kaiming | orthogonal
+            init_gain (float) -- scaling factor for normal, xavier and orthogonal
+        """
+        import os
+
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, "net" + name)
+
+                # Move to device
+                if torch.cuda.is_available():
+                    if "LOCAL_RANK" in os.environ:
+                        local_rank = int(os.environ["LOCAL_RANK"])
+                        net.to(local_rank)
+                        print(f"Initialized network {name} with device cuda:{local_rank}")
+                    else:
+                        net.to(0)
+                        print(f"Initialized network {name} with device cuda:0")
+                else:
+                    net.to("cpu")
+                    print(f"Initialized network {name} with device cpu")
+
+                # Initialize weights using networks function
+                networks.init_weights(net, init_type, init_gain)

models/colorization_model.py

@@ -0,0 +1,69 @@
+from .pix2pix_model import Pix2PixModel
+import torch
+from skimage import color  # used for lab2rgb
+import numpy as np
+
+
+class ColorizationModel(Pix2PixModel):
+    """This is a subclass of Pix2PixModel for image colorization (black & white image -> colorful images).
+
+    The model training requires '-dataset_model colorization' dataset.
+    It trains a pix2pix model, mapping from L channel to ab channels in Lab color space.
+    By default, the colorization dataset will automatically set '--input_nc 1' and '--output_nc 2'.
+    """
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train=True):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+
+        By default, we use 'colorization' dataset for this model.
+        See the original pix2pix paper (https://arxiv.org/pdf/1611.07004.pdf) and colorization results (Figure 9 in the paper)
+        """
+        Pix2PixModel.modify_commandline_options(parser, is_train)
+        parser.set_defaults(dataset_mode="colorization")
+        return parser
+
+    def __init__(self, opt):
+        """Initialize the class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+
+        For visualization, we set 'visual_names' as 'real_A' (input real image),
+        'real_B_rgb' (ground truth RGB image), and 'fake_B_rgb' (predicted RGB image)
+        We convert the Lab image 'real_B' (inherited from Pix2pixModel) to a RGB image 'real_B_rgb'.
+        we convert the Lab image 'fake_B' (inherited from Pix2pixModel) to a RGB image 'fake_B_rgb'.
+        """
+        # reuse the pix2pix model
+        Pix2PixModel.__init__(self, opt)
+        # specify the images to be visualized.
+        self.visual_names = ["real_A", "real_B_rgb", "fake_B_rgb"]
+
+    def lab2rgb(self, L, AB):
+        """Convert an Lab tensor image to a RGB numpy output
+        Parameters:
+            L  (1-channel tensor array): L channel images (range: [-1, 1], torch tensor array)
+            AB (2-channel tensor array):  ab channel images (range: [-1, 1], torch tensor array)
+
+        Returns:
+            rgb (RGB numpy image): rgb output images  (range: [0, 255], numpy array)
+        """
+        AB2 = AB * 110.0
+        L2 = (L + 1.0) * 50.0
+        Lab = torch.cat([L2, AB2], dim=1)
+        Lab = Lab[0].data.cpu().float().numpy()
+        Lab = np.transpose(Lab.astype(np.float64), (1, 2, 0))
+        rgb = color.lab2rgb(Lab) * 255
+        return rgb
+
+    def compute_visuals(self):
+        """Calculate additional output images for visdom and HTML visualization"""
+        self.real_B_rgb = self.lab2rgb(self.real_A, self.real_B)
+        self.fake_B_rgb = self.lab2rgb(self.real_A, self.fake_B)

models/cycle_gan_model.py

@@ -0,0 +1,196 @@
+import torch
+import itertools
+from util.image_pool import ImagePool
+from .base_model import BaseModel
+from . import networks
+
+
+class CycleGANModel(BaseModel):
+    """
+    This class implements the CycleGAN model, for learning image-to-image translation without paired data.
+
+    The model training requires '--dataset_mode unaligned' dataset.
+    By default, it uses a '--netG resnet_9blocks' ResNet generator,
+    a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
+    and a least-square GANs objective ('--gan_mode lsgan').
+
+    CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
+    """
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train=True):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+
+        For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
+        A (source domain), B (target domain).
+        Generators: G_A: A -> B; G_B: B -> A.
+        Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
+        Forward cycle loss:  lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
+        Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
+        Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
+        Dropout is not used in the original CycleGAN paper.
+        """
+        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
+        if is_train:
+            parser.add_argument("--lambda_A", type=float, default=10.0, help="weight for cycle loss (A -> B -> A)")
+            parser.add_argument("--lambda_B", type=float, default=10.0, help="weight for cycle loss (B -> A -> B)")
+            parser.add_argument(
+                "--lambda_identity",
+                type=float,
+                default=0.5,
+                help="use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1",
+            )
+
+        return parser
+
+    def __init__(self, opt):
+        """Initialize the CycleGAN class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseModel.__init__(self, opt)
+        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
+        self.loss_names = ["D_A", "G_A", "cycle_A", "idt_A", "D_B", "G_B", "cycle_B", "idt_B"]
+        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
+        visual_names_A = ["real_A", "fake_B", "rec_A"]
+        visual_names_B = ["real_B", "fake_A", "rec_B"]
+        if self.isTrain and self.opt.lambda_identity > 0.0:  # if identity loss is used, we also visualize idt_B=G_A(B) ad idt_A=G_B(A)
+            visual_names_A.append("idt_B")
+            visual_names_B.append("idt_A")
+
+        self.visual_names = visual_names_A + visual_names_B  # combine visualizations for A and B
+        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
+        if self.isTrain:
+            self.model_names = ["G_A", "G_B", "D_A", "D_B"]
+        else:  # during test time, only load Gs
+            self.model_names = ["G_A", "G_B"]
+
+        # define networks (both Generators and discriminators)
+        # The naming is different from those used in the paper.
+        # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
+        self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain)
+        self.netG_B = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain)
+
+        if self.isTrain:  # define discriminators
+            self.netD_A = networks.define_D(opt.output_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain)
+            self.netD_B = networks.define_D(opt.input_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain)
+
+        if self.isTrain:
+            if opt.lambda_identity > 0.0:  # only works when input and output images have the same number of channels
+                assert opt.input_nc == opt.output_nc
+            self.fake_A_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
+            self.fake_B_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
+            # define loss functions
+            self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)  # define GAN loss.
+            self.criterionCycle = torch.nn.L1Loss()
+            self.criterionIdt = torch.nn.L1Loss()
+            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
+            self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
+            self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
+            self.optimizers.append(self.optimizer_G)
+            self.optimizers.append(self.optimizer_D)
+
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input (dict): include the data itself and its metadata information.
+
+        The option 'direction' can be used to swap domain A and domain B.
+        """
+        AtoB = self.opt.direction == "AtoB"
+        self.real_A = input["A" if AtoB else "B"].to(self.device)
+        self.real_B = input["B" if AtoB else "A"].to(self.device)
+        self.image_paths = input["A_paths" if AtoB else "B_paths"]
+
+    def forward(self):
+        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
+        self.fake_B = self.netG_A(self.real_A)  # G_A(A)
+        self.rec_A = self.netG_B(self.fake_B)  # G_B(G_A(A))
+        self.fake_A = self.netG_B(self.real_B)  # G_B(B)
+        self.rec_B = self.netG_A(self.fake_A)  # G_A(G_B(B))
+
+    def backward_D_basic(self, netD, real, fake):
+        """Calculate GAN loss for the discriminator
+
+        Parameters:
+            netD (network)      -- the discriminator D
+            real (tensor array) -- real images
+            fake (tensor array) -- images generated by a generator
+
+        Return the discriminator loss.
+        We also call loss_D.backward() to calculate the gradients.
+        """
+        # Real
+        pred_real = netD(real)
+        loss_D_real = self.criterionGAN(pred_real, True)
+        # Fake
+        pred_fake = netD(fake.detach())
+        loss_D_fake = self.criterionGAN(pred_fake, False)
+        # Combined loss and calculate gradients
+        loss_D = (loss_D_real + loss_D_fake) * 0.5
+        loss_D.backward()
+        return loss_D
+
+    def backward_D_A(self):
+        """Calculate GAN loss for discriminator D_A"""
+        fake_B = self.fake_B_pool.query(self.fake_B)
+        self.loss_D_A = self.backward_D_basic(self.netD_A, self.real_B, fake_B)
+
+    def backward_D_B(self):
+        """Calculate GAN loss for discriminator D_B"""
+        fake_A = self.fake_A_pool.query(self.fake_A)
+        self.loss_D_B = self.backward_D_basic(self.netD_B, self.real_A, fake_A)
+
+    def backward_G(self):
+        """Calculate the loss for generators G_A and G_B"""
+        lambda_idt = self.opt.lambda_identity
+        lambda_A = self.opt.lambda_A
+        lambda_B = self.opt.lambda_B
+        # Identity loss
+        if lambda_idt > 0:
+            # G_A should be identity if real_B is fed: ||G_A(B) - B||
+            self.idt_A = self.netG_A(self.real_B)
+            self.loss_idt_A = self.criterionIdt(self.idt_A, self.real_B) * lambda_B * lambda_idt
+            # G_B should be identity if real_A is fed: ||G_B(A) - A||
+            self.idt_B = self.netG_B(self.real_A)
+            self.loss_idt_B = self.criterionIdt(self.idt_B, self.real_A) * lambda_A * lambda_idt
+        else:
+            self.loss_idt_A = 0
+            self.loss_idt_B = 0
+
+        # GAN loss D_A(G_A(A))
+        self.loss_G_A = self.criterionGAN(self.netD_A(self.fake_B), True)
+        # GAN loss D_B(G_B(B))
+        self.loss_G_B = self.criterionGAN(self.netD_B(self.fake_A), True)
+        # Forward cycle loss || G_B(G_A(A)) - A||
+        self.loss_cycle_A = self.criterionCycle(self.rec_A, self.real_A) * lambda_A
+        # Backward cycle loss || G_A(G_B(B)) - B||
+        self.loss_cycle_B = self.criterionCycle(self.rec_B, self.real_B) * lambda_B
+        # combined loss and calculate gradients
+        self.loss_G = self.loss_G_A + self.loss_G_B + self.loss_cycle_A + self.loss_cycle_B + self.loss_idt_A + self.loss_idt_B
+        self.loss_G.backward()
+
+    def optimize_parameters(self):
+        """Calculate losses, gradients, and update network weights; called in every training iteration"""
+        # forward
+        self.forward()  # compute fake images and reconstruction images.
+        # G_A and G_B
+        self.set_requires_grad([self.netD_A, self.netD_B], False)  # Ds require no gradients when optimizing Gs
+        self.optimizer_G.zero_grad()  # set G_A and G_B's gradients to zero
+        self.backward_G()  # calculate gradients for G_A and G_B
+        self.optimizer_G.step()  # update G_A and G_B's weights
+        # D_A and D_B
+        self.set_requires_grad([self.netD_A, self.netD_B], True)
+        self.optimizer_D.zero_grad()  # set D_A and D_B's gradients to zero
+        self.backward_D_A()  # calculate gradients for D_A
+        self.backward_D_B()  # calculate graidents for D_B
+        self.optimizer_D.step()  # update D_A and D_B's weights

models/networks.py

@@ -0,0 +1,588 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from torch.optim import lr_scheduler
+
+
+###############################################################################
+# Helper Functions
+###############################################################################
+
+
+class Identity(nn.Module):
+    def forward(self, x):
+        return x
+
+
+def get_norm_layer(norm_type="instance"):
+    """Return a normalization layer
+
+    Parameters:
+        norm_type (str) -- the name of the normalization layer: batch | instance | none
+
+    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
+    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
+    """
+    if norm_type == "batch":
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
+    elif norm_type == "syncbatch":
+        norm_layer = functools.partial(nn.SyncBatchNorm, affine=True, track_running_stats=True)
+    elif norm_type == "instance":
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+    elif norm_type == "none":
+
+        def norm_layer(x):
+            return Identity()
+
+    else:
+        raise NotImplementedError("normalization layer [%s] is not found" % norm_type)
+    return norm_layer
+
+
+def get_scheduler(optimizer, opt):
+    """Return a learning rate scheduler
+
+    Parameters:
+        optimizer          -- the optimizer of the network
+        opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．
+                              opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine
+
+    For 'linear', we keep the same learning rate for the first <opt.n_epochs> epochs
+    and linearly decay the rate to zero over the next <opt.n_epochs_decay> epochs.
+    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
+    See https://pytorch.org/docs/stable/optim.html for more details.
+    """
+    if opt.lr_policy == "linear":
+
+        def lambda_rule(epoch):
+            lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1)
+            return lr_l
+
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+    elif opt.lr_policy == "step":
+        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
+    elif opt.lr_policy == "plateau":
+        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.2, threshold=0.01, patience=5)
+    elif opt.lr_policy == "cosine":
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)
+    else:
+        return NotImplementedError("learning rate policy [%s] is not implemented", opt.lr_policy)
+    return scheduler
+
+
+def init_weights(net, init_type="normal", init_gain=0.02):
+    """Initialize network weights.
+
+    Parameters:
+        net (network)   -- network to be initialized
+        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.
+
+    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
+    work better for some applications. Feel free to try yourself.
+    """
+
+    def init_func(m):  # define the initialization function
+        classname = m.__class__.__name__
+        if hasattr(m, "weight") and (classname.find("Conv") != -1 or classname.find("Linear") != -1):
+            if init_type == "normal":
+                init.normal_(m.weight.data, 0.0, init_gain)
+            elif init_type == "xavier":
+                init.xavier_normal_(m.weight.data, gain=init_gain)
+            elif init_type == "kaiming":
+                init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
+            elif init_type == "orthogonal":
+                init.orthogonal_(m.weight.data, gain=init_gain)
+            else:
+                raise NotImplementedError("initialization method [%s] is not implemented" % init_type)
+            if hasattr(m, "bias") and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find("BatchNorm2d") != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
+            init.normal_(m.weight.data, 1.0, init_gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print("initialize network with %s" % init_type)
+    net.apply(init_func)  # apply the initialization function <init_func>
+
+
+def init_net(net, init_type="normal", init_gain=0.02):
+    """Initialize a network: 1. register CPU/GPU device; 2. initialize the network weights
+    Parameters:
+        net (network)      -- the network to be initialized
+        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        gain (float)       -- scaling factor for normal, xavier and orthogonal.
+
+    Return an initialized network.
+    """
+    import os
+
+    if torch.cuda.is_available():
+        if "LOCAL_RANK" in os.environ:
+            local_rank = int(os.environ["LOCAL_RANK"])
+            net.to(local_rank)
+            print(f"Initialized with device cuda:{local_rank}")
+        else:
+            net.to(0)
+            print("Initialized with device cuda:0")
+    init_weights(net, init_type, init_gain=init_gain)
+    return net
+
+
+def define_G(input_nc, output_nc, ngf, netG, norm="batch", use_dropout=False, init_type="normal", init_gain=0.02):
+    """Create a generator
+
+    Parameters:
+        input_nc (int) -- the number of channels in input images
+        output_nc (int) -- the number of channels in output images
+        ngf (int) -- the number of filters in the last conv layer
+        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_128 | unet_256
+        norm (str) -- the name of normalization layers used in the network: batch | instance | none
+        use_dropout (bool) -- if use dropout layers.
+        init_type (str)    -- the name of our initialization method.
+        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
+
+    Returns a generator
+    """
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netG == "resnet_9blocks":
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
+    elif netG == "resnet_6blocks":
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
+    elif netG == "unet_128":
+        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == "unet_256":
+        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    else:
+        raise NotImplementedError("Generator model name [%s] is not recognized" % netG)
+    return net
+
+
+def define_D(input_nc, ndf, netD, n_layers_D=3, norm="batch", init_type="normal", init_gain=0.02):
+    """Create a discriminator
+
+    Parameters:
+        input_nc (int)     -- the number of channels in input images
+        ndf (int)          -- the number of filters in the first conv layer
+        netD (str)         -- the architecture's name: basic | n_layers | pixel
+        n_layers_D (int)   -- the number of conv layers in the discriminator; effective when netD=='n_layers'
+        norm (str)         -- the type of normalization layers used in the network.
+        init_type (str)    -- the name of the initialization method.
+        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
+
+    Returns a discriminator
+
+    Our current implementation provides three types of discriminators:
+        [basic]: 'PatchGAN' classifier described in the original pix2pix paper.
+        It can classify whether 70×70 overlapping patches are real or fake.
+        Such a patch-level discriminator architecture has fewer parameters
+        than a full-image discriminator and can work on arbitrarily-sized images
+        in a fully convolutional fashion.
+
+        [n_layers]: With this mode, you can specify the number of conv layers in the discriminator
+        with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)
+
+        [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.
+        It encourages greater color diversity but has no effect on spatial statistics.
+
+    The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.
+    """
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netD == "basic":  # default PatchGAN classifier
+        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)
+    elif netD == "n_layers":  # more options
+        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)
+    elif netD == "pixel":  # classify if each pixel is real or fake
+        net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
+    else:
+        raise NotImplementedError("Discriminator model name [%s] is not recognized" % netD)
+    return net
+
+
+##############################################################################
+# Classes
+##############################################################################
+class GANLoss(nn.Module):
+    """Define different GAN objectives.
+
+    The GANLoss class abstracts away the need to create the target label tensor
+    that has the same size as the input.
+    """
+
+    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
+        """Initialize the GANLoss class.
+
+        Parameters:
+            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
+            target_real_label (bool) - - label for a real image
+            target_fake_label (bool) - - label of a fake image
+
+        Note: Do not use sigmoid as the last layer of Discriminator.
+        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
+        """
+        super(GANLoss, self).__init__()
+        self.register_buffer("real_label", torch.tensor(target_real_label))
+        self.register_buffer("fake_label", torch.tensor(target_fake_label))
+        self.gan_mode = gan_mode
+        if gan_mode == "lsgan":
+            self.loss = nn.MSELoss()
+        elif gan_mode == "vanilla":
+            self.loss = nn.BCEWithLogitsLoss()
+        elif gan_mode in ["wgangp"]:
+            self.loss = None
+        else:
+            raise NotImplementedError("gan mode %s not implemented" % gan_mode)
+
+    def get_target_tensor(self, prediction, target_is_real):
+        """Create label tensors with the same size as the input.
+
+        Parameters:
+            prediction (tensor) - - tpyically the prediction from a discriminator
+            target_is_real (bool) - - if the ground truth label is for real images or fake images
+
+        Returns:
+            A label tensor filled with ground truth label, and with the size of the input
+        """
+
+        if target_is_real:
+            target_tensor = self.real_label
+        else:
+            target_tensor = self.fake_label
+        return target_tensor.expand_as(prediction)
+
+    def __call__(self, prediction, target_is_real):
+        """Calculate loss given Discriminator's output and grount truth labels.
+
+        Parameters:
+            prediction (tensor) - - tpyically the prediction output from a discriminator
+            target_is_real (bool) - - if the ground truth label is for real images or fake images
+
+        Returns:
+            the calculated loss.
+        """
+        if self.gan_mode in ["lsgan", "vanilla"]:
+            target_tensor = self.get_target_tensor(prediction, target_is_real)
+            loss = self.loss(prediction, target_tensor)
+        elif self.gan_mode == "wgangp":
+            if target_is_real:
+                loss = -prediction.mean()
+            else:
+                loss = prediction.mean()
+        return loss
+
+
+def cal_gradient_penalty(netD, real_data, fake_data, device, type="mixed", constant=1.0, lambda_gp=10.0):
+    """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028
+
+    Arguments:
+        netD (network)              -- discriminator network
+        real_data (tensor array)    -- real images
+        fake_data (tensor array)    -- generated images from the generator
+        device (str)                -- GPU / CPU
+        type (str)                  -- if we mix real and fake data or not [real | fake | mixed].
+        constant (float)            -- the constant used in formula ( ||gradient||_2 - constant)^2
+        lambda_gp (float)           -- weight for this loss
+
+    Returns the gradient penalty loss
+    """
+    if lambda_gp > 0.0:
+        if type == "real":  # either use real images, fake images, or a linear interpolation of two.
+            interpolatesv = real_data
+        elif type == "fake":
+            interpolatesv = fake_data
+        elif type == "mixed":
+            alpha = torch.rand(real_data.shape[0], 1, device=device)
+            alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)
+            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
+        else:
+            raise NotImplementedError(f"{type} not implemented")
+        interpolatesv.requires_grad_(True)
+        disc_interpolates = netD(interpolatesv)
+        gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv, grad_outputs=torch.ones(disc_interpolates.size()).to(device), create_graph=True, retain_graph=True, only_inputs=True)
+        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data
+        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp  # added eps
+        return gradient_penalty, gradients
+    else:
+        return 0.0, None
+
+
+class ResnetGenerator(nn.Module):
+    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
+
+    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
+    """
+
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type="reflect"):
+        """Construct a Resnet-based generator
+
+        Parameters:
+            input_nc (int)      -- the number of channels in input images
+            output_nc (int)     -- the number of channels in output images
+            ngf (int)           -- the number of filters in the last conv layer
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers
+            n_blocks (int)      -- the number of ResNet blocks
+            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
+        """
+        assert n_blocks >= 0
+        super(ResnetGenerator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias), norm_layer(ngf), nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):  # add downsampling layers
+            mult = 2**i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias), norm_layer(ngf * mult * 2), nn.ReLU(True)]
+
+        mult = 2**n_downsampling
+        for i in range(n_blocks):  # add ResNet blocks
+
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        for i in range(n_downsampling):  # add upsampling layers
+            mult = 2 ** (n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1, bias=use_bias), norm_layer(int(ngf * mult / 2)), nn.ReLU(True)]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class ResnetBlock(nn.Module):
+    """Define a Resnet block"""
+
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Initialize the Resnet block
+
+        A resnet block is a conv block with skip connections
+        We construct a conv block with build_conv_block function,
+        and implement skip connections in <forward> function.
+        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
+        """
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Construct a convolutional block.
+
+        Parameters:
+            dim (int)           -- the number of channels in the conv layer.
+            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+            use_bias (bool)     -- if the conv layer uses bias or not
+
+        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
+        """
+        conv_block = []
+        p = 0
+        if padding_type == "reflect":
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == "replicate":
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == "zero":
+            p = 1
+        else:
+            raise NotImplementedError("padding [%s] is not implemented" % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == "reflect":
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == "replicate":
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == "zero":
+            p = 1
+        else:
+            raise NotImplementedError("padding [%s] is not implemented" % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        """Forward function (with skip connections)"""
+        out = x + self.conv_block(x)  # add skip connections
+        return out
+
+
+class UnetGenerator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(UnetGenerator, self).__init__()
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer
+        for i in range(num_downs - 5):  # add intermediate layers with ngf * 8 filters
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        # gradually reduce the number of filters from ngf * 8 to ngf
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class UnetSkipConnectionBlock(nn.Module):
+    """Defines the Unet submodule with skip connection.
+    X -------------------identity----------------------
+    |-- downsampling -- |submodule| -- upsampling --|
+    """
+
+    def __init__(self, outer_nc, inner_nc, input_nc=None, submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet submodule with skip connections.
+
+        Parameters:
+            outer_nc (int) -- the number of filters in the outer conv layer
+            inner_nc (int) -- the number of filters in the inner conv layer
+            input_nc (int) -- the number of channels in input images/features
+            submodule (UnetSkipConnectionBlock) -- previously defined submodules
+            outermost (bool)    -- if this module is the outermost module
+            innermost (bool)    -- if this module is the innermost module
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+        """
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:  # add skip connections
+            return torch.cat([x, self.model(x)], 1)
+
+
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator"""
+
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+        """Construct a PatchGAN discriminator
+
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2**n, 8)
+            sequence += [nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias), norm_layer(ndf * nf_mult), nn.LeakyReLU(0.2, True)]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2**n_layers, 8)
+        sequence += [nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias), norm_layer(ndf * nf_mult), nn.LeakyReLU(0.2, True)]
+
+        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.model = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.model(input)
+
+
+class PixelDiscriminator(nn.Module):
+    """Defines a 1x1 PatchGAN discriminator (pixelGAN)"""
+
+    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d):
+        """Construct a 1x1 PatchGAN discriminator
+
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+        """
+        super(PixelDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        self.net = [
+            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
+            norm_layer(ndf * 2),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias),
+        ]
+
+        self.net = nn.Sequential(*self.net)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.net(input)

models/pix2pix_model.py

@@ -0,0 +1,127 @@
+import torch
+from .base_model import BaseModel
+from . import networks
+
+
+class Pix2PixModel(BaseModel):
+    """This class implements the pix2pix model, for learning a mapping from input images to output images given paired data.
+
+    The model training requires '--dataset_mode aligned' dataset.
+    By default, it uses a '--netG unet256' U-Net generator,
+    a '--netD basic' discriminator (PatchGAN),
+    and a '--gan_mode' vanilla GAN loss (the cross-entropy objective used in the orignal GAN paper).
+
+    pix2pix paper: https://arxiv.org/pdf/1611.07004.pdf
+    """
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train=True):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+
+        For pix2pix, we do not use image buffer
+        The training objective is: GAN Loss + lambda_L1 * ||G(A)-B||_1
+        By default, we use vanilla GAN loss, UNet with batchnorm, and aligned datasets.
+        """
+        # changing the default values to match the pix2pix paper (https://phillipi.github.io/pix2pix/)
+        parser.set_defaults(norm="batch", netG="unet_256", dataset_mode="aligned")
+        if is_train:
+            parser.set_defaults(pool_size=0, gan_mode="vanilla")
+            parser.add_argument("--lambda_L1", type=float, default=100.0, help="weight for L1 loss")
+
+        return parser
+
+    def __init__(self, opt):
+        """Initialize the pix2pix class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseModel.__init__(self, opt)
+        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
+        self.loss_names = ["G_GAN", "G_L1", "D_real", "D_fake"]
+        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
+        self.visual_names = ["real_A", "fake_B", "real_B"]
+        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
+        if self.isTrain:
+            self.model_names = ["G", "D"]
+        else:  # during test time, only load G
+            self.model_names = ["G"]
+        self.device = opt.device
+        # define networks (both generator and discriminator)
+        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain)
+
+        if self.isTrain:  # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
+            self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain)
+
+        if self.isTrain:
+            # define loss functions
+            self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)  # move to the device for custom loss
+            self.criterionL1 = torch.nn.L1Loss()
+            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
+            self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
+            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
+            self.optimizers.append(self.optimizer_G)
+            self.optimizers.append(self.optimizer_D)
+
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input (dict): include the data itself and its metadata information.
+
+        The option 'direction' can be used to swap images in domain A and domain B.
+        """
+        AtoB = self.opt.direction == "AtoB"
+        self.real_A = input["A" if AtoB else "B"].to(self.device)
+        self.real_B = input["B" if AtoB else "A"].to(self.device)
+        self.image_paths = input["A_paths" if AtoB else "B_paths"]
+
+    def forward(self):
+        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
+        self.fake_B = self.netG(self.real_A)  # G(A)
+
+    def backward_D(self):
+        """Calculate GAN loss for the discriminator"""
+        # Fake; stop backprop to the generator by detaching fake_B
+        fake_AB = torch.cat((self.real_A, self.fake_B), 1)  # we use conditional GANs; we need to feed both input and output to the discriminator
+        pred_fake = self.netD(fake_AB.detach())
+        self.loss_D_fake = self.criterionGAN(pred_fake, False)
+        # Real
+        real_AB = torch.cat((self.real_A, self.real_B), 1)
+        pred_real = self.netD(real_AB)
+        self.loss_D_real = self.criterionGAN(pred_real, True)
+        # combine loss and calculate gradients
+        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
+        self.loss_D.backward()
+
+    def backward_G(self):
+        """Calculate GAN and L1 loss for the generator"""
+        # First, G(A) should fake the discriminator
+        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
+        pred_fake = self.netD(fake_AB)
+        self.loss_G_GAN = self.criterionGAN(pred_fake, True)
+        # Second, G(A) = B
+        self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1
+        # combine loss and calculate gradients
+        self.loss_G = self.loss_G_GAN + self.loss_G_L1
+        self.loss_G.backward()
+
+    def optimize_parameters(self):
+        self.forward()  # compute fake images: G(A)
+        # update D
+        self.set_requires_grad(self.netD, True)  # enable backprop for D
+        self.optimizer_D.zero_grad()  # set D's gradients to zero
+        self.backward_D()  # calculate gradients for D
+        self.optimizer_D.step()  # update D's weights
+        # update G
+        self.set_requires_grad(self.netD, False)  # D requires no gradients when optimizing G
+        self.optimizer_G.zero_grad()  # set G's gradients to zero
+        self.backward_G()  # calculate graidents for G
+        self.optimizer_G.step()  # update G's weights

models/template_model.py

@@ -0,0 +1,100 @@
+"""Model class template
+
+This module provides a template for users to implement custom models.
+You can specify '--model template' to use this model.
+The class name should be consistent with both the filename and its model option.
+The filename should be <model>_dataset.py
+The class name should be <Model>Dataset.py
+It implements a simple image-to-image translation baseline based on regression loss.
+Given input-output pairs (data_A, data_B), it learns a network netG that can minimize the following L1 loss:
+    min_<netG> ||netG(data_A) - data_B||_1
+You need to implement the following functions:
+    <modify_commandline_options>:　Add model-specific options and rewrite default values for existing options.
+    <__init__>: Initialize this model class.
+    <set_input>: Unpack input data and perform data pre-processing.
+    <forward>: Run forward pass. This will be called by both <optimize_parameters> and <test>.
+    <optimize_parameters>: Update network weights; it will be called in every training iteration.
+"""
+
+import torch
+from .base_model import BaseModel
+from . import networks
+
+
+class TemplateModel(BaseModel):
+    @staticmethod
+    def modify_commandline_options(parser, is_train=True):
+        """Add new model-specific options and rewrite default values for existing options.
+
+        Parameters:
+            parser -- the option parser
+            is_train -- if it is training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        parser.set_defaults(dataset_mode="aligned")  # You can rewrite default values for this model. For example, this model usually uses aligned dataset as its dataset.
+        if is_train:
+            parser.add_argument("--lambda_regression", type=float, default=1.0, help="weight for the regression loss")  # You can define new arguments for this model.
+
+        return parser
+
+    def __init__(self, opt):
+        """Initialize this model class.
+
+        Parameters:
+            opt -- training/test options
+
+        A few things can be done here.
+        - (required) call the initialization function of BaseModel
+        - define loss function, visualization images, model names, and optimizers
+        """
+        BaseModel.__init__(self, opt)  # call the initialization method of BaseModel
+        # specify the training losses you want to print out. The program will call base_model.get_current_losses to plot the losses to the console and save them to the disk.
+        self.loss_names = ["G"]
+        # specify the images you want to save and display. The program will call base_model.get_current_visuals to save and display these images.
+        self.visual_names = ["data_A", "data_B", "output"]
+        # specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks to save and load networks.
+        # you can use opt.isTrain to specify different behaviors for training and test. For example, some networks will not be used during test, and you don't need to load them.
+        self.model_names = ["G"]
+        # define networks; you can use opt.isTrain to specify different behaviors for training and test.
+        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG)
+        if self.isTrain:  # only defined during training time
+            # define your loss functions. You can use losses provided by torch.nn such as torch.nn.L1Loss.
+            # We also provide a GANLoss class "networks.GANLoss". self.criterionGAN = networks.GANLoss().to(self.device)
+            self.criterionLoss = torch.nn.L1Loss()
+            # define and initialize optimizers. You can define one optimizer for each network.
+            # If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
+            self.optimizer = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
+            self.optimizers = [self.optimizer]
+
+        # Our program will automatically call <model.setup> to define schedulers, load networks, and print networks
+
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input: a dictionary that contains the data itself and its metadata information.
+        """
+        AtoB = self.opt.direction == "AtoB"  # use <direction> to swap data_A and data_B
+        self.data_A = input["A" if AtoB else "B"].to(self.device)  # get image data A
+        self.data_B = input["B" if AtoB else "A"].to(self.device)  # get image data B
+        self.image_paths = input["A_paths" if AtoB else "B_paths"]  # get image paths
+
+    def forward(self):
+        """Run forward pass. This will be called by both functions <optimize_parameters> and <test>."""
+        self.output = self.netG(self.data_A)  # generate output image given the input data_A
+
+    def backward(self):
+        """Calculate losses, gradients, and update network weights; called in every training iteration"""
+        # caculate the intermediate results if necessary; here self.output has been computed during function <forward>
+        # calculate loss given the input and intermediate results
+        self.loss_G = self.criterionLoss(self.output, self.data_B) * self.opt.lambda_regression
+        self.loss_G.backward()  # calculate gradients of network G w.r.t. loss_G
+
+    def optimize_parameters(self):
+        """Update network weights; it will be called in every training iteration."""
+        self.forward()  # first call forward to calculate intermediate results
+        self.optimizer.zero_grad()  # clear network G's existing gradients
+        self.backward()  # calculate gradients for network G
+        self.optimizer.step()  # update gradients for network G

models/test_model.py

@@ -0,0 +1,69 @@
+from .base_model import BaseModel
+from . import networks
+
+
+class TestModel(BaseModel):
+    """This TesteModel can be used to generate CycleGAN results for only one direction.
+    This model will automatically set '--dataset_mode single', which only loads the images from one collection.
+
+    See the test instruction for more details.
+    """
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train=True):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+
+        The model can only be used during test time. It requires '--dataset_mode single'.
+        You need to specify the network using the option '--model_suffix'.
+        """
+        assert not is_train, "TestModel cannot be used during training time"
+        parser.set_defaults(dataset_mode="single")
+        parser.add_argument("--model_suffix", type=str, default="", help="In checkpoints_dir, [epoch]_net_G[model_suffix].pth will be loaded as the generator.")
+
+        return parser
+
+    def __init__(self, opt):
+        """Initialize the pix2pix class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        assert not opt.isTrain
+        BaseModel.__init__(self, opt)
+        # specify the training losses you want to print out. The training/test scripts  will call <BaseModel.get_current_losses>
+        self.loss_names = []
+        # specify the images you want to save/display. The training/test scripts  will call <BaseModel.get_current_visuals>
+        self.visual_names = ["real", "fake"]
+        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
+        self.model_names = ["G" + opt.model_suffix]  # only generator is needed.
+        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain)
+
+        # assigns the model to self.netG_[suffix] so that it can be loaded
+        # please see <BaseModel.load_networks>
+        setattr(self, "netG" + opt.model_suffix, self.netG)  # store netG in self.
+
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input: a dictionary that contains the data itself and its metadata information.
+
+        We need to use 'single_dataset' dataset mode. It only load images from one domain.
+        """
+        self.real = input["A"].to(self.device)
+        self.image_paths = input["A_paths"]
+
+    def forward(self):
+        """Run forward pass."""
+        self.fake = self.netG(self.real)  # G(real)
+
+    def optimize_parameters(self):
+        """No optimization for test model."""
+        pass

options/__init__.py

@@ -0,0 +1 @@
+"""This package options includes option modules: training options, test options, and basic options (used in both training and test)."""

options/base_options.py

@@ -0,0 +1,127 @@
+import argparse
+from pathlib import Path
+from util import util
+import torch
+import models
+import data
+
+
+class BaseOptions:
+    """This class defines options used during both training and test time.
+
+    It also implements several helper functions such as parsing, printing, and saving the options.
+    It also gathers additional options defined in <modify_commandline_options> functions in both dataset class and model class.
+    """
+
+    def __init__(self):
+        """Reset the class; indicates the class hasn't been initailized"""
+        self.initialized = False
+
+    def initialize(self, parser):
+        """Define the common options that are used in both training and test."""
+        # basic parameters
+        parser.add_argument("--dataroot", required=True, help="path to images (should have subfolders trainA, trainB, valA, valB, etc)")
+        parser.add_argument("--name", type=str, default="experiment_name", help="name of the experiment. It decides where to store samples and models")
+        parser.add_argument("--checkpoints_dir", type=str, default="./checkpoints", help="models are saved here")
+        # model parameters
+        parser.add_argument("--model", type=str, default="cycle_gan", help="chooses which model to use. [cycle_gan | pix2pix | test | colorization]")
+        parser.add_argument("--input_nc", type=int, default=3, help="# of input image channels: 3 for RGB and 1 for grayscale")
+        parser.add_argument("--output_nc", type=int, default=3, help="# of output image channels: 3 for RGB and 1 for grayscale")
+        parser.add_argument("--ngf", type=int, default=64, help="# of gen filters in the last conv layer")
+        parser.add_argument("--ndf", type=int, default=64, help="# of discrim filters in the first conv layer")
+        parser.add_argument("--netD", type=str, default="basic", help="specify discriminator architecture [basic | n_layers | pixel]. The basic model is a 70x70 PatchGAN. n_layers allows you to specify the layers in the discriminator")
+        parser.add_argument("--netG", type=str, default="resnet_9blocks", help="specify generator architecture [resnet_9blocks | resnet_6blocks | unet_256 | unet_128]")
+        parser.add_argument("--n_layers_D", type=int, default=3, help="only used if netD==n_layers")
+        parser.add_argument("--norm", type=str, default="instance", help="instance normalization or batch normalization [instance | batch | none | syncbatch]")
+        parser.add_argument("--init_type", type=str, default="normal", help="network initialization [normal | xavier | kaiming | orthogonal]")
+        parser.add_argument("--init_gain", type=float, default=0.02, help="scaling factor for normal, xavier and orthogonal.")
+        parser.add_argument("--no_dropout", action="store_true", help="no dropout for the generator")
+        # dataset parameters
+        parser.add_argument("--dataset_mode", type=str, default="unaligned", help="chooses how datasets are loaded. [unaligned | aligned | single | colorization]")
+        parser.add_argument("--direction", type=str, default="AtoB", help="AtoB or BtoA")
+        parser.add_argument("--serial_batches", action="store_true", help="if true, takes images in order to make batches, otherwise takes them randomly")
+        parser.add_argument("--num_threads", default=4, type=int, help="# threads for loading data")
+        parser.add_argument("--batch_size", type=int, default=1, help="input batch size")
+        parser.add_argument("--load_size", type=int, default=286, help="scale images to this size")
+        parser.add_argument("--crop_size", type=int, default=256, help="then crop to this size")
+        parser.add_argument("--max_dataset_size", type=int, default=float("inf"), help="Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.")
+        parser.add_argument("--preprocess", type=str, default="resize_and_crop", help="scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]")
+        parser.add_argument("--no_flip", action="store_true", help="if specified, do not flip the images for data augmentation")
+        parser.add_argument("--display_winsize", type=int, default=256, help="display window size for both visdom and HTML")
+        # additional parameters
+        parser.add_argument("--epoch", type=str, default="latest", help="which epoch to load? set to latest to use latest cached model")
+        parser.add_argument("--load_iter", type=int, default="0", help="which iteration to load? if load_iter > 0, the code will load models by iter_[load_iter]; otherwise, the code will load models by [epoch]")
+        parser.add_argument("--verbose", action="store_true", help="if specified, print more debugging information")
+        parser.add_argument("--suffix", default="", type=str, help="customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}")
+        # wandb parameters
+        parser.add_argument("--use_wandb", action="store_true", help="if specified, then init wandb logging")
+        parser.add_argument("--wandb_project_name", type=str, default="CycleGAN-and-pix2pix", help="specify wandb project name")
+        self.initialized = True
+        return parser
+
+    def gather_options(self):
+        """Initialize our parser with basic options(only once).
+        Add additional model-specific and dataset-specific options.
+        These options are defined in the <modify_commandline_options> function
+        in model and dataset classes.
+        """
+        if not self.initialized:  # check if it has been initialized
+            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+            parser = self.initialize(parser)
+
+        # get the basic options
+        opt, _ = parser.parse_known_args()
+
+        # modify model-related parser options
+        model_name = opt.model
+        model_option_setter = models.get_option_setter(model_name)
+        parser = model_option_setter(parser, self.isTrain)
+        opt, _ = parser.parse_known_args()  # parse again with new defaults
+
+        # modify dataset-related parser options
+        dataset_name = opt.dataset_mode
+        dataset_option_setter = data.get_option_setter(dataset_name)
+        parser = dataset_option_setter(parser, self.isTrain)
+
+        # save and return the parser
+        self.parser = parser
+        return parser.parse_args()
+
+    def print_options(self, opt):
+        """Print and save options
+
+        It will print both current options and default values(if different).
+        It will save options into a text file / [checkpoints_dir] / opt.txt
+        """
+        message = ""
+        message += "----------------- Options ---------------\n"
+        for k, v in sorted(vars(opt).items()):
+            comment = ""
+            default = self.parser.get_default(k)
+            if v != default:
+                comment = "\t[default: %s]" % str(default)
+            message += "{:>25}: {:<30}{}\n".format(str(k), str(v), comment)
+        message += "----------------- End -------------------"
+        print(message)
+
+        # save to the disk
+        expr_dir = Path(opt.checkpoints_dir) / opt.name
+        util.mkdirs(expr_dir)
+        file_name = expr_dir / f"{opt.phase}_opt.txt"
+        with open(file_name, "wt") as opt_file:
+            opt_file.write(message)
+            opt_file.write("\n")
+
+    def parse(self):
+        """Parse our options, create checkpoints directory suffix, and set up gpu device."""
+        opt = self.gather_options()
+        opt.isTrain = self.isTrain  # train or test
+
+        # process opt.suffix
+        if opt.suffix:
+            suffix = ("_" + opt.suffix.format(**vars(opt))) if opt.suffix != "" else ""
+            opt.name = opt.name + suffix
+
+        self.print_options(opt)
+        self.opt = opt
+        return self.opt

options/test_options.py

@@ -0,0 +1,23 @@
+from .base_options import BaseOptions
+
+
+class TestOptions(BaseOptions):
+    """This class includes test options.
+
+    It also includes shared options defined in BaseOptions.
+    """
+
+    def initialize(self, parser):
+        parser = BaseOptions.initialize(self, parser)  # define shared options
+        parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
+        parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
+        parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
+        # Dropout and Batchnorm has different behavioir during training and test.
+        parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
+        parser.add_argument('--num_test', type=int, default=50, help='how many test images to run')
+        # rewrite devalue values
+        parser.set_defaults(model='test')
+        # To avoid cropping, the load_size should be the same as crop_size
+        parser.set_defaults(load_size=parser.get_default('crop_size'))
+        self.isTrain = False
+        return parser

options/train_options.py

@@ -0,0 +1,35 @@
+from .base_options import BaseOptions
+
+
+class TrainOptions(BaseOptions):
+    """This class includes training options.
+
+    It also includes shared options defined in BaseOptions.
+    """
+
+    def initialize(self, parser):
+        parser = BaseOptions.initialize(self, parser)
+        # HTML visualization parameters
+        parser.add_argument('--display_freq', type=int, default=400, help='frequency of showing training results on screen')
+        parser.add_argument('--update_html_freq', type=int, default=1000, help='frequency of saving training results to html')
+        parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')
+        parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')
+        # network saving and loading parameters
+        parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results')
+        parser.add_argument('--save_epoch_freq', type=int, default=5, help='frequency of saving checkpoints at the end of epochs')
+        parser.add_argument('--save_by_iter', action='store_true', help='whether saves model by iteration')
+        parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
+        parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...')
+        parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
+        # training parameters
+        parser.add_argument('--n_epochs', type=int, default=100, help='number of epochs with the initial learning rate')
+        parser.add_argument('--n_epochs_decay', type=int, default=100, help='number of epochs to linearly decay learning rate to zero')
+        parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
+        parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
+        parser.add_argument('--gan_mode', type=str, default='lsgan', help='the type of GAN objective. [vanilla| lsgan | wgangp]. vanilla GAN loss is the cross-entropy objective used in the original GAN paper.')
+        parser.add_argument('--pool_size', type=int, default=50, help='the size of image buffer that stores previously generated images')
+        parser.add_argument('--lr_policy', type=str, default='linear', help='learning rate policy. [linear | step | plateau | cosine]')
+        parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations')
+
+        self.isTrain = True
+        return parser

requirements.txt

@@ -0,0 +1,8 @@
+# requirements.txt
+
+gradio
+numpy
+opencv-python
+pillow
+torch
+torchvision

util/__init__.py

@@ -0,0 +1 @@
+"""This package includes a miscellaneous collection of useful helper functions."""

util/get_data.py

@@ -0,0 +1,107 @@
+from __future__ import print_function
+from pathlib import Path
+import tarfile
+import requests
+from warnings import warn
+from zipfile import ZipFile
+from bs4 import BeautifulSoup
+
+
+class GetData(object):
+    """A Python script for downloading CycleGAN or pix2pix datasets.
+
+    Parameters:
+        technique (str) -- One of: 'cyclegan' or 'pix2pix'.
+        verbose (bool)  -- If True, print additional information.
+
+    Examples:
+        >>> from util.get_data import GetData
+        >>> gd = GetData(technique='cyclegan')
+        >>> new_data_path = gd.get(save_path='./datasets')  # options will be displayed.
+
+    Alternatively, You can use bash scripts: 'scripts/download_pix2pix_model.sh'
+    and 'scripts/download_cyclegan_model.sh'.
+    """
+
+    def __init__(self, technique="cyclegan", verbose=True):
+        url_dict = {
+            "pix2pix": "http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/",
+            "cyclegan": "http://efrosgans.eecs.berkeley.edu/pix2pix/datasets",
+        }
+        self.url = url_dict.get(technique.lower())
+        self._verbose = verbose
+
+    def _print(self, text):
+        if self._verbose:
+            print(text)
+
+    @staticmethod
+    def _get_options(r):
+        soup = BeautifulSoup(r.text, "lxml")
+        options = [h.text for h in soup.find_all("a", href=True) if h.text.endswith((".zip", "tar.gz"))]
+        return options
+
+    def _present_options(self):
+        r = requests.get(self.url)
+        options = self._get_options(r)
+        print("Options:\n")
+        for i, o in enumerate(options):
+            print("{0}: {1}".format(i, o))
+        choice = input("\nPlease enter the number of the " "dataset above you wish to download:")
+        return options[int(choice)]
+
+    def _download_data(self, dataset_url, save_path):
+        save_path = Path(save_path)
+        if not save_path.is_dir():
+            save_path.mkdir(parents=True, exist_ok=True)
+
+        base = Path(dataset_url).name
+        temp_save_path = save_path / base
+
+        with open(temp_save_path, "wb") as f:
+            r = requests.get(dataset_url)
+            f.write(r.content)
+
+        if base.endswith(".tar.gz"):
+            obj = tarfile.open(temp_save_path)
+        elif base.endswith(".zip"):
+            obj = ZipFile(temp_save_path, "r")
+        else:
+            raise ValueError("Unknown File Type: {0}.".format(base))
+
+        self._print("Unpacking Data...")
+        obj.extractall(save_path)
+        obj.close()
+        temp_save_path.unlink()
+
+    def get(self, save_path, dataset=None):
+        """
+
+        Download a dataset.
+
+        Parameters:
+            save_path (str) -- A directory to save the data to.
+            dataset (str)   -- (optional). A specific dataset to download.
+                            Note: this must include the file extension.
+                            If None, options will be presented for you
+                            to choose from.
+
+        Returns:
+            save_path_full (str) -- the absolute path to the downloaded data.
+
+        """
+        if dataset is None:
+            selected_dataset = self._present_options()
+        else:
+            selected_dataset = dataset
+
+        save_path_full = Path(save_path) / selected_dataset.split(".")[0]
+
+        if save_path_full.is_dir():
+            warn(f"\n'{save_path_full}' already exists. Voiding Download.")
+        else:
+            self._print("Downloading Data...")
+            url = f"{self.url}/{selected_dataset}"
+            self._download_data(url, save_path=save_path)
+
+        return save_path_full.resolve()

util/html.py

@@ -0,0 +1,84 @@
+import dominate
+from dominate.tags import meta, h3, table, tr, td, p, a, img, br
+from pathlib import Path
+
+
+class HTML:
+    """This HTML class allows us to save images and write texts into a single HTML file.
+
+    It consists of functions such as <add_header> (add a text header to the HTML file),
+    <add_images> (add a row of images to the HTML file), and <save> (save the HTML to the disk).
+    It is based on Python library 'dominate', a Python library for creating and manipulating HTML documents using a DOM API.
+    """
+
+    def __init__(self, web_dir, title, refresh=0):
+        """Initialize the HTML classes
+
+        Parameters:
+            web_dir (str) -- a directory that stores the webpage. HTML file will be created at <web_dir>/index.html; images will be saved at <web_dir/images/
+            title (str)   -- the webpage name
+            refresh (int) -- how often the website refresh itself; if 0; no refreshing
+        """
+        self.title = title
+        self.web_dir = Path(web_dir)
+        self.img_dir = self.web_dir / "images"
+
+        self.web_dir.mkdir(parents=True, exist_ok=True)
+        self.img_dir.mkdir(parents=True, exist_ok=True)
+
+        self.doc = dominate.document(title=title)
+        if refresh > 0:
+            with self.doc.head:
+                meta(http_equiv="refresh", content=str(refresh))
+
+    def get_image_dir(self):
+        """Return the directory that stores images"""
+        return self.img_dir
+
+    def add_header(self, text):
+        """Insert a header to the HTML file
+
+        Parameters:
+            text (str) -- the header text
+        """
+        with self.doc:
+            h3(text)
+
+    def add_images(self, ims, txts, links, width=400):
+        """add images to the HTML file
+
+        Parameters:
+            ims (str list)   -- a list of image paths
+            txts (str list)  -- a list of image names shown on the website
+            links (str list) --  a list of hyperref links; when you click an image, it will redirect you to a new page
+        """
+        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
+        self.doc.add(self.t)
+        with self.t:
+            with tr():
+                for im, txt, link in zip(ims, txts, links):
+                    with td(style="word-wrap: break-word;", halign="center", valign="top"):
+                        with p():
+                            with a(href=Path("images") / link):
+                                img(style=f"width:{width}px", src=Path("images") / im)
+                            br()
+                            p(txt)
+
+    def save(self):
+        """save the current content to the HMTL file"""
+        html_file = self.web_dir / "index.html"
+        with open(html_file, "wt") as f:
+            f.write(self.doc.render())
+
+
+if __name__ == "__main__":  # we show an example usage here.
+    html = HTML("web/", "test_html")
+    html.add_header("hello world")
+
+    ims, txts, links = [], [], []
+    for n in range(4):
+        ims.append(f"image_{n}.png")
+        txts.append(f"text_{n}")
+        links.append(f"image_{n}.png")
+    html.add_images(ims, txts, links)
+    html.save()

util/image_pool.py

@@ -0,0 +1,54 @@
+import random
+import torch
+
+
+class ImagePool:
+    """This class implements an image buffer that stores previously generated images.
+
+    This buffer enables us to update discriminators using a history of generated images
+    rather than the ones produced by the latest generators.
+    """
+
+    def __init__(self, pool_size):
+        """Initialize the ImagePool class
+
+        Parameters:
+            pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created
+        """
+        self.pool_size = pool_size
+        if self.pool_size > 0:  # create an empty pool
+            self.num_imgs = 0
+            self.images = []
+
+    def query(self, images):
+        """Return an image from the pool.
+
+        Parameters:
+            images: the latest generated images from the generator
+
+        Returns images from the buffer.
+
+        By 50/100, the buffer will return input images.
+        By 50/100, the buffer will return images previously stored in the buffer,
+        and insert the current images to the buffer.
+        """
+        if self.pool_size == 0:  # if the buffer size is 0, do nothing
+            return images
+        return_images = []
+        for image in images:
+            image = torch.unsqueeze(image.data, 0)
+            if self.num_imgs < self.pool_size:  # if the buffer is not full; keep inserting current images to the buffer
+                self.num_imgs = self.num_imgs + 1
+                self.images.append(image)
+                return_images.append(image)
+            else:
+                p = random.uniform(0, 1)
+                if p > 0.5:  # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer
+                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive
+                    tmp = self.images[random_id].clone()
+                    self.images[random_id] = image
+                    return_images.append(tmp)
+                else:  # by another 50% chance, the buffer will return the current image
+                    return_images.append(image)
+        return_images = torch.cat(return_images, 0)  # collect all the images and return
+        return return_images

util/util.py

@@ -0,0 +1,130 @@
+"""This module contains simple helper functions"""
+
+from __future__ import print_function
+import torch
+import numpy as np
+from PIL import Image
+from pathlib import Path
+import torch.distributed as dist
+import os
+
+
+def tensor2im(input_image, imtype=np.uint8):
+    """ "Converts a Tensor array into a numpy image array.
+
+    Parameters:
+        input_image (tensor) --  the input image tensor array
+        imtype (type)        --  the desired type of the converted numpy array
+    """
+    if not isinstance(input_image, np.ndarray):
+        if isinstance(input_image, torch.Tensor):  # get the data from a variable
+            image_tensor = input_image.data
+        else:
+            return input_image
+        image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array
+        if image_numpy.shape[0] == 1:  # grayscale to RGB
+            image_numpy = np.tile(image_numpy, (3, 1, 1))
+        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
+    else:  # if it is a numpy array, do nothing
+        image_numpy = input_image
+    return image_numpy.astype(imtype)
+
+
+def diagnose_network(net, name="network"):
+    """Calculate and print the mean of average absolute(gradients)
+
+    Parameters:
+        net (torch network) -- Torch network
+        name (str) -- the name of the network
+    """
+    mean = 0.0
+    count = 0
+    for param in net.parameters():
+        if param.grad is not None:
+            mean += torch.mean(torch.abs(param.grad.data))
+            count += 1
+    if count > 0:
+        mean = mean / count
+    print(name)
+    print(mean)
+
+
+# initialize ddp
+def init_ddp():
+    # Initialize DDP if LOCAL_RANK is set
+    is_ddp = "WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 1
+
+    if is_ddp:
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl")
+        local_rank = int(os.environ["LOCAL_RANK"])
+        device = torch.device(f"cuda:{local_rank}")
+        torch.cuda.set_device(local_rank)
+    elif torch.cuda.is_available():
+        device = torch.device("cuda:0")
+        torch.cuda.set_device(0)
+    else:
+        device = torch.device("cpu")
+    print(f"Initialized with device {device}")
+    return device
+
+
+# cleanup ddp
+def cleanup_ddp():
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def save_image(image_numpy, image_path, aspect_ratio=1.0):
+    """Save a numpy image to the disk
+
+    Parameters:
+        image_numpy (numpy array) -- input numpy array
+        image_path (str)          -- the path of the image
+    """
+
+    image_pil = Image.fromarray(image_numpy)
+    h, w, _ = image_numpy.shape
+
+    if aspect_ratio > 1.0:
+        image_pil = image_pil.resize((h, int(w * aspect_ratio)), Image.BICUBIC)
+    if aspect_ratio < 1.0:
+        image_pil = image_pil.resize((int(h / aspect_ratio), w), Image.BICUBIC)
+    image_pil.save(image_path)
+
+
+def print_numpy(x, val=True, shp=False):
+    """Print the mean, min, max, median, std, and size of a numpy array
+
+    Parameters:
+        val (bool) -- if print the values of the numpy array
+        shp (bool) -- if print the shape of the numpy array
+    """
+    x = x.astype(np.float64)
+    if shp:
+        print("shape,", x.shape)
+    if val:
+        x = x.flatten()
+        print("mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f" % (np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))
+
+
+def mkdirs(paths):
+    """create empty directories if they don't exist
+
+    Parameters:
+        paths (str list) -- a list of directory paths
+    """
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+
+def mkdir(path):
+    """create a single empty directory if it didn't exist
+
+    Parameters:
+        path (str) -- a single directory path
+    """
+    Path(path).mkdir(parents=True, exist_ok=True)