'Refactoring' (#35)

* 'Refactoring'

* Refactoring

* One more

* Two more

Author: SkillzZA

Diff for: README.md

@@ -17,5 +17,7 @@ Seminar project Unsupervised Image-to-Image translation using GANs
 ``` 
 ## Implementations
 
+### pix2pix
+### CycleGAN
 ### UNIT
-### pix2pix
+### TUNIT

Diff for: api/app/api.py

@@ -8,8 +8,8 @@
 api_blueprint = Blueprint("api", __name__)
 
 
-@api_blueprint.route("/coloring", methods=["POST"])
-def process_image_route():
+@api_blueprint.route("/<style>/coloring", methods=["POST"])
+def process_image_route(style):
     # if "image" not in request.files:
     #     return jsonify({"error": "No image provided"}), 400
 
@@ -18,10 +18,9 @@ def process_image_route():
     base64_image_data = request.json["image"]
     image_data = base64.b64decode(base64_image_data)
 
-    processed_image = process_image(io.BytesIO(image_data))
+    processed_image = process_image(io.BytesIO(image_data), style)
 
     # Convert the processed image to bytes
-
     processed_image.save(img_byte_array, format="JPEG")
     img_byte_array.seek(0)
 

Diff for: api/app/image_processing.py

@@ -3,12 +3,25 @@
 
 from img2img.models.pix2pix.predictor import Pix2PixPredictor
 
-predictor = Pix2PixPredictor(
-    model_path="/home/omarabdelgawad/my_workspace/projects/github_repos/image2image/out/saved_models/anime_training/gen.pth.tar"
-)
+# from img2img.models.cyclegan.predictor import CycleGANPredictor
 
+# Initialize predictors
+anime_predictor = Pix2PixPredictor(model_path="./out/saved_models/anime_training/gen.pth.tar")
+# monet_predictor = CycleGANPredictor(model_path="./out/saved_models/monet_training/gen.pth.tar")
+# yukiyoe_predictor = CycleGANPredictor(model_path="./out/saved_models/yukiyoe_training/gen.pth.tar")
+# vangogh_predictor = CycleGANPredictor(model_path="./out/saved_models/vangogh_training/gen.pth.tar")
 
-def process_image(image_file):
+
+predictors = {
+    "anime": anime_predictor,
+
+    # "monet": monet_predictor,
+    # "yukiyoe": yukiyoe_predictor,
+    # "vangogh": vangogh_predictor,
+}
+
+
+def process_image(image_file, style):
     # Open the image
     image = Image.open(image_file)
 
@@ -18,8 +31,8 @@ def process_image(image_file):
     # Ensure the array shape is correct
     assert processed_image.shape[2] == 3
 
-    # Process the image using the Pix2Pix model
-    processed_image = predictor(processed_image)
+    # Process the image using the appropriate model
+    processed_image = predictors[style](processed_image)
 
     # Convert the processed image array back to PIL Image
     processed_image = Image.fromarray(processed_image)

Diff for: pyproject.toml

@@ -6,6 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "img2img"
 dynamic = ["version"]
 dependencies = [
+    "enums",
     "requests",
     "torch >= 2.3.0 , < 3",
     "torchaudio >= 2.3.0, < 3",
@@ -18,14 +19,18 @@ requires-python = ">=3.10"
 authors = [
   {name = "Omar Abdelgawad", email = "[email protected]"},
   {name = "Eyad Hussein", email = "[email protected]"},
+  {name = "Ali Elsawy", email = "[email protected]"}
 ]
 description = "image to image translation using GANs"
 readme = "README.md"
 license = {file = "LICENSE"}
 keywords = ["image-to-image-translation", "GAN", "vision", "deep-learning"]
 classifiers = [
   "Programming Language :: Python :: 3",
-  "Programming Language :: Python :: 3.10"
+  "Programming Language :: Python :: 3.10",
+  "Operating System :: OS Independent",
+  "License :: OSI Approved :: MIT License",
+  "Topic :: Scientific/Research :: Artificial Intelligence"
 ]
 
 [project.optional-dependencies]
@@ -35,8 +40,10 @@ dev = [
     "pytest",
     "pytest-cov",
     "tox",
-    "flask",        # api dependency
-    "requests",     # api dependency
+]
+api = [
+    "flask",
+    "requests",
 ]
 
 [tool.setuptools.package-data]

Diff for: scripts/evaluate.py

@@ -1,10 +1,10 @@
 import torch
 from torch.utils.data import DataLoader
 
-from pix2pix import cfg
-from pix2pix.dataset import create_dataset
-from pix2pix.generator import Generator
-from pix2pix.utils import evaluate_val_set
+from img2img.cfg import pix2pix as cfg
+from img2img.data.pix2pix import create_dataset
+from img2img.models.pix2pix.generator import Generator
+from img2img.utils.pix2pix import evaluate_val_set
 
 
 def main() -> int:

Diff for: src/img2img/cfg/__init__.py

@@ -6,48 +6,28 @@
 
 from img2img.cli import get_main_parser
 
+from .transform import get_transforms
+
 from .enums import ActivationType, DatasetType, NormalizationType, PaddingType
 
 # TODO: Add logger instead of all the print statements.
 
+
 args = get_main_parser()
 DEVICE = "cuda" if cuda.is_available() else "cpu"
 LEARNING_RATE = args.rate
-BETA_OPTIM = (0.5, 0.999)
 BATCH_SIZE = args.batch_size
 NUM_WORKERS = args.num_workers
 IMAGE_SIZE = args.image_size
-CHANNELS_IMG = 3
-L_1_LAMBDA = 100
 NORM_MEAN = 0.5
 NORM_STD = 0.5
 CHECKPOINT_PERIOD = 5
 NUM_EPOCHS = args.num_epochs
 LOAD_MODEL = args.load_model
 SAVE_MODEL = args.save_model
-CHOSEN_DATASET = DatasetType.ANIME_DATASET
-TRAIN_DATASET_PATH = CHOSEN_DATASET.value / "train"
-VAL_DATASET_PATH = CHOSEN_DATASET.value / "val"
-OUT_PATH = Path("./out")
-NUM_IMAGES_DATASET = args.num_images_dataset
-VAL_BATCH_SIZE = args.val_batch_size
-
-# tunit config
-CHANNELS_MULTIPLIER = 64
-K = args.cluster_number
-
-# unit config
-WEIGHT_DECAY = 0.0001
-LR_POLICY = "step"
-STEP_SIZE = 100000
-GAMMA = 0.5
-INIT = "kaiming"
-GAN_WEIGHT = 1
-RECONSTRUCTION_X_WEIGHT = 10
-RECONSTRUCTION_H_WEIGHT = 0
-RECONSTRUCTION_KL_WEIGHT = 0.01
-RECONSTRUCTION_X_CYC_WEIGHT = 10
-RECONSTRUCTION_KL_CYC_WEIGHT = 0.01
+both_transform, transform_only_input, transform_only_mask, transforms, prediction_transform = get_transforms(IMAGE_SIZE,
+                                                                                                             NORM_MEAN,
+                                                                                                             NORM_STD)
 
 
 class GEN_HYPERPARAMS:
@@ -71,80 +51,3 @@ class DIS_HYPERPARAMS:
     GAN_TYPE = "lsgan"
     NUM_SCALES = 3
     PAD_TYPE = PaddingType.REFLECT
-
-
-import albumentations as A
-from albumentations.pytorch import ToTensorV2
-
-both_transform = A.Compose(
-    [
-        A.Resize(width=IMAGE_SIZE, height=IMAGE_SIZE),
-        A.HorizontalFlip(p=0.5),
-    ],
-    additional_targets={"image0": "image"},
-)
-
-transform_only_input = A.Compose(
-    [
-        A.ColorJitter(p=0.1),
-        # TODO: calculate mean and std for the dataset instead of using these values.
-        A.Normalize(
-            mean=[NORM_MEAN, NORM_MEAN, NORM_MEAN],
-            std=[NORM_STD, NORM_STD, NORM_STD],
-            max_pixel_value=255.0,
-        ),
-        ToTensorV2(),
-    ]
-)
-
-transform_only_mask = A.Compose(
-    [
-        A.Normalize(
-            mean=[NORM_MEAN, NORM_MEAN, NORM_MEAN],
-            std=[NORM_STD, NORM_STD, NORM_STD],
-            max_pixel_value=255.0,
-        ),
-        ToTensorV2(),
-    ]
-)
-
-
-"""unit transforms
-
-# TODO: understand the augmentations below and improve them (maybe add more augmentations).
-both_transform = A.Compose(
-    [
-        A.Resize(width=IMAGE_SIZE, height=IMAGE_SIZE),
-        A.HorizontalFlip(p=0.5),
-    ],
-    additional_targets={"image0": "image"},
-)
-
-# this is equivalent to first domain transform
-transform_only_input = A.Compose(
-    [
-        A.ColorJitter(p=0.1),
-        A.Normalize(
-            mean=[0.5, 0.5, 0.5],
-            std=[0.5, 0.5, 0.5],
-            max_pixel_value=255.0,
-        ),
-        A.HorizontalFlip(p=0.5),
-        ToTensorV2(),
-    ]
-)
-
-# this is equivalent to second domain transform
-transform_only_mask = A.Compose(
-    [
-        A.Normalize(
-            mean=[0.5, 0.5, 0.5],
-            std=[0.5, 0.5, 0.5],
-            max_pixel_value=255.0,
-        ),
-        ToTensorV2(),
-    ]
-)
-"""
-
-# TODO: make transforms differ from model to another

Diff for: src/img2img/cfg/cyclegan.py

@@ -0,0 +1,12 @@
+from . import *
+
+LAMBDA_IDENTITY = 0.0
+LAMBDA_CYCLE = 10
+BATCH_SIZE = 2
+CHECKPOINT_GEN_H = "genh.pth.tar"
+CHECKPOINT_GEN_Z = "genz.pth.tar"
+CHECKPOINT_CRITIC_H = "critich.pth.tar"
+CHECKPOINT_CRITIC_Z = "criticz.pth.tar"
+CHOSEN_DATASET = DatasetType.VANGOGH2PHOTO
+TRAIN_DATASET_PATH = CHOSEN_DATASET.value / "train"
+VAL_DATASET_PATH = CHOSEN_DATASET.value / "val"

Diff for: src/img2img/cfg/enums.py

@@ -7,16 +7,13 @@
 class DatasetType(Enum):
     """Enum for the dataset type."""
 
-    ANIME_DATASET = Path(
-        "/media/omarabdelgawad/New Volume/Datasets/image_coloring/anime_dataset/"
-    )
-    NATURAL_VIEW_DATASET = Path(
-        "/media/omarabdelgawad/New Volume/Datasets/image_coloring/natural_view/"
-    )
-    EDGES2SHOES = Path(
-        "/media/omarabdelgawad/New Volume/Datasets/image_coloring/edges2shoes/"
-    )
-    AFHQ_CATS_DATASET = Path("/home/eyad/Downloads/afhq/")
+    ANIME_DATASET = Path("/media/omarabdelgawad/New Volume/Datasets/Anime_Dataset")
+    NATURAL_VIEW_DATASET = Path("/media/omarabdelgawad/New Volume/Datasets/Natural_View")
+    EDGES2SHOES = Path("/media/omarabdelgawad/New Volume/Datasets/Edges2Shoes")
+    AFHQ_CATS_DATASET = Path("/media/omarabdelgawad/New Volume/Datasets/AFHQ_Cats")
+    VANGOGH2PHOTO = Path("/media/omarabdelgawad/New Volume/Datasets/vangogh2photo")
+    yukiyoe = Path("/media/omarabdelgawad/New Volume/Datasets/yukiyoe")
+    monet = Path("/media/omarabdelgawad/New Volume/Datasets/monet")
 
 
 class PaddingType(Enum):

Diff for: src/img2img/cfg/pix2pix.py

@@ -0,0 +1,38 @@
+from . import *
+
+BETA_OPTIM = (0.5, 0.999)
+NUM_WORKERS = args.num_workers
+CHANNELS_IMG = 3
+L_1_LAMBDA = 100
+NORM_MEAN = 0.5
+NORM_STD = 0.5
+CHOSEN_DATASET = DatasetType.ANIME_DATASET
+TRAIN_DATASET_PATH = CHOSEN_DATASET.value / "train"
+VAL_DATASET_PATH = CHOSEN_DATASET.value / "val"
+OUT_PATH = Path("./out")
+NUM_IMAGES_DATASET = args.num_images_dataset
+VAL_BATCH_SIZE = args.val_batch_size
+
+
+# Hyperparameters for the generator and discriminator
+class GEN_HYPERPARAMS:
+    """Hyperparameters for the generator."""
+
+    DIM = 64
+    NORM = NormalizationType.INSTANCE
+    ACTIV = ActivationType.RELU
+    N_DOWNSAMPLE = 2
+    N_RES = 4
+    PAD_TYPE = PaddingType.REFLECT
+
+
+class DIS_HYPERPARAMS:
+    """Hyperparameters for the discriminator."""
+
+    DIM = 64
+    NORM = NormalizationType.NONE
+    ACTIV = ActivationType.LEAKY_RELU
+    N_LAYER = 4
+    GAN_TYPE = "lsgan"
+    NUM_SCALES = 3
+    PAD_TYPE = PaddingType.REFLECT

Diff for: src/img2img/cfg/transform.py

@@ -0,0 +1,55 @@
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+def get_transforms(IMAGE_SIZE, NORM_MEAN, NORM_STD):
+    both_transform = A.Compose(
+        [
+            A.Resize(width=IMAGE_SIZE, height=IMAGE_SIZE),
+            A.HorizontalFlip(p=0.5),
+        ],
+        additional_targets={"image0": "image"},
+    )
+
+    transform_only_input = A.Compose(
+        [
+            A.ColorJitter(p=0.1),
+            A.Normalize(
+                mean=[NORM_MEAN, NORM_MEAN, NORM_MEAN],
+                std=[NORM_STD, NORM_STD, NORM_STD],
+                max_pixel_value=255.0,
+            ),
+            ToTensorV2(),
+        ]
+    )
+
+    transform_only_mask = A.Compose(
+        [
+            A.Normalize(
+                mean=[NORM_MEAN, NORM_MEAN, NORM_MEAN],
+                std=[NORM_STD, NORM_STD, NORM_STD],
+                max_pixel_value=255.0,
+            ),
+            ToTensorV2(),
+        ]
+    )
+
+    transforms = A.Compose(
+        [
+            A.Resize(width=256, height=256),
+            A.HorizontalFlip(p=0.5),
+            A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),
+            ToTensorV2(),
+        ],
+        additional_targets={"image0": "image"},
+    )
+
+    prediction_transform = A.Compose(
+        [
+            A.Resize(width=256, height=256),
+            A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),
+            ToTensorV2(),
+        ],
+    )
+
+    return both_transform, transform_only_input, transform_only_mask, transforms, prediction_transform