Refactor/clean up code

.gitignore

@@ -1,3 +1,7 @@
 __pycache__/
 .vscode/
 .DS_Store
+.idea/
+checkpoints/
+results/
+dataset/

.pre-commit-config.yaml

@@ -0,0 +1,9 @@
+repos:
+  - repo: https://github.com/pre-commit/mirrors-isort
+    rev: v5.10.1
+    hooks:
+      - id: isort
+  - repo: https://github.com/psf/black
+    rev: 25.1.0
+    hooks:
+      - id: black

README.md

@@ -95,12 +95,17 @@ We provide some [pretrained models](https://drive.google.com/open?id=1zVhM2VQTir
 
 ```
 python generate.py --dataroot /path/to/dataset --name your_exp_name --checkpoints_dir /path/to/checkpoints --results_dir /path/to/result --eval
+or
+(my version): python generate.py --name your_exp_name --eval 
 ```
 
 - Make gifs from generated frames.
 
 ```
 python img2gif.py --exp_names exp_name1,exp_name1,.. --dataroot /path/to/dataset
+or
+python img2gif.py --exp_names happy --dataroot ./dataset/test_image/
+
 ```
 
 - You will see the results in the specified `results_dir`.

data/affineGAN_dataset.py

@@ -17,9 +17,7 @@ def modify_commandline_options(parser, is_train):
     def pre_process_img(self, path, convertRGB, w_offset, h_offset, flip):
         if not os.path.exists(path):
             if convertRGB:
-                return np.zeros(
-                    (self.opt.fineSize, self.opt.fineSize, 3), dtype=np.float32
-                )
+                return np.zeros((self.opt.fineSize, self.opt.fineSize, 3), dtype=np.float32)
             else:
                 return np.zeros((self.opt.fineSize, self.opt.fineSize), dtype=np.float32)
 
@@ -93,9 +91,7 @@ def __getitem__(self, index):
             B_list = []
             np.random.seed()
             img_sample = range(1, len(img_names))
-            img_sample = np.random.choice(
-                img_sample, self.opt.train_imagenum, replace=True
-            )
+            img_sample = np.random.choice(img_sample, self.opt.train_imagenum, replace=True)
             for img_idx in range(self.opt.train_imagenum):
                 sample_image_idx = img_sample[img_idx]
 
@@ -104,12 +100,8 @@ def __getitem__(self, index):
 
                 B_list.append(B)
                 if not self.opt.no_patch:
-                    B_name = os.path.join(
-                        self.AB_patch[index], img_names[sample_image_idx]
-                    )
-                    B_patch = self.pre_process_img(
-                        B_name, False, w_offset, h_offset, flip
-                    )
+                    B_name = os.path.join(self.AB_patch[index], img_names[sample_image_idx])
+                    B_patch = self.pre_process_img(B_name, False, w_offset, h_offset, flip)
                     B_patch_list.append(B_patch)
 
             ret_dict["B_list"] = B_list

img2gif.py

@@ -1,7 +1,9 @@
-import imageio
+import argparse
 import os
+
+import imageio.v2 as imageio
 import numpy as np
-import argparse
+from PIL import Image
 
 base_output_dir = "gifs"
 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
@@ -11,15 +13,21 @@
 parser.add_argument("--phase", default="test", help="which phase")
 parser.add_argument("--dataroot", required=True, help="path to images")
 parser.add_argument("--interval", default=0.05, type=float, help="time interval")
+parser.add_argument("--duration", default=100, type=int, help="duration per frame in ms")
 opt, _ = parser.parse_known_args()
 
 exp_list = opt.exp_names.split(",")
 
 for exp_name in exp_list:
     current_output_dir = os.path.join(opt.results_dir, base_output_dir, exp_name)
-    if not os.path.exists(current_output_dir):
-        os.makedirs(current_output_dir)
-    for sample_idx in os.listdir(os.path.join(opt.dataroot, "test", "img")):
+    os.makedirs(current_output_dir, exist_ok=True)
+
+    test_img_dir = os.path.join(opt.dataroot, "test", "img")
+    if not os.path.exists(test_img_dir):
+        print(f"Directory not found: {test_img_dir}")
+        continue
+
+    for sample_idx in os.listdir(test_img_dir):
         filenames = []
         images = []
         num_str = sample_idx
@@ -28,17 +36,48 @@
             c_name = os.path.join(
                 opt.results_dir,
                 exp_name,
-                "%s_%s" % (opt.phase, opt.epoch),
+                f"{opt.phase}_{opt.epoch}",
                 "images",
-                "%s_fake_B_list%d.png" % (num_str, i),
+                f"{num_str}_fake_B_list{i}.png",
             )
 
-            filenames.append(c_name)
+            if os.path.exists(c_name):
+                filenames.append(c_name)
+            else:
+                print(f"File not found: {c_name}")
+                continue
+
+        if not filenames:
+            print(f"No files found for sample {sample_idx}")
+            continue
+
         for filename in filenames:
-            a = np.array(imageio.imread(filename))
-            images.append(a)
+            try:
+                img = imageio.imread(filename)
+                if isinstance(img, np.ndarray):
+                    images.append(img)
+                else:
+                    print(f"Unexpected image type in {filename}: {type(img)}")
+            except Exception as e:
+                print(f"Error reading {filename}: {e}")
+                continue
+
+        if not images:
+            print(f"No valid images for sample {sample_idx}")
+            continue
 
-        output_dir = os.path.join(
-            current_output_dir, sample_idx + "_" + str(opt.epoch) + ".gif"
-        )
-        imageio.mimsave(output_dir, images)
+        output_path = os.path.join(current_output_dir, f"{sample_idx}_{opt.epoch}.gif")
+
+        try:
+            pil_images = [Image.fromarray(img) for img in images]
+            pil_images[0].save(
+                output_path,
+                save_all=True,
+                append_images=pil_images[1:],
+                duration=opt.duration,
+                loop=0,
+                optimize=True,
+            )
+            print(f"Successfully saved with PIL: {output_path}")
+        except Exception as e:
+            print(f"Error saving GIF with PIL {output_path}: {e}")