Merge pull request #7 from georgw777/develop

v0.2.7

.github/workflows/build.yml

@@ -25,7 +25,7 @@ jobs:
         uses: actions/cache@v2
         with:
           path: .venv
-          key: venv-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('**/poetry.lock') }}
+          key: venv-${{ runner.os }}-${{ matrix.python-version }}-v2-${{ hashFiles('**/poetry.lock') }}
       - name: Install dependencies
         run: poetry install
         if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'

Dockerfile

@@ -25,6 +25,9 @@ COPY ./pyproject.toml ./poetry.lock* ./
 RUN poetry install --no-root
 ENV PYTHONPATH "/chess:${PYTHONPATH}"
 
+# Tensorboard fix
+RUN poetry run python -m pip install wheel "setuptools>=41.0.0"
+
 # Setup data mount
 RUN mkdir -p /data
 ENV DATA_DIR /data

chesscog/__version__.py

@@ -1 +1 @@
-__version__ = "0.2.6"
+__version__ = "0.2.7"

chesscog/core/dataset/__init__.py

@@ -0,0 +1,3 @@
+from .dataset import color_name, piece_name, name_to_piece, build_transforms, build_dataset, build_data_loader
+from .transforms import unnormalize, build_transforms
+from .datasets import Datasets

chesscog/core/dataset.py → chesscog/core/dataset/dataset.py

@@ -1,19 +1,17 @@
 import torch
 import torchvision
-from torchvision import transforms as T
 import typing
 import logging
 from enum import Enum
 import numpy as np
 import chess
 from recap import URI, CfgNode as CN
 
+from .transforms import build_transforms
+from .datasets import Datasets
 
 logger = logging.getLogger(__name__)
 
-_MEAN = np.array([0.485, 0.456, 0.406])
-_STD = np.array([0.229, 0.224, 0.225])
-
 
 def color_name(color: chess.Color):
     return {chess.WHITE: "white",
@@ -31,31 +29,6 @@ def name_to_piece(name: str) -> chess.Piece:
     return chess.Piece(piece_type, color)
 
 
-class Datasets(Enum):
-    TRAIN = "train"
-    VAL = "val"
-    TEST = "test"
-
-
-def build_transforms(cfg: CN, mode: Datasets) -> typing.Callable:
-    transforms = cfg.DATASET.TRANSFORMS
-    t = []
-    if transforms.CENTER_CROP:
-        t.append(T.CenterCrop(transforms.CENTER_CROP))
-    if mode == Datasets.TRAIN and transforms.RANDOM_HORIZONTAL_FLIP:
-        t.append(T.RandomHorizontalFlip(transforms.RANDOM_HORIZONTAL_FLIP))
-    if transforms.RESIZE:
-        t.append(T.Resize(tuple(reversed(transforms.RESIZE))))
-    t.extend([T.ToTensor(),
-              T.Normalize(mean=_MEAN, std=_STD)])
-    return T.Compose(t)
-
-
-def unnormalize(x: typing.Union[torch.Tensor, np.ndarray]) -> typing.Union[torch.Tensor, np.ndarray]:
-    # x must be of the form ([..., W, H, 3])
-    return x * _STD + _MEAN
-
-
 def build_dataset(cfg: CN, mode: Datasets) -> torch.utils.data.Dataset:
     transform = build_transforms(cfg, mode)
     dataset = torchvision.datasets.ImageFolder(root=URI(cfg.DATASET.PATH) / mode.value,

chesscog/core/dataset/datasets.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class Datasets(Enum):
+    TRAIN = "train"
+    VAL = "val"
+    TEST = "test"

chesscog/core/dataset/transforms.py

@@ -0,0 +1,127 @@
+from recap import CfgNode as CN
+import typing
+from torchvision import transforms as T
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+from abc import ABC
+
+from .datasets import Datasets
+
+_MEAN = np.array([0.485, 0.456, 0.406])
+_STD = np.array([0.229, 0.224, 0.225])
+
+
+def unnormalize(x: typing.Union[torch.Tensor, np.ndarray]) -> typing.Union[torch.Tensor, np.ndarray]:
+    # x must be of the form ([..., W, H, 3])
+    return x * _STD + _MEAN
+
+
+class Shear:
+    """Custom shear transform that keeps the bottom of the image invariant because for piece classification, we only want to "tilt" the top of the image.
+    """
+
+    def __init__(self, amount: typing.Union[tuple, float, int, None]):
+        self.amount = amount
+
+    @classmethod
+    def _shear(cls, img: Image, amount: float) -> Image:
+        img = ImageOps.flip(img)
+        img = img.transform(img.size, Image.AFFINE,
+                            (1, -amount, 0, 0, 1, 0))
+        img = ImageOps.flip(img)
+        return img
+
+    def __call__(self, img: Image) -> Image:
+        if not self.amount:
+            return img
+        if isinstance(self.amount, (tuple, list)):
+            min_val, max_val = sorted(self.amount)
+        else:
+            min_val = max_val = self.amount
+
+        amount = np.random.uniform(low=min_val, high=max_val)
+        return self._shear(img, amount)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.amount})"
+
+
+class _HVTransform(ABC):
+    """Base class for transforms parameterized by horizontal and vertical values.
+    """
+
+    def __init__(self, horizontal: typing.Union[float, tuple, None], vertical: typing.Union[float, tuple, None]):
+        self.horizontal = self._get_tuple(horizontal)
+        self.vertical = self._get_tuple(vertical)
+
+    _default_value = None
+
+    @classmethod
+    def _get_tuple(cls, value: typing.Union[float, tuple, None]) -> tuple:
+        if value is None:
+            return cls._default_value, cls._default_value
+        elif isinstance(value, (tuple, list)):
+            return tuple(map(float, value))
+        elif isinstance(value, (float, int)):
+            return tuple(map(float, (value, value)))
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.horizontal}, {self.vertical})"
+
+
+class Scale(_HVTransform):
+    """Custom scaling transform where the horizontal and vertical scales can independently be specified.
+
+    The center of scaling is the bottom left of the image (this makes particular sense for the piece classifier).
+    """
+
+    _default_value = 1.
+
+    def __call__(self, img: Image) -> Image:
+        w, h = img.size
+        w_scale = np.random.uniform(*self.horizontal)
+        h_scale = np.random.uniform(*self.vertical)
+        w_, h_ = map(int, (w*w_scale, h*h_scale))
+        img = img.resize((w_, h_))
+        img = img.transform((w, h), Image.AFFINE, (1, 0, 0, 0, 1, h_-h))
+        return img
+
+
+class Translate(_HVTransform):
+    """Custom translation transform for convenience.
+    """
+
+    _default_value = 0.
+
+    def __call__(self, img: Image) -> Image:
+        w, h = img.size
+        w_translate = np.random.uniform(*self.horizontal)
+        h_translate = np.random.uniform(*self.vertical)
+        w_, h_ = map(int, (w*w_translate, h*h_translate))
+        img = img.transform((w, h), Image.AFFINE, (1, 0, -w_, 0, 1, h_))
+        return img
+
+
+def build_transforms(cfg: CN, mode: Datasets) -> typing.Callable:
+    transforms = cfg.DATASET.TRANSFORMS
+    t = []
+    if transforms.CENTER_CROP:
+        t.append(T.CenterCrop(transforms.CENTER_CROP))
+    if mode == Datasets.TRAIN:
+        if transforms.RANDOM_HORIZONTAL_FLIP:
+            t.append(T.RandomHorizontalFlip(transforms.RANDOM_HORIZONTAL_FLIP))
+        t.append(T.ColorJitter(brightness=transforms.COLOR_JITTER.BRIGHTNESS,
+                               contrast=transforms.COLOR_JITTER.CONTRAST,
+                               saturation=transforms.COLOR_JITTER.SATURATION,
+                               hue=transforms.COLOR_JITTER.HUE))
+        t.append(Shear(transforms.SHEAR))
+        t.append(Scale(transforms.SCALE.HORIZONTAL,
+                       transforms.SCALE.VERTICAL))
+        t.append(Translate(transforms.TRANSLATE.HORIZONTAL,
+                           transforms.TRANSLATE.VERTICAL))
+    if transforms.RESIZE:
+        t.append(T.Resize(tuple(reversed(transforms.RESIZE))))
+    t.extend([T.ToTensor(),
+              T.Normalize(mean=_MEAN, std=_STD)])
+    return T.Compose(t)

chesscog/core/io/download.py

@@ -16,15 +16,21 @@
 def _get_members(f: zipfile.ZipFile) -> typing.Iterator[zipfile.ZipInfo]:
     parts = [name.partition("/")[0]
              for name in f.namelist()
-             if not name.endswith("/")]
+             if not name.endswith("/")
+             and not ".DS_Store" in name
+             and not "__MACOSX" in name]
 
     prefix = os.path.commonprefix(parts)
-    if prefix:
-        prefix += "/"
+    if "/" in prefix:
+        prefix = prefix[:prefix.rfind("/") + 1]
+    else:
+        prefix = ""
     offset = len(prefix)
     # Alter file names
     for zipinfo in f.infolist():
         name = zipinfo.filename
+        if ".DS_Store" in name or "__MACOSX" in name:
+            continue
         if len(name) > offset:
             zipinfo.filename = name[offset:]
             yield zipinfo

chesscog/core/training/train.py

@@ -21,8 +21,15 @@
 
 
 def train(cfg: CN, run_dir: Path) -> nn.Module:
+    model = build_model(cfg)
+    is_inception = "inception" in cfg.TRAINING.MODEL.NAME.lower()
+    train_model(cfg, run_dir, model, is_inception)
+
+
+def train_model(cfg: CN, run_dir: Path, model: torch.nn.Module, is_inception: bool = False, model_name: str = None, eval_on_train: bool = False):
     logger.info(f"Starting training in {run_dir}")
-    model_name = run_dir.name
+    if not model_name:
+        model_name = run_dir.name
 
     # Create folder
     if run_dir.exists():
@@ -35,18 +42,22 @@ def train(cfg: CN, run_dir: Path) -> nn.Module:
     with (run_dir / f"{model_name}.yaml").open("w") as f:
         cfg.dump(stream=f)
 
-    model = build_model(cfg)
+    # Move model to device
     device(model)
-    is_inception = "inception" in cfg.TRAINING.MODEL.NAME.lower()
 
     best_weights, best_accuracy, best_step = copy.deepcopy(
         model.state_dict()), 0., 0
 
     criterion = nn.CrossEntropyLoss()
 
     modes = {Datasets.TRAIN, Datasets.VAL}
-    datasets = {mode: build_dataset(cfg, mode)
-                for mode in modes}
+    if eval_on_train:
+        dataset = build_dataset(cfg, Datasets.TRAIN)
+        datasets = {mode: dataset
+                    for mode in modes}
+    else:
+        datasets = {mode: build_dataset(cfg, mode)
+                    for mode in modes}
     classes = datasets[Datasets.TRAIN].classes
     loader = {mode: build_data_loader(cfg, datasets[mode], mode)
               for mode in modes}
@@ -150,7 +161,7 @@ def perform_iteration(data: typing.Tuple[torch.Tensor, torch.Tensor], mode: Data
 
                 # Save weights if we get a better performance
                 accuracy = aggregator[Datasets.VAL].accuracy()
-                if accuracy > best_accuracy:
+                if accuracy >= best_accuracy:
                     best_accuracy = accuracy
                     best_weights = copy.deepcopy(model.state_dict())
                     best_step = step

chesscog/corner_detection/detect_corners.py

@@ -148,7 +148,16 @@ def _absolute_angle_difference(x, y):
     return np.min(np.stack([diff, np.pi - diff], axis=-1), axis=-1)
 
 
+def _sort_lines(lines: np.ndarray) -> np.ndarray:
+    if lines.ndim == 0 or lines.shape[-2] == 0:
+        return lines
+    rhos = lines[..., 0]
+    sorted_indices = np.argsort(rhos)
+    return lines[sorted_indices]
+
+
 def cluster_horizontal_and_vertical_lines(lines: np.ndarray):
+    lines = _sort_lines(lines)
     thetas = lines[..., 1].reshape(-1, 1)
     distance_matrix = pairwise_distances(
         thetas, thetas, metric=_absolute_angle_difference)

chesscog/occupancy_classifier/create_dataset.py

@@ -41,9 +41,10 @@ def warp_chessboard_image(img: np.ndarray, corners: np.ndarray) -> np.ndarray:
     return cv2.warpPerspective(img, transformation_matrix, (IMG_SIZE, IMG_SIZE))
 
 
-def extract_squares_from_sample(id: str, subset: str = ""):
-    img = cv2.imread(str(RENDERS_DIR / subset / (id + ".png")))
-    with (RENDERS_DIR / subset / (id + ".json")).open("r") as f:
+def extract_squares_from_sample(id: str, subset: str = "", input_dir: Path = RENDERS_DIR, output_dir: Path = OUT_DIR):
+    img = cv2.imread(str(input_dir / subset / (id + ".png")))
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    with (input_dir / subset / (id + ".json")).open("r") as f:
         label = json.load(f)
 
     corners = np.array(label["corners"], dtype=np.float)
@@ -56,18 +57,23 @@ def extract_squares_from_sample(id: str, subset: str = ""):
             square) is None else "occupied"
         piece_img = crop_square(unwarped, square, label["white_turn"])
         with Image.fromarray(piece_img, "RGB") as piece_img:
-            piece_img.save(OUT_DIR / subset / target_class /
+            piece_img.save(output_dir / subset / target_class /
                            f"{id}_{chess.square_name(square)}.png")
 
 
-if __name__ == "__main__":
+def create_dataset(input_dir: Path = RENDERS_DIR, output_dir: Path = OUT_DIR):
     for subset in ("train", "val", "test"):
         for c in ("empty", "occupied"):
-            folder = OUT_DIR / subset / c
+            folder = output_dir / subset / c
             shutil.rmtree(folder, ignore_errors=True)
             os.makedirs(folder, exist_ok=True)
-        samples = list((RENDERS_DIR / subset).glob("*.png"))
+        samples = list((input_dir / subset).glob("*.png"))
         for i, img_file in enumerate(samples):
-            if i % int(len(samples) / 100) == 0:
+            if len(samples) > 100 and i % int(len(samples) / 100) == 0:
                 print(f"{i / len(samples)*100:.0f}%")
-            extract_squares_from_sample(img_file.stem, subset)
+            extract_squares_from_sample(img_file.stem, subset,
+                                        input_dir, output_dir)
+
+
+if __name__ == "__main__":
+    create_dataset()