Merge pull request #6 from georgw777/develop

v0.2.6

Dockerfile

@@ -18,20 +18,13 @@ RUN pip install --upgrade pip && \
     pip install poetry && \
     poetry config virtualenvs.create false
 
-# Install dependencies
-COPY ./pyproject.toml ./poetry.lock* /app/
-RUN poetry install --no-root --no-dev
-
 # Install dependencies
 RUN mkdir -p /chess
 WORKDIR /chess
-COPY poetry.lock pyproject.toml ./
-RUN poetry install
+COPY ./pyproject.toml ./poetry.lock* ./
+RUN poetry install --no-root
 ENV PYTHONPATH "/chess:${PYTHONPATH}"
 
-# Fix for tensorboard
-RUN poetry run pip install wheel
-
 # Setup data mount
 RUN mkdir -p /data
 ENV DATA_DIR /data

chesscog/__init__.py

@@ -1,7 +1,7 @@
 import sys
 import logging
 
-from .utils import io as _
+from .core import io as _
 from .__version__ import __version__
 
 

chesscog/__version__.py

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

chesscog/utils/evaluation.py → chesscog/core/evaluation.py

@@ -8,9 +8,9 @@
 from PIL import Image
 from recap import URI, CfgNode as CN
 
-from chesscog.utils.dataset import build_dataset, build_data_loader, Datasets, unnormalize
-from chesscog.utils.statistics import StatsAggregator
-from chesscog.utils import device, DEVICE
+from chesscog.core.dataset import build_dataset, build_data_loader, Datasets, unnormalize
+from chesscog.core.statistics import StatsAggregator
+from chesscog.core import device, DEVICE
 
 logger = logging.getLogger(__name__)
 

chesscog/core/exceptions.py

@@ -0,0 +1,11 @@
+class RecognitionException(Exception):
+    def __init__(self, message: str = "unknown error"):
+        super().__init__("chess recognition error: " + message)
+
+
+class ChessboardNotLocatedException(RecognitionException):
+    def __init__(self, reason: str = None):
+        message = "chessboard could not be located"
+        if reason:
+            message += ": " + reason
+        super().__init__(message)

chesscog/utils/io/__init__.py → chesscog/core/io/__init__.py

@@ -2,7 +2,7 @@
 import os
 from recap.path_manager import register_translator
 
-from .download import download_file, download_zip_folder
+from .download import download_file, download_zip_folder, download_zip_folder_from_google_drive
 
 
 _DATA_DIR = Path(os.getenv("DATA_DIR",

chesscog/utils/models.py → chesscog/core/models.py

@@ -1,7 +1,7 @@
 from torch import nn
 from recap import CfgNode as CN
 
-from chesscog.utils.registry import Registry
+from chesscog.core.registry import Registry
 
 MODELS_REGISTRY = Registry()
 

chesscog/utils/training/create_configs.py → chesscog/core/training/create_configs.py

@@ -1,7 +1,7 @@
 import logging
 from recap import URI, CfgNode as CN
 
-from chesscog.utils.models import MODELS_REGISTRY
+from chesscog.core.models import MODELS_REGISTRY
 
 logger = logging.getLogger(__name__)
 

chesscog/utils/training/train.py → chesscog/core/training/train.py

@@ -11,11 +11,11 @@
 import shutil
 from recap import CfgNode as CN
 
-from chesscog.utils import device
-from chesscog.utils.training import build_optimizer_from_config
-from chesscog.utils.statistics import StatsAggregator
-from chesscog.utils.dataset import build_dataset, build_data_loader, Datasets
-from chesscog.utils.models import build_model
+from chesscog.core import device
+from chesscog.core.training import build_optimizer_from_config
+from chesscog.core.statistics import StatsAggregator
+from chesscog.core.dataset import build_dataset, build_data_loader, Datasets
+from chesscog.core.models import build_model
 
 logger = logging.getLogger(__name__)
 

chesscog/utils/training/train_classifier.py → chesscog/core/training/train_classifier.py

@@ -2,7 +2,7 @@
 import logging
 from recap import URI, CfgNode as CN
 
-from chesscog.utils.training import train
+from chesscog.core.training import train
 
 logger = logging.getLogger(__name__)
 

chesscog/corner_detection/create_configs.py

@@ -3,7 +3,7 @@
 import functools
 from recap import URI, CfgNode as CN
 
-from chesscog.utils import listify
+from chesscog.core import listify
 
 parameters = {
     "EDGE_DETECTION.LOW_THRESHOLD": np.arange(80, 151, 10),

chesscog/corner_detection/detect_corners.py

@@ -5,18 +5,19 @@
 import typing
 from recap import URI, CfgNode as CN
 
-from chesscog.utils import sort_corner_points
-from chesscog.utils.coordinates import from_homogenous_coordinates, to_homogenous_coordinates
+from chesscog.core import sort_corner_points
+from chesscog.core.coordinates import from_homogenous_coordinates, to_homogenous_coordinates
+from chesscog.core.exceptions import ChessboardNotLocatedException
 
 
 def find_corners(cfg: CN, img: np.ndarray) -> np.ndarray:
     img, img_scale = resize_image(cfg, img)
 
     gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    edges = detect_edges(cfg, gray)
+    edges = detect_edges(cfg.EDGE_DETECTION, gray)
     lines = detect_lines(cfg, edges)
     if lines.shape[0] > 400:
-        return None
+        raise ChessboardNotLocatedException("too many lines in the image")
     all_horizontal_lines, all_vertical_lines = cluster_horizontal_and_vertical_lines(
         lines)
 
@@ -31,7 +32,7 @@ def find_corners(cfg: CN, img: np.ndarray) -> np.ndarray:
     best_num_inliers = 0
     best_configuration = None
     iterations = 0
-    while iterations < 50 or best_num_inliers < 40:
+    while iterations < 200 or best_num_inliers < 30:
         row1, row2 = _choose_from_range(len(horizontal_lines))
         col1, col2 = _choose_from_range(len(vertical_lines))
         transformation_matrix = compute_homography(all_intersection_points,
@@ -55,8 +56,9 @@ def find_corners(cfg: CN, img: np.ndarray) -> np.ndarray:
                 best_num_inliers = num_inliers
                 best_configuration = configuration
         iterations += 1
-        if iterations > 1000:
-            return None
+        if iterations > 10000:
+            raise ChessboardNotLocatedException(
+                "RANSAC produced no viable results")
 
     # Retrieve best configuration
     (xmin, xmax, ymin, ymax), scale, quantized_points, intersection_points, warped_img_size = best_configuration
@@ -67,12 +69,20 @@ def find_corners(cfg: CN, img: np.ndarray) -> np.ndarray:
     inverse_transformation_matrix = np.linalg.inv(transformation_matrix)
 
     # Warp grayscale image
-    warped = cv2.warpPerspective(
-        gray, transformation_matrix, tuple(warped_img_size.astype(np.int)))
+    dims = tuple(warped_img_size.astype(np.int))
+    warped = cv2.warpPerspective(gray, transformation_matrix, dims)
+    borders = np.zeros_like(gray)
+    borders[3:-3, 3:-3] = 1
+    warped_borders = cv2.warpPerspective(borders, transformation_matrix, dims)
+    warped_mask = warped_borders == 1
 
     # Refine board boundaries
-    xmin, xmax = compute_vertical_borders(cfg, warped, scale, xmin, xmax)
-    ymin, ymax = compute_horizontal_borders(cfg, warped, scale, ymin, ymax)
+    xmin, xmax = compute_vertical_borders(
+        cfg, warped, warped_mask, scale, xmin, xmax)
+    scaled_xmin, scaled_xmax = (int(x * scale[0]) for x in (xmin, xmax))
+    warped_mask[:, :scaled_xmin] = warped_mask[:, scaled_xmax:] = False
+    ymin, ymax = compute_horizontal_borders(
+        cfg, warped, warped_mask, scale, ymin, ymax)
 
     # Transform boundaries to image space
     corners = np.array([[xmin, ymin],
@@ -96,11 +106,14 @@ def resize_image(cfg: CN, img: np.ndarray) -> typing.Tuple[np.ndarray, float]:
     return img, scale
 
 
-def detect_edges(cfg: CN, gray: np.ndarray) -> np.ndarray:
+def detect_edges(edge_detection_cfg: CN, gray: np.ndarray) -> np.ndarray:
+    if gray.dtype != np.uint8:
+        gray = gray / gray.max() * 255
+        gray = gray.astype(np.uint8)
     edges = cv2.Canny(gray,
-                      cfg.EDGE_DETECTION.LOW_THRESHOLD,
-                      cfg.EDGE_DETECTION.HIGH_THRESHOLD,
-                      cfg.EDGE_DETECTION.APERTURE)
+                      edge_detection_cfg.LOW_THRESHOLD,
+                      edge_detection_cfg.HIGH_THRESHOLD,
+                      edge_detection_cfg.APERTURE)
     return edges
 
 
@@ -213,10 +226,10 @@ def compute_homography(intersection_points: np.ndarray, row1: int, row2: int, co
     p4 = intersection_points[row2, col1]  # bottom left
 
     src_points = np.stack([p1, p2, p3, p4])
-    dst_points = np.array([[0, 1],  # top left
-                           [1, 1],  # top right
-                           [1, 0],  # bottom right
-                           [0, 0]])  # bottom left
+    dst_points = np.array([[0, 0],  # top left
+                           [1, 0],  # top right
+                           [1, 1],  # bottom right
+                           [0, 1]])  # bottom left
     return compute_transformation_matrix(src_points, dst_points)
 
 
@@ -311,43 +324,50 @@ def quantize_points(cfg: CN, warped_scaled_points: np.ndarray, intersection_poin
     return (xmin, xmax, ymin, ymax), scale, scaled_quantized_points, intersection_points, warped_img_size
 
 
-def compute_vertical_borders(cfg: CN, warped: np.ndarray, scale: np.ndarray, xmin: int, xmax: int) -> typing.Tuple[int, int]:
-    G_y = np.abs(cv2.Sobel(warped, cv2.CV_64F, 1, 0,
+def compute_vertical_borders(cfg: CN, warped: np.ndarray, mask: np.ndarray, scale: np.ndarray, xmin: int, xmax: int) -> typing.Tuple[int, int]:
+    G_x = np.abs(cv2.Sobel(warped, cv2.CV_64F, 1, 0,
                            ksize=cfg.BORDER_REFINEMENT.SOBEL_KERNEL_SIZE))
+    G_x[~mask] = 0
+    G_x = detect_edges(cfg.BORDER_REFINEMENT.EDGE_DETECTION.VERTICAL, G_x)
+    G_x[~mask] = 0
 
-    def get_vertical_line_score(x):
+    def get_nonmax_supressed(x):
         x = (x * scale[0]).astype(np.int)
         thresh = cfg.BORDER_REFINEMENT.LINE_WIDTH // 2
-        return G_y[:, x-thresh:x+thresh+1].sum()
+        return G_x[:, x-thresh:x+thresh+1].max(axis=1)
 
     while xmax - xmin < 8:
-        top_score = get_vertical_line_score(xmax + 1)
-        bottom_score = get_vertical_line_score(xmin - 1)
-        if top_score > bottom_score:
+        top = get_nonmax_supressed(xmax + 1)
+        bottom = get_nonmax_supressed(xmin - 1)
+
+        if top.sum() > bottom.sum():
             xmax += 1
         else:
             xmin -= 1
 
     return xmin, xmax
 
 
-def compute_horizontal_borders(cfg: CN, warped: np.ndarray, scale: np.ndarray, ymin: int, ymax: int) -> typing.Tuple[int, int]:
+def compute_horizontal_borders(cfg: CN, warped: np.ndarray, mask: np.ndarray, scale: np.ndarray, ymin: int, ymax: int) -> typing.Tuple[int, int]:
     G_y = np.abs(cv2.Sobel(warped, cv2.CV_64F, 0, 1,
                            ksize=cfg.BORDER_REFINEMENT.SOBEL_KERNEL_SIZE))
+    G_y[~mask] = 0
+    G_y = detect_edges(cfg.BORDER_REFINEMENT.EDGE_DETECTION.HORIZONTAL, G_y)
+    G_y[~mask] = 0
 
-    def get_horizontal_line_score(y):
+    def get_nonmax_supressed(y):
         y = (y * scale[1]).astype(np.int)
         thresh = cfg.BORDER_REFINEMENT.LINE_WIDTH // 2
-        return G_y[y-thresh:y+thresh+1].sum()
+        return G_y[y-thresh:y+thresh+1].max(axis=0)
 
     while ymax - ymin < 8:
-        top_score = get_horizontal_line_score(ymax + 1)
-        bottom_score = get_horizontal_line_score(ymin - 1)
-        if top_score > bottom_score:
+        top = get_nonmax_supressed(ymax + 1)
+        bottom = get_nonmax_supressed(ymin - 1)
+
+        if top.sum() > bottom.sum():
             ymax += 1
         else:
             ymin -= 1
-
     return ymin, ymax
 
 

chesscog/corner_detection/evaluate.py

@@ -7,8 +7,8 @@
 import logging
 from recap import URI, CfgNode as CN
 
-from chesscog.utils.dataset import Datasets
-from chesscog.utils import sort_corner_points
+from chesscog.core.dataset import Datasets
+from chesscog.core import sort_corner_points
 from chesscog.corner_detection import find_corners
 
 logger = logging.getLogger(__name__)

chesscog/data_synthesis/download_dataset.py

@@ -1,7 +1,7 @@
 """Script to download the rendered dataset."""
 
 import functools
-from chesscog.utils.io import download_zip_folder_from_google_drive
+from chesscog.core.io import download_zip_folder_from_google_drive
 
 ensure_dataset = functools.partial(download_zip_folder_from_google_drive,
                                    "1XClmGJwEWNcIkwaH0VLuvvAY3qk_CRJh",

chesscog/occupancy_classifier/create_configs.py

@@ -1,4 +1,4 @@
-from chesscog.utils.training import create_configs
+from chesscog.core.training import create_configs
 
 if __name__ == "__main__":
     create_configs("occupancy_classifier", include_centercrop=True)

chesscog/occupancy_classifier/create_dataset.py

@@ -9,7 +9,7 @@
 import shutil
 from recap import URI
 
-from chesscog.utils import sort_corner_points
+from chesscog.core import sort_corner_points
 
 RENDERS_DIR = URI("data://render")
 OUT_DIR = URI("data://occupancy")

chesscog/occupancy_classifier/download_model.py

@@ -1,7 +1,7 @@
 """Script to download the best occupancy classifier."""
 
 import functools
-from chesscog.utils.io import download_zip_folder
+from chesscog.core.io import download_zip_folder
 
 
 ensure_model = functools.partial(download_zip_folder,

chesscog/occupancy_classifier/evaluate.py

@@ -1,4 +1,4 @@
-from chesscog.utils.evaluation import perform_evaluation
+from chesscog.core.evaluation import perform_evaluation
 
 if __name__ == "__main__":
     perform_evaluation("occupancy_classifier")

chesscog/occupancy_classifier/models.py

@@ -4,8 +4,8 @@
 import functools
 from recap import CfgNode as CN
 
-from chesscog.utils.registry import Registry
-from chesscog.utils.models import MODELS_REGISTRY
+from chesscog.core.registry import Registry
+from chesscog.core.models import MODELS_REGISTRY
 
 NUM_CLASSES = 2
 

chesscog/occupancy_classifier/train.py

@@ -1,4 +1,4 @@
-from chesscog.utils.training import train_classifier
+from chesscog.core.training import train_classifier
 
 if __name__ == "__main__":
     train_classifier("occupancy_classifier")

chesscog/piece_classifier/create_configs.py

@@ -1,4 +1,4 @@
-from chesscog.utils.training import create_configs
+from chesscog.core.training import create_configs
 
 if __name__ == "__main__":
     create_configs("piece_classifier")

chesscog/piece_classifier/create_dataset.py

@@ -9,8 +9,8 @@
 import shutil
 from recap import URI
 
-from chesscog.utils import sort_corner_points
-from chesscog.utils.dataset import piece_name
+from chesscog.core import sort_corner_points
+from chesscog.core.dataset import piece_name
 
 RENDERS_DIR = URI("data://render")
 OUT_DIR = URI("data://pieces")

chesscog/piece_classifier/download_model.py

@@ -1,7 +1,7 @@
 """Script to download the best piece classifier."""
 
 import functools
-from chesscog.utils.io import download_zip_folder
+from chesscog.core.io import download_zip_folder
 
 
 ensure_model = functools.partial(download_zip_folder,

chesscog/piece_classifier/evaluate.py

@@ -1,4 +1,4 @@
-from chesscog.utils.evaluation import perform_evaluation
+from chesscog.core.evaluation import perform_evaluation
 
 if __name__ == "__main__":
     perform_evaluation("piece_classifier")