├── .gitignore ├── Brutesac.py ├── CITATION.cff ├── Chessboard Detect.ipynb ├── FindChessboards.ipynb ├── FindChessboards.py ├── FindChessboardsWithML.py ├── GradientAngleBasedChessboardDetection.ipynb ├── InformedHoughOnPoints.ipynb ├── LICENSE ├── README.md ├── RunBrutesacOnVideo.py ├── RunContourSac.py ├── RunContourSacLK.py ├── RunExportedMLOnImage.py ├── RunMLOnImage.py ├── SaddlePoints.py ├── StoreChessboardPoints.py ├── XcornerRansac.ipynb ├── base_imgload.py ├── board_detect.py ├── buildMLSavedModel.py ├── buildMLSavedModel2.py ├── centralSymmetryTile.py ├── chess_detect_helper.py ├── chessboard_detect.py ├── chessboard_detect2.py ├── contour_detect.py ├── countEntriesInTfRecords.py ├── data ├── bad_tiles.png └── good_tiles.png ├── example_unity_grid.png ├── generateFullDataset.py ├── generateMLDataset.py ├── generateMLTiles.py ├── helpers.py ├── hog.py ├── hough_visualize.py ├── image_segment.py ├── inform_hough_on_pts.py ├── line_intersection.py ├── makeGif.sh ├── make_videos.sh ├── oriented_convolve.py ├── outlier_point_removal.py ├── output_ml.gif ├── processChessPoints.py ├── quickCheck_deleteeasily.py ├── quickFix.py ├── readme_find_warp_example.png ├── readme_labeled.png ├── readme_output.png ├── readme_rectified.gif ├── rectify_refine.py ├── result.png ├── run_xcorner_model_on_img.py ├── run_xcorner_on_saddle_tiles.py ├── saddle ├── 0.jpg ├── 1.jpg ├── 2.jpg └── saddle_tutorial.ipynb ├── sam2_composite.gif ├── speedchess1_composite.gif ├── speedchess1_ml.gif ├── tile_segment.py ├── trainML_model.py ├── training_pipeline ├── doTestSweep.sh ├── model.py ├── preprocess.py ├── saveModel.py ├── testModelOnVideo.py ├── trainCNN_tfrecords.py ├── trainML_CNN_pipeline.py └── trainML_pipeline.py ├── triangle_mesh.png └── view_xcorner_heatmap.py /.gitignore: -------------------------------------------------------------------------------- 1 | # folders 2 | datasets/ 3 | ml/ 4 | results/ 5 | training_pipeline/training_models/ 6 | 7 | # images/videos 8 | *.avi 9 | *.mp4 10 | #*.png 11 | #*.jpg 12 | *.gif 13 | 14 | # Byte-compiled / optimized / DLL files 15 | __pycache__/ 16 | *.py[cod] 17 | *$py.class 18 | 19 | # C extensions 20 | *.so 21 | 22 | # Distribution / packaging 23 | .Python 24 | env/ 25 | build/ 26 | develop-eggs/ 27 | dist/ 28 | downloads/ 29 | eggs/ 30 | .eggs/ 31 | lib/ 32 | lib64/ 33 | parts/ 34 | sdist/ 35 | var/ 36 | *.egg-info/ 37 | .installed.cfg 38 | *.egg 39 | 40 | # PyInstaller 41 | # Usually these files are written by a python script from a template 42 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 43 | *.manifest 44 | *.spec 45 | 46 | # Installer logs 47 | pip-log.txt 48 | pip-delete-this-directory.txt 49 | 50 | # Unit test / coverage reports 51 | htmlcov/ 52 | .tox/ 53 | .coverage 54 | .coverage.* 55 | .cache 56 | nosetests.xml 57 | coverage.xml 58 | *,cover 59 | .hypothesis/ 60 | 61 | # Translations 62 | *.mo 63 | *.pot 64 | 65 | # Django stuff: 66 | *.log 67 | local_settings.py 68 | 69 | # Flask stuff: 70 | instance/ 71 | .webassets-cache 72 | 73 | # Scrapy stuff: 74 | .scrapy 75 | 76 | # Sphinx documentation 77 | docs/_build/ 78 | 79 | # PyBuilder 80 | target/ 81 | 82 | # IPython Notebook 83 | .ipynb_checkpoints 84 | 85 | # pyenv 86 | .python-version 87 | 88 | # celery beat schedule file 89 | celerybeat-schedule 90 | 91 | # dotenv 92 | .env 93 | 94 | # virtualenv 95 | venv/ 96 | ENV/ 97 | 98 | # Spyder project settings 99 | .spyderproject 100 | 101 | # Rope project settings 102 | .ropeproject 103 | -------------------------------------------------------------------------------- /Brutesac.py: -------------------------------------------------------------------------------- 1 | 2 | # coding: utf-8 3 | 4 | # We have a bunch of probably chess x-corners output by our ML detector. Let's do the following routine: 5 | # 6 | # 1. Take 4 random points, (check no 3 points are colinear, else retry) then warp them to a unity grid. 7 | # 2. Count distances of each point to closest integer grid corner, sum of these will be the score. 8 | # 3. Do this for an exhaustive set of random points? 9 | # 10 | # The best score will likely be the correct transform to a chessboard. 11 | 12 | # In[731]: 13 | 14 | import numpy as np 15 | from IPython.display import Image, display 16 | import PIL.Image 17 | import matplotlib.pyplot as plt 18 | import scipy.ndimage 19 | import cv2 # For Sobel etc 20 | import glob 21 | import SaddlePoints 22 | from scipy.spatial import Delaunay 23 | from functools import wraps 24 | import time 25 | 26 | def timed(f): 27 | @wraps(f) 28 | def wrapper(*args, **kwds): 29 | start = time.time() 30 | result = f(*args, **kwds) 31 | elapsed = time.time() - start 32 | print "%s took %.2f ms to finish" % (f.__name__, elapsed*1e3) 33 | return result 34 | return wrapper 35 | 36 | # # RANSAC to find the best 4 points signifying a chess tile. 37 | # Idea to find quads from triangles. For each triangle, the triangle that completes the quad should always be on the longest edge of the triangle and shares the two vertices of that edge. No image gradient checks needed. 38 | 39 | # Build up a list of quads from input delaunay triangles, returns an Nx4 list of indices on the points used. 40 | def getAllQuads(tri): 41 | pairings = set() 42 | quads = [] 43 | # Build up a list of all possible unique quads from triangle neighbor pairings. 44 | # In general the worst common case with a fully visible board is 6*6*2=36 triangles, each with 3 neighbor 45 | # so around ~100 quads. 46 | for i, neighbors in enumerate(tri.neighbors): 47 | for k in range(3): # For each potential neighbor (3, one opposing each vertex of triangle) 48 | nk = neighbors[k] 49 | if nk != -1: 50 | # There is a neighbor, create a quad unless it already exists in set 51 | pair = (i, nk) 52 | reverse_pair = (nk, i) 53 | if reverse_pair not in pairings: 54 | # New pair, add and create a quad 55 | pairings.add(pair) 56 | b = tri.simplices[i] 57 | d = tri.simplices[nk] 58 | nk_vtx = (set(d) - set(b)).pop() 59 | insert_mapping = [2,3,1] 60 | b = np.insert(b,insert_mapping[k], nk_vtx) 61 | quads.append(b) 62 | return np.array(quads) 63 | 64 | def countHits(given_pts, x_offset, y_offset): 65 | # Check the given integer points (in unity grid coordinate space) for matches 66 | # to an ideal chess grid with given initial offsets 67 | pt_set = set((a,b) for a,b in given_pts) 68 | [X,Y] = np.meshgrid(np.arange(7) + x_offset,np.arange(7) + y_offset) 69 | matches = 0 70 | # count matching points in set 71 | matches = sum(1 for x,y in zip(X.flatten(), Y.flatten()) if (x,y) in pt_set) 72 | return matches 73 | 74 | def getBestBoardMatchup(given_pts): 75 | best_score = 0 76 | best_offset = None 77 | # scores = np.zeros([7,7]) 78 | for i in range(7): 79 | for j in range(7): 80 | # Offsets from -6 to 0 for both 81 | score = countHits(given_pts, i-6, j-6) 82 | # scores[i,j] = score 83 | if score > best_score: 84 | best_score = score 85 | best_offset = [i-6, j-6] 86 | # print scores 87 | return best_score, best_offset 88 | 89 | def scoreQuad(quad, pts, prevBestScore=0): 90 | idealQuad = np.array([[0,1],[1,1],[1,0],[0,0]],dtype=np.float32) 91 | M = cv2.getPerspectiveTransform(quad.astype(np.float32), idealQuad) 92 | 93 | # warped_to_ideal = cv2.perspectiveTransform(np.expand_dims(quad.astype(float),0), M) 94 | 95 | # Warp points and score error 96 | pts_warped = cv2.perspectiveTransform(np.expand_dims(pts.astype(float),0), M)[0,:,:] 97 | 98 | # Get their closest idealized grid point 99 | pts_warped_int = pts_warped.round().astype(int) 100 | 101 | # Refine 102 | M_refined, _ = cv2.findHomography(pts, pts_warped_int, cv2.RANSAC) 103 | if (M_refined is None): 104 | M_refined = M 105 | 106 | # Re-Warp points with the refined M and score error 107 | pts_warped = cv2.perspectiveTransform(np.expand_dims(pts.astype(float),0), M_refined)[0,:,:] 108 | 109 | # Get their closest idealized grid point 110 | pts_warped_int = pts_warped.round().astype(int) 111 | 112 | 113 | # Count matches 114 | score, offset = getBestBoardMatchup(pts_warped_int) 115 | if score < prevBestScore: 116 | return score, None, None, None 117 | 118 | # Sum of distances from closest integer value for each point 119 | # Use this error score for tie-breakers where number of matches is the same. 120 | error_score = np.sum(np.linalg.norm((pts_warped - pts_warped_int), axis=1)) 121 | 122 | return score, error_score, M, offset 123 | 124 | # brutesacChessboard ~752ms 125 | # getAllQuads 12ms 126 | 127 | @timed 128 | def brutesacChessboard(xcorner_pts): 129 | # Build a list of quads to try. 130 | tri = Delaunay(xcorner_pts) 131 | quads = getAllQuads(tri) # 132 | 133 | # For each quad, keep track of the best fitting chessboard. 134 | best_score = 0 135 | best_error_score = None 136 | best_M = None 137 | best_quad = None 138 | best_offset = None 139 | for quad in xcorner_pts[quads]: 140 | score, error_score, M, offset = scoreQuad(quad, xcorner_pts, best_score) 141 | if score > best_score or (score == best_score and error_score < best_error_score): 142 | best_score = score 143 | best_error_score = error_score 144 | best_M = M 145 | best_quad = quad 146 | best_offset = offset 147 | if best_score > (len(xcorner_pts)*0.8): 148 | break 149 | 150 | return best_M, best_quad, best_offset, best_score, best_error_score 151 | 152 | 153 | # @timed 154 | def refineHomography(pts, M, best_offset): 155 | pts_warped = cv2.perspectiveTransform(np.expand_dims(pts.astype(float),0), M)[0,:,:] 156 | a = pts_warped.round() - best_offset 157 | b = pts 158 | 159 | # Only use points from within chessboard region. 160 | outliers = np.any((a < 0) | (a > 7), axis=1) 161 | a = a[~outliers] 162 | b = b[~outliers] 163 | 164 | # Least-Median robust homography 165 | M_homog, _ = cv2.findHomography(b,a, cv2.LMEDS) 166 | 167 | return M_homog 168 | 169 | @timed 170 | def predictOnTiles(tiles, predict_fn): 171 | predictions = predict_fn( 172 | {"x": tiles}) 173 | 174 | # Return array of probability of tile being an xcorner. 175 | # return np.array([p[1] for p in predictions['probabilities']]) 176 | return np.array([p[1] for p in predictions['probabilities']]) 177 | 178 | @timed 179 | def predictOnImage(pts, img_gray, predict_fn, WINSIZE = 10): 180 | # Build tiles to run classifier on. (23 ms) 181 | tiles = getTilesFromImage(pts, img_gray, WINSIZE=WINSIZE) 182 | 183 | # Classify tiles. (~137ms) 184 | probs = predictOnTiles(tiles, predict_fn) 185 | 186 | return probs 187 | 188 | @timed 189 | def getTilesFromImage(pts, img_gray, WINSIZE=10): 190 | # NOTE : Assumes no point is within WINSIZE of an edge! 191 | 192 | # Points Nx2, columns should be x and y, not r and c. 193 | # Build tiles 194 | img_shape = np.array([img_gray.shape[1], img_gray.shape[0]]) 195 | tiles = np.zeros([len(pts), WINSIZE*2+1, WINSIZE*2+1]) 196 | for i, pt in enumerate(np.round(pts).astype(np.int64)): 197 | tiles[i,:,:] = img_gray[pt[1]-WINSIZE:pt[1]+WINSIZE+1, pt[0]-WINSIZE:pt[0]+WINSIZE+1] 198 | 199 | return tiles 200 | 201 | @timed 202 | def classifyPoints(pts, img_gray, predict_fn, WINSIZE = 10): 203 | tiles = getTilesFromImage(pts, img_gray, WINSIZE=WINSIZE) 204 | 205 | # Classify tiles. 206 | probs = predictOnTiles(tiles, predict_fn) 207 | 208 | return probs 209 | 210 | @timed 211 | def classifyImage(img_gray, predict_fn, WINSIZE = 10, prob_threshold=0.5): 212 | spts = SaddlePoints.getFinalSaddlePoints(img_gray, WINSIZE) 213 | 214 | return spts[predictOnImage(spts, img_gray, predict_fn, WINSIZE) > prob_threshold, :] 215 | 216 | @timed 217 | def processFrame(gray): 218 | # and return M for chessboard from image 219 | pts = classifyImage(gray) 220 | 221 | # pts = np.loadtxt('example_pts.txt') 222 | pts = pts[:,[1,0]] # Switch rows/cols to x/y for plotting on an image 223 | 224 | if (len(pts)) < 6: 225 | print("Probably not enough points (%d) for a good fit, skipping" % len(pts)) 226 | return None, None 227 | 228 | raw_M, best_quad, best_offset, best_score, best_error_score = brutesacChessboard(pts) 229 | 230 | # The chess outline is only a rough guess, which we can easily refine now that we have an understanding of which points are meant for which chess corner. One could do distortion correction with the points available (as normal chessboard algorithms are used for in computer vision). If we wanted to get fancy we could apply a warp flow field on a per-tile basis. 231 | 232 | # ## Refine with Homography on all points 233 | # Now we can refine our best guess using all the found x-corner points to their ideal points with opencv findHomography, using the regular method with all points. 234 | if (raw_M is not None): 235 | M_homog = refineHomography(pts, raw_M, best_offset) 236 | else: 237 | M_homog = None 238 | 239 | return M_homog, pts 240 | 241 | 242 | 243 | if __name__ == '__main__': 244 | np.set_printoptions(suppress=True) # Better printing of arrays 245 | plt.rcParams['image.cmap'] = 'jet' # Default colormap is jet 246 | 247 | # img = PIL.Image.open('frame900.png') 248 | # img = PIL.Image.open('input/img_35.jpg') 249 | filenames = glob.glob('input_yt/*.jpg') 250 | filenames.extend(glob.glob('input_yt/*.png')) 251 | filenames.extend(glob.glob('input/*.jpg')) 252 | filenames.extend(glob.glob('input/*.png')) 253 | # filenames.extend(glob.glob('input_bad/*.png')) 254 | # filenames.extend(glob.glob('input_bad/*.jpg')) 255 | filename = np.random.choice(filenames,1)[0] 256 | filename= 'weird.jpg' 257 | print(filename) 258 | img = PIL.Image.open(filename) 259 | if (img.size[0] > 800): 260 | img = img.resize([600,400]) 261 | gray = np.array(img.convert('L')) 262 | 263 | M_homog, pts = processFrame(gray) 264 | 265 | if M_homog is None: 266 | print("Failed on image, not enough points") 267 | exit() 268 | 269 | pts_warped2 = cv2.perspectiveTransform(np.expand_dims(pts.astype(float),0), M_homog)[0,:,:] 270 | 271 | fig = plt.figure(figsize=(10,5)) 272 | ax = fig.add_subplot(1, 2, 1) 273 | 274 | # Major ticks every 20, minor ticks every 5 275 | major_ticks = np.arange(-2, 9, 2) 276 | minor_ticks = np.arange(-2, 9, 1) 277 | 278 | ax.axis('square') 279 | ax.set_xticks(major_ticks) 280 | ax.set_xticks(minor_ticks, minor=True) 281 | ax.set_yticks(major_ticks) 282 | ax.set_yticks(minor_ticks, minor=True) 283 | 284 | # And a corresponding grid 285 | ax.grid(which='both') 286 | 287 | # Or if you want different settings for the grids: 288 | ax.grid(which='minor', alpha=0.2) 289 | ax.grid(which='major', alpha=0.5) 290 | 291 | ideal_grid_pts = np.vstack([np.array([0,0,1,1,0])*8-1, np.array([0,1,1,0,0])*8-1]).T 292 | 293 | plt.plot(ideal_grid_pts[:,0], ideal_grid_pts[:,1], 'r:') 294 | plt.plot(pts_warped2[:,0], pts_warped2[:,1], 'r.') 295 | 296 | plt.subplot(122) 297 | plt.imshow(img) 298 | plt.plot(pts[:,0], pts[:,1], 'ro') 299 | 300 | # Refined via homography of all valid points 301 | unwarped_ideal_chess_corners_homography = cv2.perspectiveTransform( 302 | np.expand_dims(ideal_grid_pts.astype(float),0), np.linalg.inv(M_homog))[0,:,:] 303 | 304 | plt.plot(unwarped_ideal_chess_corners_homography[:,0], unwarped_ideal_chess_corners_homography[:,1], 'r-', lw=4); 305 | plt.show() -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- 1 | cff-version: 1.2.0 2 | message: "If you use this software, please cite it as below." 3 | authors: 4 | - family-names: "Ansari" 5 | given-names: "Sameer" 6 | orcid: "https://orcid.org/0000-0003-2096-5048" 7 | title: "ChessboardDetect" 8 | version: 1.0.0 9 | date-released: 2018-06-26 10 | url: "https://github.com/Elucidation/ChessboardDetect" 11 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2016 Sam 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Chessboard Detection 2 | 3 | *2024:* I've started a small youtube series explaining this repository: [Teaching AI To See](https://www.youtube.com/playlist?list=PLvPT8yiPYHne3qB84Vb2MujuIHcCTPSaW) 4 | 5 | *Note* : This repo is a staging ground of half-baked hacky code and different approaches to chessboard detection. Several different algorithms have been implemented each with varying tradeoffs. 6 | 7 | ![SpeedChessTrackedLK](sam2_composite.gif) 8 | ![Live Detection Speed](speedchess1_ml.gif) 9 | ![Live Detection](output_ml.gif) 10 | 11 | More Tracking videos: https://photos.app.goo.gl/QYBZH4PrmR1FKaUa9 12 | 13 | Desktop real-time demo (~30-200ms per frame) of the chessboard detector running, each frame is processed initially by the classifier, and once a match is found, Lucas-Kanade tracking is used to follow it onwards, occassionaly resetting if enough points are lost in the transition. It's written all in python (portions were written in C++/Halide but I've ended up not using them yet due to build issues). 14 | 15 | The algorithm is a combination of opencv for video capture and several basic vision algorithms, finding saddle points in the image, which are then classified using a Tensorflow DNNClassifier. After that all potential chess tile quads are used to warp the points onto an ideal unity grid and they are scored for grid-i-ness, a kind of brutesac (ransac without the random). Lots of opportunities for optimization to be had. 16 | 17 | Once the best scoring grid is found, we try and fit a chessboard to the points in unity grid space, then do one more refinement with opencv's findHomography and the final chessboard transform is shown as a red outline. 18 | 19 | ![Example Unity Grid](example_unity_grid.png) 20 | 21 | ## Goal 22 | 23 | * Given a photo with a chess board in it, find the chessboard grid. The goal is to make this fast enough to run real-time on an android phone. 24 | 25 | ![Labeled](readme_labeled.png) 26 | 27 | This is to be an evolution on [Tensorflow Chessbot](https://github.com/Elucidation/tensorflow_chessbot), working with real images. 28 | 29 | ## Algorithm #4 CNN Classifier + Lucas-Kanade tracking for consistency through video frames (~50-200ms per frame) 30 | 31 | Trained both a CNN and a DNN, found the CNN to work better. Trained with with ~12k tiles (6k each of true/false chess corners) gets us to roughly 96% success rate before over-training. Combining Algorithm #3 with a Lucas Kanade tracking of found chessboard points gives more consistent chessboards across frames. 32 | 33 | ![SpeedChessTrackedLK](speedchess1_composite.gif) 34 | 35 | More Tracking videos: https://photos.app.goo.gl/QYBZH4PrmR1FKaUa9 36 | 37 | Notes: It does not yet track off-by-one errors, this is a reasonably simple fix (TODO). 38 | 39 | The dataset can be found [here](data/). Please cite this repository if you use it, you can pull the citation from the about section. 40 | 41 | ## Algorithm #3 (DNNClassifier) (~100-200ms per image) 42 | 43 | One separate track is real-time chessboard classification on video feeds such as off of youtube videos. Using a combination of x-corner saddle detection and an ML DNN Classifier trained off of the previous algorithms on tiles of saddle points, we can find a triangle mesh for 'mostly' chessboard corners in realtime' (~20ms per 960x554 px frame). This is with python and opencv, the saddle point detection takes ~15ms, using a C++ Halide routine we've gotten this as low as 4ms, so there's lots of room for optimization in the future. 44 | 45 | ![Triangle Mesh](triangle_mesh.png) 46 | 47 | ## Algorithm #2 (~30 sec per image) 48 | 49 | 1. Find potential quad contours within the image 50 | 1. Grow out a grid of points from potential contours, vote for best match to saddle points 51 | 1. Warp image to grid, find set of 7 hough lines that maximize alternating chess-tile gradients 52 | 1. Build rectified chess image and overlay image with final transform 53 | 54 | 55 | Here are several results, 36 successes and 3 failures, red lines overlay the board outline and internal saddle points. 56 | 57 | ![Results](result.png) 58 | 59 | ## Algorithm #1 (~minute per image) 60 | 61 | Animation of several rectified boards that were found from images such as the one below 62 | 63 | ![Animated Rectified Images](readme_rectified.gif) 64 | 65 | ![Find chessboard and warp image](readme_find_warp_example.png) 66 | 67 | 1. Use Probabilistic Hough Transform to find lines 68 | 2. Prune lines based on strong alternating normal gradient angle frequency (checkerboard pattern) 69 | 3. Cluster line sets into segments, and choose top two corresponding to two axes of chessboard pattern 70 | 4. Find set of line intersections to define grid points 71 | 5. Take bounding corner grid points and perspective warp image 72 | 6. Re-center tile-map and refine corner points with cornerSubpix 73 | 7. Refine final transform with updated corners & rectify tile image 74 | 8. Correlate chessboard with tiled pattern, rotate 90 deg if orientation of dark/light tiles is off (A1 of H8 tiles must always be black in legal games, turns out a lot of stock images online don't follow this) 75 | 76 | ### Old Example 77 | 78 | We find the chessboard and warp the image 79 | 80 | ![Example warpimage](readme_output.png) 81 | 82 | *TODO, split warped image into tiles, predict chess pieces on tiles* 83 | 84 | ## Some constraints: 85 | 86 | * Chessboard can be populated with pieces, possibly partially occluded by hands etc. too 87 | * Background can be crowded 88 | * No user input besides the image taken 89 | 90 | ## Prior Work 91 | 92 | There exists checkerboard detection algorithms, but they assume no occlusions, consistent black/white pattern, and clean demarcation with the background. We may be dealing with textured/patterned surfaces, heavy occlusion due to pieces or people's hands, etc., and the background may interrupt to the edges of the chessboard. 93 | 94 | ## Current Work 95 | 96 | A promising approach is Gradient Angle Informed Hough transforms to find chessboard lines, once again taking advantage of the alternating gradient angles of the internal chessboard tiles. 97 | -------------------------------------------------------------------------------- /RunBrutesacOnVideo.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import os 3 | import glob 4 | import PIL.Image 5 | import cv2 6 | import skvideo.io 7 | import numpy as np 8 | import Brutesac 9 | 10 | def videostream(filename='carlsen_match.mp4', SAVE_FRAME=True): 11 | print("Loading video %s" % filename) 12 | vidstream = skvideo.io.vread(filename)#, num_frames=1000) 13 | print("Finished loading") 14 | print(vidstream.shape) 15 | 16 | # ffmpeg -i vidstream_frames/ml_frame_%03d.jpg -c:v libx264 -vf "fps=25,format=yuv420p" test.avi -y 17 | 18 | output_folder = "%s_vidstream_frames" % (filename[:-4]) 19 | if not os.path.exists(output_folder): 20 | os.mkdir(output_folder) 21 | 22 | for i, frame in enumerate(vidstream): 23 | # if i < 900: 24 | # continue 25 | print("Frame %d" % i) 26 | # if (i%5!=0): 27 | # continue 28 | 29 | # frame = cv2.resize(frame, (320,240), interpolation = cv2.INTER_CUBIC) 30 | 31 | # Our operations on the frame come here 32 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) 33 | gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) 34 | 35 | M_homog, pts = Brutesac.processFrame(gray) 36 | 37 | if M_homog is not None: 38 | ideal_grid_pts = np.vstack([np.array([0,0,1,1,0])*8-1, np.array([0,1,1,0,0])*8-1]).T 39 | unwarped_ideal_chess_corners_homography = cv2.perspectiveTransform( 40 | np.expand_dims(ideal_grid_pts.astype(float),0), np.linalg.inv(M_homog))[0,:,:] 41 | 42 | 43 | for pt in pts: 44 | cv2.circle(frame, tuple(pt), 3, (0,255,0), -1) 45 | 46 | # for pt in unwarped_ideal_chess_corners_homography: 47 | # cv2.circle(frame, tuple(pt[::-1]), 3, (0,0,255), -1) 48 | cv2.polylines(frame, [unwarped_ideal_chess_corners_homography.astype(np.int32)], isClosed=True, thickness=3, color=(0,0,255)) 49 | 50 | cv2.putText(frame, 'Frame %d' % i, (5,15), cv2.FONT_HERSHEY_PLAIN, 1.0,(255,255,255),0,cv2.LINE_AA) 51 | 52 | # Display the resulting frame 53 | cv2.imshow('frame',frame) 54 | output_filepath = '%s/ml_frame_%03d.jpg' % (output_folder, i) 55 | if SAVE_FRAME: 56 | cv2.imwrite(output_filepath, frame) 57 | 58 | if cv2.waitKey(1) & 0xFF == ord('q'): 59 | break 60 | 61 | # When everything done, release the capture 62 | # cap.release() 63 | cv2.destroyAllWindows() 64 | 65 | def main(): 66 | # filenames = glob.glob('input_bad/*') 67 | # filenames = glob.glob('input/img_*') filenames = sorted(filenames) 68 | # n = len(filenames) 69 | # filename = filenames[0] 70 | filename = 'input/img_01.jpg' 71 | 72 | print ("Processing %s" % (filename)) 73 | gray_img = PIL.Image.open(filename).convert('L').resize([600,400]) 74 | gray = np.array(gray_img) 75 | 76 | 77 | 78 | cv2.imshow('frame',gray) 79 | cv2.waitKey() 80 | 81 | print('Finished') 82 | 83 | if __name__ == '__main__': 84 | # main() 85 | # filename = 'carlsen_match.mp4' 86 | # filename = 'carlsen_match2.mp4' 87 | # filename = 'output.avi' 88 | # filename = 'output2.avi' 89 | # filename = 'random1.mp4' 90 | filename = 'speedchess1.mp4' 91 | # filename = 'match1.mp4' 92 | # filename = 'match2.mp4' 93 | videostream(filename, True) 94 | 95 | 96 | 97 | -------------------------------------------------------------------------------- /RunContourSac.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import os 3 | import cv2 4 | import PIL.Image 5 | import skvideo.io 6 | import numpy as np 7 | import Brutesac 8 | from functools import wraps 9 | import time 10 | from scipy.spatial import ConvexHull 11 | 12 | @Brutesac.timed 13 | def calculateOnFrame(gray): 14 | # and return M for chessboard from image 15 | pts = Brutesac.classifyImage(gray) 16 | 17 | # pts = np.loadtxt('example_pts.txt') 18 | pts = pts[:,[1,0]] # Switch rows/cols to x/y for plotting on an image 19 | 20 | # Get contours 21 | contours, hierarchy = getContours(gray, pts) 22 | 23 | # xcorner_map = np.zeros(gray.shape, dtype=np.uint8) 24 | # for pt in pts: 25 | # cv2.circle(xcorner_map, tuple(pt), 5, 1, -1) 26 | 27 | contours, hierarchy = pruneContours(contours, hierarchy, pts) 28 | 29 | return pts, contours 30 | 31 | @Brutesac.timed 32 | def simplifyContours(contours): 33 | for i in range(len(contours)): 34 | # Approximate contour and update in place 35 | contours[i] = cv2.approxPolyDP(contours[i],0.04*cv2.arcLength(contours[i],True),True) 36 | 37 | def updateCorners(contour, pts): 38 | new_contour = contour.copy() 39 | for i in range(len(contour)): 40 | cc,rr = contour[i,0,:] 41 | r = np.all(np.abs(pts - [cc,rr]) < 4, axis=1) 42 | closest_xpt = np.argwhere(r) 43 | if len(closest_xpt) > 0: 44 | new_contour[i,0,:] = pts[closest_xpt[0]][0] 45 | else: 46 | return [] 47 | return new_contour 48 | 49 | 50 | @Brutesac.timed 51 | def pruneContours(contours, hierarchy, xpts): 52 | new_contours = [] 53 | new_hierarchies = [] 54 | for i in range(len(contours)): 55 | cnt = contours[i] 56 | h = hierarchy[i] 57 | 58 | # Must be child 59 | if h[2] != -1: 60 | continue 61 | 62 | # Only rectangular contours allowed 63 | if len(cnt) != 4: 64 | continue 65 | 66 | # Only contours that fill an area of at least 8x8 pixels 67 | if cv2.contourArea(cnt) < 8*8: 68 | continue 69 | 70 | # if not is_square(cnt): 71 | # continue 72 | 73 | # TODO : Remove those where internal luma variance is greater than threshold 74 | 75 | cnt = updateCorners(cnt, xpts) 76 | # If not all saddle corners 77 | if len(cnt) != 4: 78 | continue 79 | 80 | new_contours.append(cnt) 81 | new_hierarchies.append(h) 82 | 83 | new_contours = np.array(new_contours) 84 | new_hierarchy = np.array(new_hierarchies) 85 | if len(new_contours) == 0: 86 | return new_contours, new_hierarchy 87 | 88 | # Prune contours below median area 89 | areas = [cv2.contourArea(c) for c in new_contours] 90 | mask = [areas >= np.median(areas)*0.25] and [areas <= np.median(areas)*2.0] 91 | new_contours = new_contours[mask] 92 | new_hierarchy = new_hierarchy[mask] 93 | return np.array(new_contours), np.array(new_hierarchy) 94 | 95 | @Brutesac.timed 96 | def getContours(gray, pts, iters=10): 97 | edges = cv2.Canny(gray, 20, 250) 98 | 99 | # Mask edges to only those in convex hull of points (dilated) 100 | if len(pts) >= 3: 101 | xcorner_mask = np.zeros(gray.shape, dtype=np.uint8) 102 | hull = ConvexHull(pts) 103 | hull_pts = pts[hull.vertices] 104 | xcorner_mask = cv2.fillConvexPoly(xcorner_mask, hull_pts, 255) 105 | # Dilate mask a bit 106 | element = np.ones([21, 21], np.uint8) 107 | xcorner_mask = cv2.dilate(xcorner_mask, element) 108 | 109 | edges = cv2.bitwise_and(edges,edges,mask = xcorner_mask) 110 | 111 | # Morphological Gradient to get internal squares of canny edges. 112 | kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3)) 113 | edges_gradient = cv2.morphologyEx(edges, cv2.MORPH_GRADIENT, kernel) 114 | _, contours, hierarchy = cv2.findContours(edges_gradient, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE) 115 | 116 | simplifyContours(contours,) 117 | 118 | return np.array(contours), hierarchy[0] 119 | 120 | @Brutesac.timed 121 | def contourSacChessboard(xcorner_pts, quads): 122 | # For each quad, keep track of the best fitting chessboard. 123 | best_score = 0 124 | best_error_score = None 125 | best_M = None 126 | best_quad = None 127 | best_offset = None 128 | for quad in quads: 129 | score, error_score, M, offset = Brutesac.scoreQuad(quad, xcorner_pts, best_score) 130 | if score > best_score or (score == best_score and error_score < best_error_score): 131 | best_score = score 132 | best_error_score = error_score 133 | best_M = M 134 | best_quad = quad 135 | best_offset = offset 136 | if best_score > (len(xcorner_pts)*0.9): 137 | break 138 | 139 | return best_M, best_quad, best_offset, best_score, best_error_score 140 | 141 | 142 | @Brutesac.timed 143 | def processFrame(frame, gray): 144 | pts, contours = calculateOnFrame(gray) 145 | 146 | raw_M, best_quad, best_offset, best_score, best_error_score = contourSacChessboard(pts, contours) 147 | if raw_M is not None: 148 | M_homog = Brutesac.refineHomography(pts, raw_M, best_offset) 149 | else: 150 | M_homog = None 151 | 152 | # Draw tiles found 153 | # cv2.drawContours(frame,contours,-1,(0,255,255),2) 154 | 155 | # Draw xcorner points 156 | for pt in pts: 157 | cv2.circle(frame, tuple(pt), 3, (0,0,255), -1) 158 | 159 | 160 | ideal_grid_pts = np.vstack([np.array([0,0,1,1,0])*8-1, np.array([0,1,1,0,0])*8-1]).T 161 | 162 | if M_homog is not None: 163 | # Refined via homography of all valid points 164 | unwarped_ideal_chess_corners_homography = cv2.perspectiveTransform( 165 | np.expand_dims(ideal_grid_pts.astype(float),0), np.linalg.inv(M_homog))[0,:,:] 166 | 167 | cv2.polylines(frame, 168 | [unwarped_ideal_chess_corners_homography.astype(np.int32)], 169 | isClosed=True, thickness=4, color=(0,0,255)) 170 | 171 | # if best_quad is not None: 172 | # cv2.polylines(frame, 173 | # [best_quad.astype(np.int32)], 174 | # isClosed=True, thickness=4, color=(255,0,255)) 175 | 176 | # Visualize mask used by getContours 177 | # if len(pts) >= 3: 178 | # xcorner_mask = np.zeros(gray.shape, dtype=np.uint8) 179 | # hull = ConvexHull(pts) 180 | # hull_pts = pts[hull.vertices] 181 | # xcorner_mask = cv2.fillConvexPoly(xcorner_mask, hull_pts, 255) 182 | # # Dilate mask a bit 183 | # element = np.ones([21, 21], np.uint8) 184 | # xcorner_mask = cv2.dilate(xcorner_mask, element) 185 | 186 | # frame = cv2.bitwise_and(frame,frame,mask = xcorner_mask) 187 | return frame 188 | 189 | # M_homog, pts = Brutesac.calculateOnFrame(gray) 190 | 191 | # if M_homog is not None: 192 | # ideal_grid_pts = np.vstack([np.array([0,0,1,1,0])*8-1, np.array([0,1,1,0,0])*8-1]).T 193 | # unwarped_ideal_chess_corners_homography = cv2.perspectiveTransform( 194 | # np.expand_dims(ideal_grid_pts.astype(float),0), np.linalg.inv(M_homog))[0,:,:] 195 | 196 | 197 | 198 | # # for pt in unwarped_ideal_chess_corners_homography: 199 | # # cv2.circle(frame, tuple(pt[::-1]), 3, (0,0,255), -1) 200 | # cv2.polylines(frame, [unwarped_ideal_chess_corners_homography.astype(np.int32)], isClosed=True, thickness=3, color=(0,0,255)) 201 | 202 | # cv2.putText(frame, 'Frame %d' % i, (5,15), cv2.FONT_HERSHEY_PLAIN, 1.0,(255,255,255),0,cv2.LINE_AA) 203 | 204 | def getWarpedChessboard(img, M, tile_px=32): 205 | # Given a the 4 points of a chessboard, get a warped image of just the chessboard 206 | 207 | # board_pts = np.vstack([ 208 | # np.array([0,0,1,1])*tile_px, 209 | # np.array([0,1,1,0])*tile_px 210 | # ]).T 211 | img_warp = cv2.warpPerspective(img, M, (8*tile_px, 8*tile_px)) 212 | return img_warp 213 | 214 | 215 | 216 | def videostream(filename='carlsen_match.mp4', SAVE_FRAME=True): 217 | print("Loading video %s" % filename) 218 | # vidstream = skvideo.io.vread(filename, num_frames=4000) 219 | # Load frame-by-frame 220 | vidstream = skvideo.io.vreader(filename) 221 | print("Finished loading") 222 | # print(vidstream.shape) 223 | 224 | # ffmpeg -i vidstream_frames/ml_frame_%03d.jpg -c:v libx264 -vf "fps=25,format=yuv420p" test.avi -y 225 | 226 | output_folder = "%s_vidstream_frames" % (filename[:-4]) 227 | if not os.path.exists(output_folder): 228 | os.mkdir(output_folder) 229 | 230 | for i, frame in enumerate(vidstream): 231 | # if i < 2000: 232 | # continue 233 | print("Frame %d" % i) 234 | # if (i%5!=0): 235 | # continue 236 | 237 | # frame = cv2.resize(frame, (320,240), interpolation = cv2.INTER_CUBIC) 238 | 239 | # Our operations on the frame come here 240 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) 241 | gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) 242 | 243 | # if i == 63: 244 | # cv2.imwrite('weird.png', frame) 245 | # break; 246 | 247 | a = time.time() 248 | frame = processFrame(frame, gray) 249 | t_proc = time.time() - a 250 | 251 | # Add frame counter 252 | cv2.putText(frame, 'Frame % 4d (Processed in % 6.1f ms)' % (i, t_proc*1e3), (5,15), cv2.FONT_HERSHEY_PLAIN, 1.0,(255,255,255),0) 253 | 254 | # Display the resulting frame 255 | cv2.imshow('frame',frame) 256 | output_filepath = '%s/ml_frame_%03d.jpg' % (output_folder, i) 257 | if SAVE_FRAME: 258 | cv2.imwrite(output_filepath, frame) 259 | 260 | if cv2.waitKey(1) & 0xFF == ord('q'): 261 | break 262 | 263 | # When everything done, release the capture 264 | # cap.release() 265 | cv2.destroyAllWindows() 266 | 267 | 268 | def main(): 269 | # filenames = glob.glob('input_bad/*') 270 | # filenames = glob.glob('input/img_*') filenames = sorted(filenames) 271 | # n = len(filenames) 272 | # filename = filenames[0] 273 | # filename = 'input/img_01.jpg' 274 | filename = 'weird.jpg' 275 | filename = 'chess_out1.png' 276 | 277 | print ("Processing %s" % (filename)) 278 | img = PIL.Image.open(filename).resize([600,400]) 279 | # img = PIL.Image.open(filename) 280 | rgb = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) 281 | gray = np.array(img.convert('L')) 282 | 283 | ### 284 | rgb = processFrame(rgb, gray) 285 | ### 286 | 287 | cv2.imshow('frame',rgb) 288 | cv2.waitKey() 289 | 290 | print('Finished') 291 | 292 | 293 | if __name__ == '__main__': 294 | # main() 295 | # filename = 'carlsen_match.mp4' 296 | # filename = 'carlsen_match2.mp4' 297 | # filename = 'output2.avi' 298 | # filename = 'random1.mp4' 299 | filename = 'match2.mp4' 300 | # filename = 'output.avi' 301 | # filename = 'speedchess1.mp4' 302 | # filename = 'chess_beer.mp4' 303 | # filename = 'john1.mp4' 304 | # filename = 'john2.mp4' 305 | # filename = 'john3.mp4' 306 | videostream(filename, False) 307 | 308 | 309 | 310 | -------------------------------------------------------------------------------- /RunExportedMLOnImage.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import PIL.Image 3 | import matplotlib.image as mpimg 4 | import scipy.ndimage 5 | import cv2 # For Sobel etc 6 | import glob 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | import os 10 | import time 11 | import sys 12 | import skvideo.io 13 | from functools import wraps 14 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 15 | 16 | from scipy.spatial import Delaunay 17 | 18 | def timed(f): 19 | @wraps(f) 20 | def wrapper(*args, **kwds): 21 | start = time.time() 22 | result = f(*args, **kwds) 23 | elapsed = time.time() - start 24 | print "%s took %.2f ms to finish" % (f.__name__, elapsed*1e3) 25 | return result 26 | return wrapper 27 | 28 | # export_dir = 'ml/model/001/1521934334' # old 29 | # export_dir = 'ml/model/002/1528405915' # newer (same dataset, but random image changes) 30 | # export_dir = 'ml/model/003/1528406613' # newer still 31 | # export_dir = 'ml/model/004/1528441286' # win 21x21, 95% accuracy 32 | # export_dir = 'ml/model/005/1528489968' # win 21x21 96% accuracy 33 | 34 | # export_dir = 'ml/model/006/1528565066' # win 21x21 97% accuracy 35 | # predict_fn = predictor.from_saved_model(export_dir, signature_def_key='predict') 36 | 37 | def getModel(export_dir='ml/model/006/1528565066'): 38 | from tensorflow.contrib import predictor 39 | return predictor.from_saved_model(export_dir, signature_def_key='predict') 40 | 41 | # Saddle 42 | def getSaddle(gray_img): 43 | img = gray_img#.astype(np.float64) 44 | gx = cv2.Sobel(img,cv2.CV_32F,1,0) 45 | gy = cv2.Sobel(img,cv2.CV_32F,0,1) 46 | gxx = cv2.Sobel(gx,cv2.CV_32F,1,0) 47 | gyy = cv2.Sobel(gy,cv2.CV_32F,0,1) 48 | gxy = cv2.Sobel(gx,cv2.CV_32F,0,1) 49 | 50 | # Inverse everything so positive equals more likely. 51 | S = -gxx*gyy + gxy**2 52 | 53 | # Calculate subpixel offsets 54 | denom = (gxx*gyy - gxy*gxy) 55 | sub_s = (gy*gxy - gx*gyy) / denom 56 | sub_t = (gx*gxy - gy*gxx) / denom 57 | return S, sub_s, sub_t 58 | 59 | def fast_nonmax_sup(img, win=11): 60 | element = np.ones([win, win], np.uint8) 61 | img_dilate = cv2.dilate(img, element) 62 | peaks = cv2.compare(img, img_dilate, cv2.CMP_EQ) 63 | img[peaks == 0] = 0 64 | 65 | 66 | # Deprecated for fast_nonmax_sup 67 | def nonmax_sup(img, win=10): 68 | w, h = img.shape 69 | # img = cv2.blur(img, ksize=(5,5)) 70 | img_sup = np.zeros_like(img, dtype=np.float64) 71 | for i,j in np.argwhere(img): 72 | # Get neigborhood 73 | ta=max(0,i-win) 74 | tb=min(w,i+win+1) 75 | tc=max(0,j-win) 76 | td=min(h,j+win+1) 77 | cell = img[ta:tb,tc:td] 78 | val = img[i,j] 79 | # if np.sum(cell.max() == cell) > 1: 80 | # print(cell.argmax()) 81 | if cell.max() == val: 82 | img_sup[i,j] = val 83 | return img_sup 84 | 85 | def pruneSaddle(s, init=128): 86 | thresh = init 87 | score = (s>0).sum() 88 | while (score > 10000): 89 | thresh = thresh*2 90 | s[s0).sum() 92 | 93 | 94 | def loadImage(filepath): 95 | img_orig = PIL.Image.open(filepath).convert('RGB') 96 | img_width, img_height = img_orig.size 97 | 98 | # Resize 99 | aspect_ratio = min(500.0/img_width, 500.0/img_height) 100 | new_width, new_height = ((np.array(img_orig.size) * aspect_ratio)).astype(int) 101 | img = img_orig.resize((new_width,new_height), resample=PIL.Image.BILINEAR) 102 | gray_img = img.convert('L') # grayscale 103 | img = np.array(img) 104 | gray_img = np.array(gray_img) 105 | 106 | return img, gray_img 107 | 108 | @timed 109 | def getFinalSaddlePoints(img, WINSIZE=10): # 32ms -> 15ms 110 | img = cv2.blur(img, (3,3)) # Blur it (.5ms) 111 | saddle, sub_s, sub_t = getSaddle(img) # 6ms 112 | fast_nonmax_sup(saddle) # ~6ms 113 | saddle[saddle<10000]=0 # Hardcoded ~1ms 114 | sub_idxs = np.nonzero(saddle) 115 | idxs = np.argwhere(saddle).astype(np.float64) 116 | spts = idxs[:,[1,0]] # Return in x,y order instead or row-col 117 | spts = spts + np.array([sub_s[sub_idxs], sub_t[sub_idxs]]).transpose() 118 | # Remove those points near win_size edges 119 | spts = clipBoundingPoints(spts, img.shape, WINSIZE) 120 | return spts # returns in x,y column order 121 | 122 | def clipBoundingPoints(pts, img_shape, WINSIZE=10): # ~100us 123 | # Points are given in x,y coords, not r,c of the image shape 124 | a = ~np.any(np.logical_or(pts <= WINSIZE, pts[:,[1,0]] >= np.array(img_shape)-WINSIZE-1), axis=1) 125 | return pts[a,:] 126 | 127 | def removeOutlierSimplices(tri): 128 | dists = np.zeros([tri.nsimplex, 3]) 129 | for i,triangle in enumerate(tri.points[tri.simplices]): 130 | # We want the distance of the edge opposite the vertex k, so r_k. 131 | r0 = (triangle[2,:] - triangle[1,:]) 132 | r1 = (triangle[0,:] - triangle[2,:]) 133 | r2 = (triangle[1,:] - triangle[0,:]) 134 | dists[i,:] = np.linalg.norm(np.vstack([r0,r1,r2]), axis=1) 135 | # Threshold based on twice the smallest edge. 136 | threshold = 2*np.median(dists) 137 | # Find edges that are too long 138 | long_edges = dists > threshold 139 | long_edged_simplices = np.any(long_edges,axis=1) 140 | # Good simplices are those that don't contain any long edges 141 | good_simplices_mask = ~long_edged_simplices 142 | # good_simplices = tri.simplices[good_simplices_mask] 143 | return dists, good_simplices_mask 144 | 145 | def findQuadSimplices(tri, dists, simplices_mask=None): 146 | vertex_idx_opposite_longest_edge = dists.argmax(axis=1) 147 | # The neighboring tri for each tri about the longest edge 148 | potential_neighbor = tri.neighbors[ 149 | np.arange(tri.nsimplex), 150 | vertex_idx_opposite_longest_edge] 151 | good_neighbors = [] 152 | for i,j in enumerate(potential_neighbor): 153 | if j == -1: 154 | # Skip those that don't have a neighbor 155 | continue 156 | # If these tris both agree that they're good neighbors, keep them. 157 | if (potential_neighbor[j] == i and i < j): 158 | if simplices_mask is not None: 159 | if simplices_mask[i]: 160 | good_neighbors.append(i) 161 | if simplices_mask[j]: 162 | good_neighbors.append(j) 163 | else: 164 | good_neighbors.extend([i,j]) 165 | return good_neighbors 166 | 167 | # def videostream(filename='carlsen_match.mp4', SAVE_FRAME=True): 168 | # # vidstream = skvideo.io.vread('VID_20170427_003836.mp4') 169 | # # vidstream = skvideo.io.vread('VID_20170109_183657.mp4') 170 | # print("Loading video %s" % filename) 171 | # # vidstream = skvideo.io.vread('output2.avi') 172 | # vidstream = skvideo.io.vread(filename)#, num_frames=1000) 173 | # # vidstream = skvideo.io.vread('output.avi') 174 | # print("Finished loading") 175 | # # vidstream = skvideo.io.vread(0) 176 | # print(vidstream.shape) 177 | 178 | # # ffmpeg -i vidstream_frames/ml_frame_%03d.jpg -c:v libx264 -vf "fps=25,format=yuv420p" test.avi -y 179 | 180 | # output_folder = "%s_vidstream_frames" % (filename[:-4]) 181 | # if not os.path.exists(output_folder): 182 | # os.mkdir(output_folder) 183 | 184 | 185 | # for i, frame in enumerate(vidstream): 186 | # print("Frame %d" % i) 187 | # # frame = cv2.resize(frame, (320,240), interpolation = cv2.INTER_CUBIC) 188 | 189 | # # Our operations on the frame come here 190 | # frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) 191 | # gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) 192 | 193 | # inlier_pts, outlier_pts, pred_pts, final_predictions, prediction_levels, tri, simplices_mask = processImage(gray) 194 | 195 | 196 | # for pt in inlier_pts: 197 | # cv2.circle(frame, tuple(pt[::-1]), 3, (0,255,0), -1) 198 | 199 | # for pt in outlier_pts: 200 | # cv2.circle(frame, tuple(pt[::-1]), 1, (0,0,255), -1) 201 | 202 | # # Draw triangle mesh 203 | # if tri is not None: 204 | # cv2.polylines(frame, 205 | # np.flip(inlier_pts[tri.simplices].astype(np.int32), axis=2), 206 | # isClosed=True, color=(255,0,0)) 207 | # cv2.polylines(frame, 208 | # np.flip(inlier_pts[tri.simplices[simplices_mask]].astype(np.int32), axis=2), 209 | # isClosed=True, color=(0,255,0), thickness=3) 210 | 211 | # cv2.putText(frame, 'Frame %d' % i, (5,15), cv2.FONT_HERSHEY_PLAIN, 1.0,(255,255,255),0,cv2.LINE_AA) 212 | 213 | # # Display the resulting frame 214 | # cv2.imshow('frame',frame) 215 | # output_filepath = '%s/ml_frame_%03d.jpg' % (output_folder, i) 216 | # if SAVE_FRAME: 217 | # cv2.imwrite(output_filepath, frame) 218 | 219 | # if cv2.waitKey(1) & 0xFF == ord('q'): 220 | # break 221 | 222 | # # When everything done, release the capture 223 | # cv2.destroyAllWindows() 224 | 225 | # def calculateOutliers(pts, threshold_mult = 2.5): 226 | # N = len(pts) 227 | # std = np.std(pts, axis=0) 228 | # ctr = np.mean(pts, axis=0) 229 | # return (np.any(np.abs(pts-ctr) > threshold_mult * std, axis=1)) 230 | 231 | 232 | 233 | 234 | -------------------------------------------------------------------------------- /RunMLOnImage.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import PIL.Image 3 | import matplotlib.image as mpimg 4 | import scipy.ndimage 5 | import cv2 # For Sobel etc 6 | import glob 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | import os 10 | import tensorflow as tf 11 | import time 12 | import sys 13 | import skvideo.io 14 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 15 | 16 | # Set up model 17 | featureA = tf.feature_column.numeric_column("x", shape=[11,11], dtype=tf.uint8) 18 | 19 | estimator = tf.estimator.DNNClassifier( 20 | feature_columns=[featureA], 21 | hidden_units=[256, 32], 22 | n_classes=2, 23 | dropout=0.1, 24 | model_dir='./xcorner_model_6k', 25 | ) 26 | 27 | # Saddle 28 | 29 | def getSaddle(gray_img): 30 | img = gray_img.astype(np.float64) 31 | gx = cv2.Sobel(img,cv2.CV_64F,1,0) 32 | gy = cv2.Sobel(img,cv2.CV_64F,0,1) 33 | gxx = cv2.Sobel(gx,cv2.CV_64F,1,0) 34 | gyy = cv2.Sobel(gy,cv2.CV_64F,0,1) 35 | gxy = cv2.Sobel(gx,cv2.CV_64F,0,1) 36 | 37 | S = gxx*gyy - gxy**2 38 | return S 39 | 40 | 41 | # void nonmaxSupress(Mat &img) { 42 | # int dilation_size = 5; 43 | # Mat img_dilate; 44 | # Mat peaks; 45 | # Mat notPeaks; 46 | # Mat nonzeroImg; 47 | # Mat element = getStructuringElement(MORPH_RECT, 48 | # Size( 2*dilation_size + 1, 2*dilation_size+1 ), 49 | # Point( dilation_size, dilation_size ) ); 50 | # // Dilate max value by window size 51 | # dilate(img, img_dilate, element); 52 | # // Compare and find where values of dilated vs original image are NOT the same. 53 | # compare(img, img_dilate, peaks, CMP_EQ); 54 | # // compare(img, img_dilate, notPeaks, CMP_NE); 55 | # compare(img, 0, nonzeroImg, CMP_NE); 56 | # bitwise_and(nonzeroImg, peaks, peaks); // Only keep peaks that are non-zero 57 | 58 | # // Remove peaks that are zero 59 | # // Also set max to 255 60 | # // compare(img, 0.8, nonzeroImg, CMP_GT); 61 | # // bitwise_and(nonzeroImg, peaks, peaks); // Only keep peaks that are non-zero 62 | # // bitwise_not(peaks, notPeaks); 63 | # // Set all values where not the same to zero. Non-max suppress. 64 | # bitwise_not(peaks, notPeaks); 65 | # img.setTo(0, notPeaks); 66 | # // img.setTo(255, peaks); 67 | 68 | # } 69 | 70 | def fast_nonmax_sup(img, win=21): 71 | element = np.ones([win, win], np.uint8) 72 | img_dilate = cv2.dilate(img, element) 73 | peaks = cv2.compare(img, img_dilate, cv2.CMP_EQ) 74 | # nonzeroImg = cv2.compare(img, 0, cv2.CMP_NE) 75 | # peaks = cv2.bitwise_and(peaks, nonzeroImg) 76 | peaks[img == 0] = 0 77 | # notPeaks = cv2.bitwise_not(peaks) 78 | 79 | img[peaks == 0] = 0 80 | return img 81 | 82 | 83 | 84 | def nonmax_sup(img, win=10): 85 | w, h = img.shape 86 | # img = cv2.blur(img, ksize=(5,5)) 87 | img_sup = np.zeros_like(img, dtype=np.float64) 88 | for i,j in np.argwhere(img): 89 | # Get neigborhood 90 | ta=max(0,i-win) 91 | tb=min(w,i+win+1) 92 | tc=max(0,j-win) 93 | td=min(h,j+win+1) 94 | cell = img[ta:tb,tc:td] 95 | val = img[i,j] 96 | # if np.sum(cell.max() == cell) > 1: 97 | # print(cell.argmax()) 98 | if cell.max() == val: 99 | img_sup[i,j] = val 100 | return img_sup 101 | 102 | def pruneSaddle(s): 103 | thresh = 128 104 | score = (s>0).sum() 105 | while (score > 10000): 106 | thresh = thresh*2 107 | s[s0).sum() 109 | 110 | 111 | def loadImage(filepath): 112 | img_orig = PIL.Image.open(filepath) 113 | img_width, img_height = img_orig.size 114 | 115 | # Resize 116 | aspect_ratio = min(500.0/img_width, 500.0/img_height) 117 | new_width, new_height = ((np.array(img_orig.size) * aspect_ratio)).astype(int) 118 | img = img_orig.resize((new_width,new_height), resample=PIL.Image.BILINEAR) 119 | gray_img = img.convert('L') # grayscale 120 | img = np.array(img) 121 | gray_img = np.array(gray_img) 122 | 123 | return img, gray_img 124 | 125 | def getFinalSaddlePoints(img): 126 | blur_img = cv2.blur(img, (3,3)) # Blur it 127 | saddle = getSaddle(blur_img) 128 | saddle = -saddle 129 | saddle[saddle<0] = 0 130 | pruneSaddle(saddle) 131 | s2 = fast_nonmax_sup(saddle) 132 | s2[s2<100000]=0 133 | spts = np.argwhere(s2) 134 | return spts 135 | 136 | # def processSingle(filename='input/img_10.png'): 137 | # img = loadImage(filename) 138 | # spts = getFinalSaddlePoints(img) 139 | def input_fn_predict(img_data): # returns x, None 140 | def ret_func(): 141 | dataset = tf.data.Dataset.from_tensor_slices( 142 | { 143 | 'x':img_data 144 | } 145 | ) 146 | dataset = dataset.batch(25) 147 | return dataset.make_one_shot_iterator().get_next(), None 148 | return ret_func 149 | 150 | def videostream(): 151 | # vidstream = skvideo.io.vread('VID_20170427_003836.mp4') 152 | # vidstream = skvideo.io.vread('VID_20170109_183657.mp4') 153 | print("Loading video") 154 | # vidstream = skvideo.io.vread('output2.avi') 155 | vidstream = skvideo.io.vread('output.avi') 156 | print("Finished loading") 157 | # vidstream = skvideo.io.vread(0) 158 | print(vidstream.shape) 159 | 160 | output_folder = "vidstream_frames" 161 | if not os.path.exists(output_folder): 162 | os.mkdir(output_folder) 163 | 164 | # cap.set(3,320) 165 | # cap.set(4,240) 166 | 167 | # while(True): 168 | # Capture frame-by-frame 169 | # ret, frame = cap.read() 170 | 171 | # if not ret: 172 | # print("No frame, stopping") 173 | # break 174 | for i, frame in enumerate(vidstream): 175 | # if (i%2!=0): 176 | # continue 177 | 178 | # frame = cv2.resize(frame, (320,240), interpolation = cv2.INTER_CUBIC) 179 | 180 | # Our operations on the frame come here 181 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) 182 | gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) 183 | 184 | final_predictions, pred_pts = processImage(gray) 185 | 186 | for c,pt in zip(final_predictions, pred_pts): 187 | if (c == 1): 188 | # Good 189 | cv2.circle(frame, tuple(pt[::-1]), 4, (0,255,0), -1) 190 | else: 191 | cv2.circle(frame, tuple(pt[::-1]), 2, (0,0,255), -1) 192 | 193 | cv2.putText(frame, 'Frame %d' % i, (5,15), cv2.FONT_HERSHEY_PLAIN, 1.0,(255,255,255),0,cv2.LINE_AA) 194 | 195 | # Display the resulting frame 196 | cv2.imshow('frame',frame) 197 | output_filepath = '%s/ml_frame_%03d.jpg' % (output_folder, i) 198 | cv2.imwrite(output_filepath, frame) 199 | 200 | if cv2.waitKey(1) & 0xFF == ord('q'): 201 | break 202 | 203 | # When everything done, release the capture 204 | # cap.release() 205 | cv2.destroyAllWindows() 206 | 207 | def processImage(img_gray): 208 | a = time.time() 209 | spts = getFinalSaddlePoints(img_gray) 210 | b = time.time() 211 | print("getFinalSaddlePoints() took %.2f ms" % ((b-a)*1e3)) 212 | WINSIZE = 5 213 | 214 | tiles = [] 215 | pred_pts = [] 216 | for pt in spts: 217 | # Build tiles 218 | if (np.any(pt <= WINSIZE) or np.any(pt >= np.array(img_gray.shape[:2]) - WINSIZE)): 219 | continue 220 | else: 221 | tile = img_gray[pt[0]-WINSIZE:pt[0]+WINSIZE+1, pt[1]-WINSIZE:pt[1]+WINSIZE+1] 222 | tiles.append(tile) 223 | pred_pts.append(pt) 224 | tiles = np.array(tiles, dtype=np.uint8) 225 | 226 | a = time.time() 227 | predictions = estimator.predict(input_fn=input_fn_predict(tiles)) 228 | 229 | good_pts = [] 230 | bad_pts = [] 231 | final_predictions = [] 232 | for i, prediction in enumerate(predictions): 233 | c = prediction['probabilities'].argmax() 234 | pt = pred_pts[i] 235 | final_predictions.append(c) 236 | b = time.time() 237 | print("predict() took %.2f ms" % ((b-a)*1e3)) 238 | return final_predictions, pred_pts 239 | 240 | 241 | 242 | 243 | def main(): 244 | filenames = glob.glob('input_bad/*') 245 | # filenames = glob.glob('input/img_*') 246 | # filenames.extend(glob.glob('input_yt/*.jpg')) 247 | filenames = sorted(filenames) 248 | n = len(filenames) 249 | n = 1 250 | 251 | WINSIZE = 5 252 | 253 | for i in range(n): 254 | filename = filenames[i] 255 | print ("Processing %d/%d : %s" % (i+1,n,filename)) 256 | img, img_gray = loadImage(filename) 257 | final_predictions, pred_pts = processImage(img_gray) 258 | 259 | b,g,r = cv2.split(img) # get b,g,r 260 | rgb_img = cv2.merge([r,g,b]) # switch it to rgb 261 | 262 | for c,pt in zip(final_predictions, pred_pts): 263 | if (c == 1): 264 | # Good 265 | cv2.circle(rgb_img, tuple(pt[::-1]), 4, (0,255,0), -1) 266 | else: 267 | cv2.circle(rgb_img, tuple(pt[::-1]), 3, (0,0,255), -1) 268 | 269 | cv2.imshow('frame',rgb_img) 270 | if cv2.waitKey() & 0xFF == ord('q'): 271 | break 272 | 273 | print('Finished') 274 | 275 | 276 | 277 | 278 | 279 | 280 | if __name__ == '__main__': 281 | # main() 282 | videostream() 283 | 284 | 285 | 286 | -------------------------------------------------------------------------------- /SaddlePoints.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Saddle point calculator. 3 | import cv2 # For Sobel etc 4 | import numpy as np 5 | 6 | 7 | # Saddle 8 | def getSaddle(gray_img): 9 | img = gray_img 10 | gx = cv2.Sobel(img,cv2.CV_32F,1,0) 11 | gy = cv2.Sobel(img,cv2.CV_32F,0,1) 12 | gxx = cv2.Sobel(gx,cv2.CV_32F,1,0) 13 | gyy = cv2.Sobel(gy,cv2.CV_32F,0,1) 14 | gxy = cv2.Sobel(gx,cv2.CV_32F,0,1) 15 | 16 | # Inverse everything so positive equals more likely. 17 | S = -gxx*gyy + gxy**2 18 | 19 | # Calculate subpixel offsets 20 | denom = (gxx*gyy - gxy*gxy) 21 | sub_s = np.divide(gy*gxy - gx*gyy, denom, out=np.zeros_like(denom), where=denom!=0) 22 | sub_t = np.divide(gx*gxy - gy*gxx, denom, out=np.zeros_like(denom), where=denom!=0) 23 | return S, sub_s, sub_t, gx, gy 24 | 25 | def fast_nonmax_sup(img, win=11): 26 | element = np.ones([win, win], np.uint8) 27 | img_dilate = cv2.dilate(img, element) 28 | peaks = cv2.compare(img, img_dilate, cv2.CMP_EQ) 29 | img[peaks == 0] = 0 30 | 31 | def getFinalSaddlePoints(img, WINSIZE=10): # 32ms -> 15ms 32 | # Get all saddle points that are not closer than WINSIZE to the boundaries. 33 | 34 | img = cv2.blur(img, (3,3)) # Blur it (.5ms) 35 | saddle, sub_s, sub_t, gx, gy = getSaddle(img) # 6ms 36 | fast_nonmax_sup(saddle) # ~6ms 37 | 38 | # Strip off low points 39 | saddle[saddle<10000]=0 # Hardcoded ~1ms 40 | sub_idxs = np.nonzero(saddle) 41 | spts = np.argwhere(saddle).astype(np.float64)[:,[1,0]] # Return in x,y order instead or row-col 42 | 43 | # Add on sub-pixel offsets 44 | subpixel_offset = np.array([sub_s[sub_idxs], sub_t[sub_idxs]]).transpose() 45 | spts = spts + subpixel_offset 46 | 47 | # Remove those points near win_size edges 48 | spts = clipBoundingPoints(spts, img.shape, WINSIZE) 49 | 50 | return spts, gx, gy # returns in x,y column order 51 | 52 | def clipBoundingPoints(pts, img_shape, WINSIZE=10): # ~100us 53 | # Points are given in x,y coords, not r,c of the image shape 54 | a = ~np.any(np.logical_or(pts <= WINSIZE, pts[:,[1,0]] >= np.array(img_shape)-WINSIZE-1), axis=1) 55 | return pts[a,:] -------------------------------------------------------------------------------- /base_imgload.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from pylab import * 3 | import numpy as np 4 | import cv2 5 | import sys 6 | 7 | 8 | def main(filenames): 9 | for filename in filenames: 10 | print("Processing %s" % filename) 11 | img = cv2.imread(filename) 12 | gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY) 13 | 14 | # Using opencv 15 | cv2.imshow('image %dx%d' % (img.shape[1],img.shape[0]),img) 16 | cv2.waitKey(0) 17 | cv2.destroyAllWindows() 18 | 19 | if __name__ == '__main__': 20 | if len(sys.argv) > 1: 21 | filenames = sys.argv[1:] 22 | else: 23 | filenames = ['input/6.jpg'] 24 | main(filenames) 25 | -------------------------------------------------------------------------------- /buildMLSavedModel.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import glob 3 | import numpy as np 4 | import os 5 | import tensorflow as tf 6 | import sys 7 | from tensorflow.python.saved_model import signature_constants 8 | from tensorflow.python.saved_model import signature_def_utils 9 | from tensorflow.python.saved_model import tag_constants 10 | from tensorflow.python.saved_model import utils 11 | 12 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 13 | 14 | export_dir = 'ml/model/001' 15 | 16 | # builder = tf.saved_model.builder.SavedModelBuilder(export_dir) 17 | with tf.Session(graph=tf.Graph()) as sess: 18 | # Set up model 19 | featureA = tf.feature_column.numeric_column("x", shape=[11,11], dtype=tf.uint8) 20 | 21 | estimator = tf.estimator.DNNClassifier( 22 | feature_columns=[featureA], 23 | hidden_units=[256, 32], 24 | n_classes=2, 25 | dropout=0.1, 26 | model_dir='./xcorner_model_6k', 27 | ) 28 | 29 | def serving_input_receiver_fn(): 30 | """Build the serving inputs.""" 31 | # The outer dimension (None) allows us to batch up inputs for 32 | # efficiency. However, it also means that if we want a prediction 33 | # for a single instance, we'll need to wrap it in an outer list. 34 | inputs = {"x": tf.placeholder(shape=[None, 11, 11], dtype=tf.uint8)} 35 | return tf.estimator.export.ServingInputReceiver(inputs, inputs) 36 | 37 | estimator.export_savedmodel(export_dir, serving_input_receiver_fn) 38 | 39 | # prediction_signature = signature_def_utils.build_signature_def( 40 | # inputs={'x': tensor_info_x}, 41 | # outputs={'output': tensor_info_y}, 42 | # method_name=signature_constants.PREDICT_METHOD_NAME) 43 | 44 | # builder.add_meta_graph_and_variables(sess, 45 | # [tag_constants.TRAINING], 46 | # signature_def_map= 47 | # { 48 | # 'predict_images': 49 | # prediction_signature, 50 | # signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: 51 | # classification_signature, 52 | # }) 53 | 54 | # # Add a MetaGraphDef for inference. 55 | # with tf.Session(graph=tf.Graph()) as sess: 56 | # builder.add_meta_graph([tag_constants.SERVING]) 57 | 58 | # builder.save() -------------------------------------------------------------------------------- /buildMLSavedModel2.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import glob 3 | import numpy as np 4 | import os 5 | import tensorflow as tf 6 | import sys 7 | from tensorflow.python.saved_model import signature_constants 8 | from tensorflow.python.saved_model import signature_def_utils 9 | from tensorflow.python.saved_model import tag_constants 10 | from tensorflow.python.saved_model import utils 11 | 12 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 13 | 14 | input_model_dir = './training_pipeline/training_models/run1_512_256_128_orig_5k_dataset_10000' 15 | export_dir = 'ml/model/run97pct' 16 | 17 | # builder = tf.saved_model.builder.SavedModelBuilder(export_dir) 18 | with tf.Session(graph=tf.Graph()) as sess: 19 | # Set up model 20 | feature_img = tf.feature_column.numeric_column("x", shape=[21,21], dtype=tf.uint8) 21 | 22 | # units = [1024,512,256] 23 | units = [512,256,128] 24 | estimator = tf.estimator.DNNClassifier( 25 | feature_columns=[feature_img], 26 | hidden_units=units, 27 | n_classes=2, 28 | model_dir=input_model_dir 29 | ) 30 | 31 | def serving_input_receiver_fn(): 32 | """Build the serving inputs.""" 33 | # The outer dimension (None) allows us to batch up inputs for 34 | # efficiency. However, it also means that if we want a prediction 35 | # for a single instance, we'll need to wrap it in an outer list. 36 | inputs = {"x": tf.placeholder(shape=[None, 21, 21], dtype=tf.uint8)} 37 | return tf.estimator.export.ServingInputReceiver(inputs, inputs) 38 | 39 | estimator.export_savedmodel(export_dir, serving_input_receiver_fn) -------------------------------------------------------------------------------- /centralSymmetryTile.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import PIL.Image 3 | import matplotlib.image as mpimg 4 | import scipy.ndimage 5 | import cv2 # For Sobel etc 6 | import glob 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | from random import shuffle 10 | import os 11 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 12 | 13 | 14 | def getRingIndices(radius): 15 | # Bottom 16 | row1 = np.ones(radius*2+1, dtype=int)*radius 17 | col1 = np.arange(radius*2+1)-radius 18 | 19 | # Right 20 | row2 = -np.arange(1,radius*2+1)+radius 21 | col2 = np.ones(radius*2, dtype=int)*radius 22 | 23 | # Top 24 | row3 = -np.ones(radius*2, dtype=int)*radius 25 | col3 = -np.arange(1,radius*2+1)+radius 26 | 27 | # Left 28 | row4 = np.arange(1,radius*2+1-1)-radius 29 | col4 = -np.ones(radius*2-1, dtype=int)*radius 30 | 31 | rows = np.hstack([row1, row2, row3, row4]) 32 | cols = np.hstack([col1, col2, col3, col4]) 33 | return (rows,cols) 34 | 35 | def countSteps(ring): 36 | # Build a big ring so we can handle circular edges 37 | bigring = np.hstack([ring,ring,ring]) 38 | n = len(ring) 39 | # Go through middle portion of ring 40 | count = 0 41 | for i in (np.arange(n) + n): 42 | if (bigring[i] != bigring[i-1] and (bigring[i-1] == bigring[i-2]) and (bigring[i] == bigring[i+1])): 43 | count += 1 44 | return count 45 | 46 | 47 | 48 | # Load a tile image and check the central symmetry around a ring 49 | def main(): 50 | bad_tile_filepaths = sorted(glob.glob('dataset_binary_5/bad/img_*.png')) 51 | good_tile_filepaths = sorted(glob.glob('dataset_binary_5/good/img_*.png')) 52 | 53 | # shuffle(bad_tile_filepaths) 54 | # shuffle(good_tile_filepaths) 55 | 56 | # Setup 57 | tile_radius = (PIL.Image.open(good_tile_filepaths[0]).size[0]-1)/2 #(img.shape[0]-1)/2 58 | radius = 5 59 | 60 | # filepath = 'dataset_binary_5/bad/img_01_008.png' 61 | # plt.figure(figsize=(20,20)) 62 | # plt.subplot(121) 63 | # plt.title('False Positives') 64 | rows, cols = getRingIndices(radius) 65 | # Center in tile 66 | rows += tile_radius 67 | cols += tile_radius 68 | 69 | # for i in range(20): 70 | # filepath = bad_tile_filepaths[i] 71 | # img = PIL.Image.open(filepath).convert('L') 72 | # img = np.array(img) 73 | # # img = cv2.adaptiveThreshold(img,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2) 74 | # ring = img[rows,cols] 75 | # plt.plot(ring + i*255*2, '.-') 76 | # plt.plot([0,len(ring)-1], np.ones(2) + 127 + i*255*2, 'k:', alpha=0.2) 77 | # plt.text(0, i*255*2, countSteps(ring)) 78 | 79 | # # Good tiles 80 | # plt.subplot(122) 81 | # plt.title('True Positives') 82 | # for i in range(20): 83 | # filepath = good_tile_filepaths[i] 84 | # img = PIL.Image.open(filepath).convert('L') 85 | # img = np.array(img) 86 | # # img = cv2.adaptiveThreshold(img,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2) 87 | # ring = img[rows,cols] 88 | # plt.plot(ring + i*255*2, '.-') 89 | # plt.plot([0,len(ring)-1], np.ones(2) + 127 + i*255*2, 'k:', alpha=0.2) 90 | # plt.text(0, i*255*2, countSteps(ring)) 91 | 92 | # plt.show() 93 | 94 | good_steps = [] 95 | bad_steps = [] 96 | for i in range(len(bad_tile_filepaths)): 97 | filepath = bad_tile_filepaths[i] 98 | img = PIL.Image.open(filepath).convert('L') 99 | img = np.array(img) 100 | ring = img[rows,cols] 101 | steps = countSteps(ring) 102 | bad_steps.append(steps) 103 | for i in range(len(good_tile_filepaths)): 104 | filepath = good_tile_filepaths[i] 105 | img = PIL.Image.open(filepath).convert('L') 106 | img = np.array(img) 107 | ring = img[rows,cols] 108 | steps = countSteps(ring) 109 | good_steps.append(steps) 110 | 111 | # print(good_steps) 112 | # print(bad_steps) 113 | 114 | plt.subplot(121) 115 | plt.hist(bad_steps) 116 | plt.title('False Positives') 117 | plt.subplot(122) 118 | plt.hist(good_steps) 119 | plt.title('True Positives') 120 | plt.show() 121 | 122 | 123 | 124 | 125 | 126 | if __name__ == '__main__': 127 | main() 128 | 129 | 130 | 131 | -------------------------------------------------------------------------------- /chess_detect_helper.py: -------------------------------------------------------------------------------- 1 | import cv2 # For Sobel etc 2 | import numpy as np 3 | from helpers import * 4 | from line_intersection import * 5 | 6 | def getChessLinesCorners(img, chessboard_to_screen_ratio=0.2): 7 | # Edges 8 | edges = cv2.Canny(img,200,500,apertureSize = 3, L2gradient=False) # Better thresholds 9 | 10 | # Gradients 11 | sobelx = cv2.Sobel(img,cv2.CV_64F,1,0,ksize=5) 12 | sobely = cv2.Sobel(img,cv2.CV_64F,0,1,ksize=5) 13 | grad_mag = np.sqrt(sobelx**2+sobely**2) 14 | 15 | # Hough Lines Probabilistic 16 | # chessboard_to_screen_ratio = 0.2 17 | min_chessboard_line_length = chessboard_to_screen_ratio * min(img.shape) 18 | # TODO: This varys based on the chessboard to screen ratio, for chessboards filling the screen, we want to hop further 19 | max_line_gap = min_chessboard_line_length / 8.0 * 1.5 # Can hop up to one missing square 20 | 21 | lines = cv2.HoughLinesP(edges,1,np.pi/360.0, 30, minLineLength = min_chessboard_line_length, maxLineGap = max_line_gap)[:,0,:] 22 | # lines = cv2.HoughLinesP(edges,1,np.pi/360.0, 30, minLineLength=30, maxLineGap=30)[:,0,:] 23 | 24 | good_lines = np.zeros(lines.shape[0]) 25 | norm_grads = np.zeros(lines.shape[0]) 26 | angles = np.zeros(lines.shape[0]) 27 | 28 | for idx in range(lines.shape[0]): 29 | line = lines[idx,:] 30 | is_good, _, _, _, _, avg_normal_gradient = getLineGradients(line, sobelx, sobely, grad_mag) 31 | good_lines[idx] = is_good 32 | norm_grads[idx] = avg_normal_gradient 33 | angles[idx] = getSegmentAngle(line) 34 | 35 | # Get angles and segment lines up 36 | segments = segmentAngles(angles, good_lines, angle_threshold=15*np.pi/180) 37 | top_two_segments = chooseBestSegments(segments, norm_grads) 38 | if (top_two_segments.size < 2): 39 | print("Couldn't find enough segments") 40 | return [], [], [], [] 41 | 42 | # Update good_mask to only include top two groups 43 | a_segment = segments == top_two_segments[0] 44 | b_segment = segments == top_two_segments[1] 45 | good_mask = a_segment | b_segment 46 | 47 | a_segment_idxs = np.argwhere(a_segment).flatten() 48 | b_segment_idxs = np.argwhere(b_segment).flatten() 49 | 50 | # Rotate angles 45 degrees first 51 | angle_a = np.mean(angles[a_segment_idxs]) + np.pi/4 52 | dir_a = np.array([np.cos(angle_a), np.sin(angle_a)]) 53 | angle_b = np.mean(angles[b_segment_idxs]) + np.pi/4 54 | dir_b = np.array([np.cos(angle_b), np.sin(angle_b)]) 55 | 56 | # Plot intersections 57 | chess_pts = getAllLineIntersections(lines[a_segment_idxs], lines[b_segment_idxs]) 58 | pruned_chess_pts = prunePoints(chess_pts,max_dist2=5**2) 59 | return lines[a_segment_idxs], lines[b_segment_idxs], pruned_chess_pts, (dir_a, dir_b) -------------------------------------------------------------------------------- /chessboard_detect2.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import cv2 3 | import PIL.Image 4 | import numpy as np 5 | import sys 6 | from time import time 7 | from matplotlib import pyplot as plt 8 | from contour_detect import * 9 | from line_intersection import * 10 | from rectify_refine import * 11 | 12 | np.set_printoptions(suppress=True, precision=2, linewidth=200) 13 | 14 | def processFile(filename): 15 | img = cv2.imread(filename) 16 | # img = scaleImageIfNeeded(img, 600, 480) 17 | img = scaleImageIfNeeded(img, 1024, 768) 18 | img_orig = img.copy() 19 | img_orig2 = img.copy() 20 | 21 | # Edges 22 | edges = cv2.Canny(img, 100, 550) 23 | 24 | # Get mask for where we think chessboard is 25 | mask, top_two_angles, min_area_rect, median_contour = getEstimatedChessboardMask(img, edges,iters=3) # More iters gives a finer mask 26 | print("Top two angles (in image coord system): %s" % top_two_angles) 27 | 28 | # Get hough lines of masked edges 29 | edges_masked = cv2.bitwise_and(edges,edges,mask = (mask > 0.5).astype(np.uint8)) 30 | img_orig = cv2.bitwise_and(img_orig,img_orig,mask = (mask > 0.5).astype(np.uint8)) 31 | 32 | lines = getHoughLines(edges_masked, min_line_size=0.25*min(min_area_rect[1])) 33 | print("Found %d lines." % len(lines)) 34 | 35 | lines_a, lines_b = parseHoughLines(lines, top_two_angles, angle_threshold_deg=35) 36 | 37 | # plotHoughLines(img, lines, color=(255,255,255), line_thickness=1) 38 | # plotHoughLines(img, lines_a, color=(0,0,255)) 39 | # plotHoughLines(img, lines_b, color=(0,255,0)) 40 | if len(lines_a) < 2 or len(lines_b) < 2: 41 | return img_orig, edges_masked, img_orig 42 | 43 | a = time() 44 | for i2 in range(10): 45 | for i in range(100): 46 | corners = chooseRandomGoodQuad(lines_a, lines_b, median_contour) 47 | 48 | # warp_img, M = getTileImage(img_orig, corners.astype(np.float32),tile_buffer=16, tile_res=16) 49 | M = getTileTransform(corners.astype(np.float32),tile_buffer=16, tile_res=16) 50 | 51 | # Warp lines and draw them on warped image 52 | all_lines = np.vstack([lines_a[:,:2], lines_a[:,2:], lines_b[:,:2], lines_b[:,2:]]).astype(np.float32) 53 | warp_pts = cv2.perspectiveTransform(all_lines[None,:,:], M) 54 | warp_pts = warp_pts[0,:,:] 55 | warp_lines_a = np.hstack([warp_pts[:len(lines_a),:], warp_pts[len(lines_a):2*len(lines_a),:]]) 56 | warp_lines_b = np.hstack([warp_pts[2*len(lines_a):2*len(lines_a)+len(lines_b),:], warp_pts[2*len(lines_a)+len(lines_b):,:]]) 57 | 58 | 59 | # Get thetas of warped lines 60 | thetas_a = np.array([getSegmentTheta(line) for line in warp_lines_a]) 61 | thetas_b = np.array([getSegmentTheta(line) for line in warp_lines_b]) 62 | median_theta_a = (np.median(thetas_a*180/np.pi)) 63 | median_theta_b = (np.median(thetas_b*180/np.pi)) 64 | 65 | # Gradually relax angle threshold over N iterations 66 | if i < 20: 67 | warp_angle_threshold = 0.03 68 | elif i < 30: 69 | warp_angle_threshold = 0.1 70 | elif i < 50: 71 | warp_angle_threshold = 0.3 72 | elif i < 70: 73 | warp_angle_threshold = 0.5 74 | elif i < 80: 75 | warp_angle_threshold = 1.0 76 | else: 77 | warp_angle_threshold = 2.0 78 | if ((angleCloseDeg(abs(median_theta_a), 0, warp_angle_threshold) and 79 | angleCloseDeg(abs(median_theta_b), 90, warp_angle_threshold)) or 80 | (angleCloseDeg(abs(median_theta_a), 90, warp_angle_threshold) and 81 | angleCloseDeg(abs(median_theta_b), 0, warp_angle_threshold))): 82 | print('Found good match (%d): %.2f %.2f' % (i, abs(median_theta_a), abs(median_theta_b))) 83 | break 84 | # else: 85 | # print('iter %d: %.2f %.2f' % (i, abs(median_theta_a), abs(median_theta_b))) 86 | 87 | warp_img, M = getTileImage(img_orig, corners.astype(np.float32),tile_buffer=16, tile_res=16) 88 | 89 | # Recalculate warp now that we're using a different tile_buffer/res 90 | # warp_pts = cv2.perspectiveTransform(all_lines[None,:,:], M) 91 | # warp_pts = warp_pts[0,:,:] 92 | # warp_lines_a = np.hstack([warp_pts[:len(lines_a),:], warp_pts[len(lines_a):2*len(lines_a),:]]) 93 | # warp_lines_b = np.hstack([warp_pts[2*len(lines_a):2*len(lines_a)+len(lines_b),:], warp_pts[2*len(lines_a)+len(lines_b):,:]]) 94 | 95 | lines_x, lines_y, step_x, step_y = getWarpCheckerLines(warp_img) 96 | if len(lines_x) > 0: 97 | print('Found good chess lines (%d): %s %s' % (i2, lines_x, lines_y)) 98 | break 99 | print("Ransac corner detection took %.4f seconds." % (time() - a)) 100 | 101 | print(lines_x, lines_y) 102 | warp_img, M = getTileImage(img_orig, corners.astype(np.float32),tile_buffer=16, tile_res=16) 103 | 104 | for corner in corners: 105 | cv2.circle(img, tuple(map(int,corner)), 5, (255,150,150),-1) 106 | 107 | if len(lines_x) > 0: 108 | print('Found chessboard?') 109 | warp_corners, all_warp_corners = getRectChessCorners(lines_x, lines_y) 110 | tile_centers = all_warp_corners + np.array([step_x/2.0, step_y/2.0]) # Offset from corner to tile centers 111 | M_inv = np.matrix(np.linalg.inv(M)) 112 | real_corners, all_real_tile_centers = getOrigChessCorners(warp_corners, tile_centers, M_inv) 113 | 114 | tile_res = 64 # Each tile has N pixels per side 115 | tile_buffer = 1 116 | warp_img, better_M = getTileImage(img_orig2, real_corners, tile_buffer=tile_buffer, tile_res=tile_res) 117 | # Further refine rectified image 118 | warp_img, was_rotated, refine_M = reRectifyImages(warp_img) 119 | # combined_M = better_M 120 | combined_M = np.matmul(refine_M,better_M) 121 | M_inv = np.matrix(np.linalg.inv(combined_M)) 122 | 123 | # Get better_M based corners 124 | hlines = vlines = (np.arange(8)+tile_buffer)*tile_res 125 | hcorner = (np.array([0,8,8,0])+tile_buffer)*tile_res 126 | vcorner = (np.array([0,0,8,8])+tile_buffer)*tile_res 127 | ideal_corners = np.vstack([hcorner,vcorner]).T 128 | ideal_all_corners = np.array(list(itertools.product(hlines, vlines))) 129 | ideal_tile_centers = ideal_all_corners + np.array([tile_res/2.0, tile_res/2.0]) # Offset from corner to tile centers 130 | 131 | real_corners, all_real_tile_centers = getOrigChessCorners(ideal_corners, ideal_tile_centers, M_inv) 132 | 133 | # Get final refined rectified warped image for saving 134 | warp_img, _ = getTileImage(img_orig2, real_corners, tile_buffer=tile_buffer, tile_res=tile_res) 135 | 136 | cv2.polylines(img, [real_corners.astype(np.int32)], True, (150,50,255), thickness=3) 137 | cv2.polylines(img, [all_real_tile_centers.astype(np.int32)], False, (0,50,255), thickness=1) 138 | 139 | # Update mask with predicted chessboard 140 | cv2.drawContours(mask,[real_corners.astype(int)],0,1,-1) 141 | 142 | 143 | img_masked_full = cv2.bitwise_and(img,img,mask = (mask > 0.5).astype(np.uint8)) 144 | img_masked = cv2.addWeighted(img,0.2,img_masked_full,0.8,0) 145 | 146 | drawMinAreaRect(img_masked, min_area_rect) 147 | 148 | return img_masked, edges_masked, warp_img 149 | 150 | 151 | def other(): 152 | # vals = np.array([224, 231, 238, 257, 271, 278, 300, 321, 342, 358, 362, 383, 404, 425, 436, 463, 474]) 153 | # vals_wrong = np.array([ 257., 278., 300., 321., 342., 358., 362., 383., 404.]) 154 | # vals = np.array([206, 222, 239, 256, 268, 273, 286, 290, 307, 324, 341, 345, 357, 373]) 155 | # vals_wrong = np.array([ 226.5, 239., 256., 268., 273., 286., 290., 307., 319.5]) 156 | # vals = np.array([252, 260, 272, 278, 294, 300, 314, 336, 357, 379, 400]) 157 | # vals = np.array([272, 283, 298, 306, 324, 331, 349, 374, 399, 424, 449]) 158 | # vals = np.array([13, 29, 49, 64, 82, 88, 96, 150, 159, 167, 179, 204, 212, 218, 228, 235, 247, 260, 272, 285, 305, 338, 363, 370, 380, 389, 402, 411, 432, 463, 478]) 159 | vals = np.array([67, 93, 100, 111, 122, 140, 147, 158, 172, 184, 209, 219, 228, 237, 249, 273, 298, 317, 324, 344, 349, 356, 374, 400, 414, 426]) 160 | 161 | print(vals) 162 | print(np.diff(vals)) 163 | # sub_arr = np.abs(vals[:,None] - vals) 164 | # print(sub_arr) 165 | 166 | n_pts = 3 167 | n = scipy.special.binom(len(vals),n_pts) 168 | # devs = np.zeros(n) 169 | # plt.plot(vals_wrong,np.zeros(len(vals_wrong)),'rs') 170 | 171 | a = time() 172 | best_spacing = getBestEqualSpacing(vals) 173 | print("iter cost took %.4f seconds for %d combinations." % (time() - a, n)) 174 | print(best_spacing) 175 | plt.plot(best_spacing,0.05+np.zeros(len(best_spacing)),'gx') 176 | 177 | # plt.hist(devs, 50) 178 | 179 | plt.plot(vals,-0.1 + np.zeros(len(vals)),'k.', ms=10) 180 | plt.show() 181 | 182 | 183 | 184 | def main(filenames): 185 | for filename in filenames: 186 | a = time() 187 | img_masked, edges_masked, warp_img = processFile(filename) 188 | print("Full image file process took %.4f seconds." % (time() - a)) 189 | cv2.imshow('img %s' % filename,img_masked) 190 | cv2.imshow('warp %s' % filename, warp_img) 191 | out_filename = filename[:-4].replace('/','_').replace('\\','_') 192 | print(filename[:-4], out_filename) 193 | PIL.Image.fromarray(cv2.cvtColor(warp_img,cv2.COLOR_BGR2RGB)).save("rectified2/%s.png" % out_filename) 194 | # cv2.imshow('edges %s' % filename, edges_masked) 195 | 196 | 197 | cv2.waitKey(0) 198 | cv2.destroyAllWindows() 199 | plt.show() 200 | 201 | if __name__ == '__main__': 202 | if len(sys.argv) > 1: 203 | filenames = sys.argv[1:] 204 | else: 205 | # filenames = ['input2/02.jpg'] 206 | # filenames = ['input2/01.jpg'] 207 | filenames = ['input/30.jpg'] 208 | print("Loading", filenames) 209 | main(filenames) 210 | # other() -------------------------------------------------------------------------------- /countEntriesInTfRecords.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import glob 3 | import numpy as np 4 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 5 | 6 | 7 | filenames = glob.glob('datasets/tfrecords/winsize_10_color/*.tfrecords') 8 | 9 | c = 0 10 | for filename in filenames: 11 | k = 0 12 | for record in tf.python_io.tf_record_iterator(filename): 13 | k += 1 14 | print('%s : %d entries' % (filename, k)) 15 | c += k 16 | 17 | print('Total : %d entries' % c) -------------------------------------------------------------------------------- /data/bad_tiles.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/data/bad_tiles.png -------------------------------------------------------------------------------- /data/good_tiles.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/data/good_tiles.png -------------------------------------------------------------------------------- /example_unity_grid.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/example_unity_grid.png -------------------------------------------------------------------------------- /generateFullDataset.py: -------------------------------------------------------------------------------- 1 | # Given a list of pts text files, build a complete dataset from it. 2 | import glob 3 | import os 4 | import PIL.Image 5 | import cv2 6 | import numpy as np 7 | from time import time 8 | from argparse import ArgumentParser 9 | from scipy.spatial import cKDTree 10 | import tensorflow as tf 11 | import SaddlePoints 12 | import errno 13 | 14 | def mkdir_p(path): 15 | try: 16 | os.makedirs(path) 17 | except OSError as exc: # Python >2.5 18 | if os.path.isdir(path): 19 | pass 20 | else: 21 | raise 22 | 23 | 24 | # Given chessboard corners, get all 7x7 = 49 internal x-corner positions. 25 | def getXcorners(corners): 26 | # Get Xcorners for image 27 | ideal_corners = np.array([[0,1],[1,1],[1,0],[0,0]],dtype=np.float32) 28 | M = cv2.getPerspectiveTransform(ideal_corners, corners) # From ideal to real. 29 | 30 | # 7x7 internal grid of 49 x-corners/ 31 | xx,yy = np.meshgrid(np.arange(7, dtype=np.float32), np.arange(7, dtype=np.float32)) 32 | all_ideal_grid_pts = np.vstack([xx.flatten(), yy.flatten()]).T 33 | all_ideal_grid_pts = (all_ideal_grid_pts + 1) / 8.0 34 | 35 | chess_xcorners = cv2.perspectiveTransform(np.expand_dims(all_ideal_grid_pts,0), M)[0,:,:] 36 | return chess_xcorners 37 | 38 | 39 | def getPointsNearPoints(ptsA, ptsB, MIN_DIST_PX=3): 40 | # Returns a mask for points in A that are close by MIN_DIST_PX to points in B 41 | min_dists, min_dist_idx = cKDTree(ptsB).query(ptsA, 1) 42 | mask = min_dists < MIN_DIST_PX 43 | return mask 44 | 45 | # Load image from path 46 | def loadImage(img_filepath): 47 | print ("Processing %s" % (img_filepath)) 48 | 49 | img = PIL.Image.open(img_filepath) 50 | if (img.size[0] > 640): 51 | img = img.resize((640, 480), PIL.Image.BICUBIC) 52 | gray = np.array(img.convert('L')) 53 | rgb = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) 54 | return rgb, gray 55 | 56 | def getTiles(pts, img_gray, WINSIZE=10): 57 | # NOTE : Assumes no point is within WINSIZE of an edge! 58 | # Points Nx2, columns should be x and y, not r and c. 59 | # WINSIZE = the number of pixels out from the point that a tile should be. 60 | 61 | # Build tiles of size Nx(2*WINSIZE+1)x(2*WINSIZE+1) 62 | img_shape = np.array([img_gray.shape[1], img_gray.shape[0]]) 63 | tiles = np.zeros([len(pts), WINSIZE*2+1, WINSIZE*2+1], dtype=img_gray.dtype) 64 | for i, pt in enumerate(np.round(pts).astype(np.int64)): 65 | tiles[i,:,:] = img_gray[pt[1]-WINSIZE:pt[1]+WINSIZE+1, 66 | pt[0]-WINSIZE:pt[0]+WINSIZE+1] 67 | return tiles 68 | 69 | def getTilesColor(pts, img, WINSIZE=10): 70 | # NOTE : Assumes no point is within WINSIZE of an edge! 71 | # Points Nx2, columns should be x and y, not r and c. 72 | # WINSIZE = the number of pixels out from the point that a tile should be. 73 | 74 | # Build tiles of size Nx(2*WINSIZE+1)x(2*WINSIZE+1) 75 | img_shape = np.array([img.shape[1], img.shape[0]]) 76 | tiles = np.zeros([len(pts), WINSIZE*2+1, WINSIZE*2+1, 3], dtype=img.dtype) 77 | for i, pt in enumerate(np.round(pts).astype(np.int64)): 78 | tiles[i,:,:,:] = img[pt[1]-WINSIZE:pt[1]+WINSIZE+1, 79 | pt[0]-WINSIZE:pt[0]+WINSIZE+1, :] 80 | return tiles 81 | 82 | # View image with chessboard lines overlaid. 83 | def addOverlay(idx, img, corners, good_xcorners, bad_pts): 84 | for pt in np.round(bad_pts).astype(np.int64): 85 | cv2.rectangle(img, tuple(pt-2),tuple(pt+2), (0,0,255), -1) 86 | 87 | for pt in np.round(good_xcorners).astype(np.int64): 88 | cv2.rectangle(img, tuple(pt-2),tuple(pt+2), (0,255,0), -1) 89 | 90 | 91 | cv2.polylines(img, 92 | [np.round(corners).astype(np.int32)], 93 | isClosed=True, thickness=2, color=(255,0,255)) 94 | 95 | cv2.putText(img, 96 | 'Frame % 4d' % (idx), 97 | (5,15), cv2.FONT_HERSHEY_PLAIN, 1.0,(255,255,255),0) 98 | 99 | def visualizeTiles(tiles): 100 | # Assumes no more than 49 tiles, only plots the first 49 101 | N = len(tiles) 102 | # assert N <= 49 103 | assert tiles.shape[1] == tiles.shape[2] # square tiles 104 | side = tiles.shape[1] 105 | cols = 7#int(np.ceil(np.sqrt(N))) 106 | rows = 7#int(np.ceil(N/(cols)))+1 107 | tile_img = np.zeros([rows*side, cols*side, 3], dtype=tiles.dtype) 108 | for i in range(min(N,49)): 109 | r, c = side*(int(i/cols)), side*(i%cols) 110 | tile_img[r:r+side, c:c+side,:] = tiles[i,:,:,:] 111 | return tile_img 112 | 113 | # Converting the values into features 114 | # _int64 is used for numeric values 115 | def _int64_feature(value): 116 | return tf.train.Feature(int64_list=tf.train.Int64List(value=[value])) 117 | 118 | # _bytes is used for string/char values 119 | def _bytes_feature(value): 120 | return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) 121 | 122 | 123 | def main(args): 124 | for pointfile in args.pointfiles: 125 | with open(pointfile, 'r') as f: 126 | lines = f.readlines() 127 | video_filepath = lines[0] 128 | images_path = os.path.dirname(pointfile) 129 | 130 | # Writing to TFrecord 131 | video_filename = os.path.basename(video_filepath)[:-5] 132 | folder_path = "%s/winsize_%s_color" % (args.tfrecords_path, args.winsize) 133 | mkdir_p(folder_path) 134 | 135 | tfrecord_path = "%s/%s_ws%d.tfrecords" % (folder_path, video_filename, args.winsize) 136 | with tf.python_io.TFRecordWriter(tfrecord_path) as writer: 137 | for line in lines[1:]: 138 | tA = time() 139 | parts = line.split(',') 140 | idx = int(parts[0]) 141 | 142 | # if (idx < 260): 143 | # continue 144 | 145 | corners = np.array(parts[1:], dtype=np.float32).reshape([4,2]) 146 | xcorners = getXcorners(corners) 147 | 148 | filename = "%s/frame_%03d.jpg" % (images_path, idx) 149 | img, gray = loadImage(filename) 150 | 151 | # Saddle points 152 | spts, gx, gy = SaddlePoints.getFinalSaddlePoints(gray, WINSIZE=args.winsize) 153 | 154 | good_spt_mask = getPointsNearPoints(spts, xcorners) 155 | good_xcorners = spts[good_spt_mask] 156 | bad_spts = spts[~good_spt_mask] 157 | 158 | # Only keep the same # of bad points as good 159 | # Shuffle bad points so we get a good smattering. 160 | N = len(good_xcorners) 161 | np.random.shuffle(bad_spts) 162 | bad_spts = bad_spts[:N] 163 | 164 | # good_xcorners, bad_xcorners, bad_spts, spts, keep_mask = getXcornersNearSaddlePts(gray, xcorners) 165 | 166 | tiles = getTilesColor(good_xcorners, img, WINSIZE=args.winsize) 167 | bad_tiles = getTilesColor(bad_spts, img, WINSIZE=args.winsize) 168 | 169 | # Write tiles to tf-records 170 | for tile in tiles: 171 | feature = { 'label': _int64_feature(1), 172 | 'image': _bytes_feature(tf.compat.as_bytes(tile.tostring())) } 173 | example = tf.train.Example(features=tf.train.Features(feature=feature)) 174 | writer.write(example.SerializeToString()) 175 | 176 | for tile in bad_tiles: 177 | feature = { 'label': _int64_feature(0), 178 | 'image': _bytes_feature(tf.compat.as_bytes(tile.tostring())) } 179 | example = tf.train.Example(features=tf.train.Features(feature=feature)) 180 | writer.write(example.SerializeToString()) 181 | 182 | 183 | if args.viztiles: 184 | tile_img = visualizeTiles(tiles) 185 | bad_tile_img = visualizeTiles(bad_tiles) 186 | 187 | print('\t Took %.1f ms.' % ((time() - tA)*1000)) 188 | 189 | if args.vizoverlay: 190 | overlay_img = img.copy() 191 | addOverlay(idx, overlay_img, corners, good_xcorners, bad_spts) 192 | 193 | cv2.imshow('frame',overlay_img) 194 | 195 | if args.viztiles: 196 | cv2.imshow('tiles', tile_img) 197 | cv2.imshow('bad_tiles', bad_tile_img) 198 | 199 | if (args.vizoverlay or args.viztiles): 200 | if (cv2.waitKey(1) & 0xFF == ord('q')): 201 | break 202 | 203 | if __name__ == '__main__': 204 | parser = ArgumentParser() 205 | parser.add_argument("pointfiles", nargs='+', 206 | help="All pts.txt points files containing filename and chessboard coordinates.") 207 | parser.add_argument("-savetf", 208 | action='store_true', help="Whether to save tfrecords") 209 | parser.add_argument("-viztiles", 210 | action='store_true', help="Whether to visualize tiles or not") 211 | parser.add_argument("-vizoverlay", 212 | action='store_true', help="Whether to visualize overlay") 213 | parser.add_argument("--tfrecords_path", default='datasets/tfrecords', 214 | help="Folder to store tfrecord output") 215 | parser.add_argument("-ws", "--winsize", dest="winsize", default=10, type=int, 216 | help="Half window size (full kernel = 2*winsize + 1)") 217 | args = parser.parse_args() 218 | print(args) 219 | main(args) -------------------------------------------------------------------------------- /generateMLDataset.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import PIL.Image 3 | import matplotlib.image as mpimg 4 | import scipy.ndimage 5 | import cv2 # For Sobel etc 6 | import glob 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | import os 10 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 11 | 12 | # Load pt_dataset.txt and generate the windowed tiles for all the good and bad 13 | # points in folders dataset/good dataset/bad 14 | 15 | 16 | def loadImage(filepath, doGrayscale=False): 17 | img_orig = PIL.Image.open(filepath) 18 | img_width, img_height = img_orig.size 19 | 20 | # Resize 21 | aspect_ratio = min(500.0/img_width, 500.0/img_height) 22 | new_width, new_height = ((np.array(img_orig.size) * aspect_ratio)).astype(int) 23 | img = img_orig.resize((new_width,new_height), resample=PIL.Image.BILINEAR) 24 | if (doGrayscale): 25 | img = img.convert('L') # grayscale 26 | img = np.array(img) 27 | 28 | return img 29 | 30 | import errno 31 | def mkdir_p(path): 32 | try: 33 | os.makedirs(path) 34 | except OSError as exc: # Python >2.5 35 | if os.path.isdir(path): 36 | pass 37 | else: 38 | raise 39 | 40 | def main(): 41 | input_data = 'pt_dataset2.txt' 42 | 43 | WINSIZE = 5 44 | dataset_folder = 'dataset_gray_%d' % WINSIZE 45 | 46 | DO_GRAYSCALE = True 47 | DO_BINARIZATION = False 48 | DO_OPENING = False 49 | 50 | if (DO_BINARIZATION and not DO_GRAYSCALE): 51 | raise('Error, must be grayscale if doing binarization.') 52 | 53 | count_good = 0 54 | count_bad = 0 55 | 56 | good_features = [] 57 | good_labels = [] 58 | bad_features = [] 59 | bad_labels = [] 60 | 61 | # save all points to a file 62 | with open(input_data, 'r') as f: 63 | lines = [x.strip() for x in f.readlines()] 64 | n = len(lines)/5 65 | # n = 1 66 | for i in range(n): 67 | print("On %d/%d" % (i+1, n)) 68 | filename = lines[i*5] 69 | s0 = lines[i*5+1].split() 70 | s1 = lines[i*5+2].split() 71 | s2 = lines[i*5+3].split() 72 | s3 = lines[i*5+4].split() 73 | good_pts = np.array([s1, s0], dtype=np.int).T 74 | bad_pts = np.array([s3, s2], dtype=np.int).T 75 | 76 | img_filepath = 'input/%s.png' % filename 77 | if not os.path.exists(img_filepath): 78 | img_filepath = 'input/%s.jpg' % filename 79 | if not os.path.exists(img_filepath): 80 | img_filepath = 'input_yt/%s.jpg' % filename 81 | if not os.path.exists(img_filepath): 82 | img_filepath = 'input_yt/%s.png' % filename 83 | img = loadImage(img_filepath, DO_GRAYSCALE) 84 | 85 | kernel = np.ones((3,3),np.uint8) 86 | 87 | # Good points 88 | for i in range(good_pts.shape[0]): 89 | pt = good_pts[i,:] 90 | if (np.any(pt <= WINSIZE) or np.any(pt >= np.array(img.shape[:2]) - WINSIZE)): 91 | continue 92 | else: 93 | tile = img[pt[0]-WINSIZE:pt[0]+WINSIZE+1, pt[1]-WINSIZE:pt[1]+WINSIZE+1] 94 | 95 | if DO_BINARIZATION: 96 | tile = cv2.adaptiveThreshold(tile,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2) 97 | 98 | if DO_OPENING: 99 | tile = cv2.morphologyEx(tile, cv2.MORPH_OPEN, kernel) 100 | 101 | 102 | good_features.append(tile) 103 | good_labels.append(1) 104 | 105 | count_good += 1 106 | 107 | # Bad points 108 | for i in range(bad_pts.shape[0]): 109 | pt = bad_pts[i,:] 110 | if (np.any(pt <= WINSIZE) or np.any(pt >= np.array(img.shape[:2]) - WINSIZE)): 111 | continue 112 | else: 113 | tile = img[pt[0]-WINSIZE:pt[0]+WINSIZE+1, pt[1]-WINSIZE:pt[1]+WINSIZE+1] 114 | if DO_BINARIZATION: 115 | tile = cv2.adaptiveThreshold(tile,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2) 116 | if DO_OPENING: 117 | tile = cv2.morphologyEx(tile, cv2.MORPH_OPEN, kernel) 118 | 119 | bad_features.append(tile) 120 | bad_labels.append(0) 121 | 122 | count_bad += 1 123 | 124 | features = np.array(good_features + bad_features) 125 | print(features.shape) 126 | labels = np.array(good_labels + bad_labels, dtype=np.float32) 127 | print(labels.shape) 128 | 129 | np.savez('dataset2_%d' % WINSIZE, features=features, labels=labels) 130 | # Example to use: print(np.load('dataset_5.npz')['features']) 131 | 132 | print ("Finished %d good and %d bad tiles" % (count_good, count_bad)) 133 | 134 | 135 | 136 | if __name__ == '__main__': 137 | main() 138 | 139 | 140 | 141 | -------------------------------------------------------------------------------- /generateMLTiles.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import PIL.Image 3 | import matplotlib.image as mpimg 4 | import scipy.ndimage 5 | import cv2 # For Sobel etc 6 | import glob 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | import tensorflow as tf 10 | import os 11 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 12 | 13 | # Load pt_dataset.txt and generate the windowed tiles for all the good and bad 14 | # points in folders dataset/good dataset/bad 15 | 16 | 17 | def loadImage(filepath, doGrayscale=False): 18 | img_orig = PIL.Image.open(filepath) 19 | img_width, img_height = img_orig.size 20 | 21 | # Resize 22 | aspect_ratio = min(500.0/img_width, 500.0/img_height) 23 | new_width, new_height = ((np.array(img_orig.size) * aspect_ratio)).astype(int) 24 | img = img_orig.resize((new_width,new_height), resample=PIL.Image.BILINEAR) 25 | if (doGrayscale): 26 | img = img.convert('L') # grayscale 27 | else: 28 | # color, because sometimes it could be RGBA 29 | img = img.convert('RGB') 30 | img = np.array(img) 31 | 32 | return img 33 | 34 | import errno 35 | def mkdir_p(path): 36 | try: 37 | os.makedirs(path) 38 | except OSError as exc: # Python >2.5 39 | if os.path.isdir(path): 40 | pass 41 | else: 42 | raise 43 | 44 | # Converting the values into features 45 | # _int64 is used for numeric values 46 | def _int64_feature(value): 47 | return tf.train.Feature(int64_list=tf.train.Int64List(value=[value])) 48 | 49 | # _bytes is used for string/char values 50 | def _bytes_feature(value): 51 | return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) 52 | 53 | def main(): 54 | input_data = 'pt_dataset2.txt' 55 | 56 | WINSIZE = 10 57 | DO_BINARIZATION = False 58 | DO_OPENING = False 59 | DO_GRAYSCALE = False 60 | DO_TFRECORD = True 61 | 62 | 63 | if DO_GRAYSCALE: 64 | dataset_folder = 'dataset_gray_%d' % WINSIZE 65 | else: 66 | dataset_folder = 'dataset_rgb_%d' % WINSIZE 67 | folder_good = '%s/good' % dataset_folder 68 | folder_bad = '%s/bad' % dataset_folder 69 | mkdir_p(folder_good) 70 | mkdir_p(folder_bad) 71 | 72 | 73 | if (DO_BINARIZATION and not DO_GRAYSCALE): 74 | raise('Error, must be grayscale if doing binarization.') 75 | 76 | count_good = 0 77 | count_bad = 0 78 | 79 | # save all points to a file 80 | with open(input_data, 'r') as f: 81 | lines = [x.strip() for x in f.readlines()] 82 | n = len(lines)/5 83 | # n = 1 84 | if DO_TFRECORD: 85 | tfrecord_path = "%s/%s_ws%d.tfrecords" % ('datasets/raw','input_images', WINSIZE) 86 | else: 87 | tfrecord_path = '/tmp/tmpdelete.tfrecords' 88 | 89 | with tf.python_io.TFRecordWriter(tfrecord_path) as writer: 90 | for i in range(n): 91 | print("On %d/%d" % (i+1, n)) 92 | filename = lines[i*5] 93 | s0 = lines[i*5+1].split() 94 | s1 = lines[i*5+2].split() 95 | s2 = lines[i*5+3].split() 96 | s3 = lines[i*5+4].split() 97 | good_pts = np.array([s1, s0], dtype=np.int).T 98 | bad_pts = np.array([s3, s2], dtype=np.int).T 99 | 100 | img_filepath = 'datasets/raw/input/%s.png' % filename 101 | if not os.path.exists(img_filepath): 102 | img_filepath = 'datasets/raw/input/%s.jpg' % filename 103 | if not os.path.exists(img_filepath): 104 | img_filepath = 'datasets/raw/input_yt/%s.jpg' % filename 105 | if not os.path.exists(img_filepath): 106 | img_filepath = 'datasets/raw/input_yt/%s.png' % filename 107 | img = loadImage(img_filepath, DO_GRAYSCALE) 108 | 109 | kernel = np.ones((3,3),np.uint8) 110 | 111 | # Good points 112 | for i in range(good_pts.shape[0]): 113 | pt = good_pts[i,:] 114 | if (np.any(pt <= WINSIZE) or np.any(pt >= np.array(img.shape[:2]) - WINSIZE)): 115 | # print("Skipping point %s" % pt) 116 | continue 117 | else: 118 | tile = img[pt[0]-WINSIZE:pt[0]+WINSIZE+1, pt[1]-WINSIZE:pt[1]+WINSIZE+1] 119 | # print(tile) 120 | out_filename = '%s/%s_%03d.png' % (folder_good, filename, i) 121 | if DO_BINARIZATION: 122 | tile = cv2.adaptiveThreshold(tile,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2) 123 | 124 | if DO_OPENING: 125 | tile = cv2.morphologyEx(tile, cv2.MORPH_OPEN, kernel) 126 | 127 | if DO_GRAYSCALE: 128 | im = PIL.Image.fromarray(tile).convert('L') 129 | else: 130 | im = PIL.Image.fromarray(tile).convert('RGB') 131 | 132 | if DO_TFRECORD: 133 | feature = { 'label': _int64_feature(1), 134 | 'image': _bytes_feature(tf.compat.as_bytes(tile.tostring())) } 135 | example = tf.train.Example(features=tf.train.Features(feature=feature)) 136 | writer.write(example.SerializeToString()) 137 | else: 138 | im.save(out_filename) 139 | count_good += 1 140 | 141 | # Bad points 142 | for i in range(bad_pts.shape[0]): 143 | pt = bad_pts[i,:] 144 | if (np.any(pt <= WINSIZE) or np.any(pt >= np.array(img.shape[:2]) - WINSIZE)): 145 | # print("Skipping point %s" % pt) 146 | continue 147 | else: 148 | tile = img[pt[0]-WINSIZE:pt[0]+WINSIZE+1, pt[1]-WINSIZE:pt[1]+WINSIZE+1] 149 | out_filename = '%s/%s_%03d.png' % (folder_bad, filename, i) 150 | if DO_BINARIZATION: 151 | tile = cv2.adaptiveThreshold(tile,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2) 152 | if DO_OPENING: 153 | tile = cv2.morphologyEx(tile, cv2.MORPH_OPEN, kernel) 154 | 155 | if DO_GRAYSCALE: 156 | im = PIL.Image.fromarray(tile).convert('L') 157 | else: 158 | im = PIL.Image.fromarray(tile).convert('RGB') 159 | 160 | if DO_TFRECORD: 161 | feature = { 'label': _int64_feature(0), 162 | 'image': _bytes_feature(tf.compat.as_bytes(tile.tostring())) } 163 | example = tf.train.Example(features=tf.train.Features(feature=feature)) 164 | writer.write(example.SerializeToString()) 165 | else: 166 | im.save(out_filename) 167 | count_bad += 1 168 | print ("Finished %d good and %d bad tiles" % (count_good, count_bad)) 169 | 170 | 171 | 172 | if __name__ == '__main__': 173 | main() 174 | 175 | 176 | 177 | -------------------------------------------------------------------------------- /helpers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def scaleImageIfNeeded(img, max_width=1000, max_height=800): 4 | """Scale image down to max_width / max_height keeping aspect ratio if needed. Do nothing otherwise.""" 5 | # Input and Output is a PIL Image 6 | img_width, img_height = img.size 7 | # print("Image size %dx%d" % (img_width, img_height)) 8 | aspect_ratio = min(max_width/img_width, max_height/img_height) 9 | if aspect_ratio < 1.0: 10 | new_width, new_height = ((np.array(img.size) * aspect_ratio)).astype(int) 11 | # print(" Resizing to %dx%d" % (new_width, new_height)) 12 | return img.resize((new_width,new_height)) 13 | return img 14 | 15 | def getSegmentAngle(line): 16 | x1,y1,x2,y2 = line 17 | return np.math.atan2(y2-y1, x2-x1) 18 | 19 | def getLineGradients(line, gradient_x, gradient_y, gradient_mag, sampling_rate=0.5, 20 | normal_magnitude_threshold=0.9, normal_strong_gradient_ratio_threshold=0.45, pos_neg_ratio_threshold=0.2): 21 | """Calculate normal gradient values along line given x/y gradients and a line segment.""" 22 | 23 | # 1 - Get gradient values 24 | line = np.array(line) 25 | ptA = line[:2] 26 | ptB = line[2:] 27 | 28 | # unit vector in direction of line 29 | line_length = np.linalg.norm(ptB - ptA) 30 | line_direction = (ptB - ptA) / line_length 31 | 32 | # Convert to normal 33 | line_normal = np.array([-line_direction[1], line_direction[0]]) # -y, x for normal in one direction 34 | 35 | # Get points along line, choosing number of points giving a sampling rate in pixels per points (1-1 is good) 36 | num_pts_on_line = np.ceil(np.sqrt(np.sum((ptB - ptA)**2)) / sampling_rate) 37 | guessx = np.linspace(ptA[1],ptB[1],num_pts_on_line) 38 | guessy = np.linspace(ptA[0],ptB[0],num_pts_on_line) 39 | 40 | line_indices = np.floor(np.vstack((guessx, guessy)).T).astype(int) 41 | gradients = np.vstack( 42 | [gradient_x[line_indices[:,0], line_indices[:,1]], 43 | gradient_y[line_indices[:,0], line_indices[:,1]]]) 44 | gradient_mags = gradient_mag[line_indices[:,0], line_indices[:,1]] 45 | 46 | # Calculate average strength of gradient along line as a score 47 | gradient_on_normal = line_normal.dot(gradients) 48 | avg_normal_gradient = np.abs(gradient_on_normal).mean() 49 | 50 | # Magnitude of gradient along normal, normalized by total gradient magnitude at that point 51 | # ex. 1.0 means strong + gradient perfectly normal to line 52 | with np.errstate(divide='ignore', invalid='ignore'): 53 | # some gradient magnitudes are zero, set normal gradients to zero for those 54 | normal_gradients = gradient_on_normal / gradient_mags 55 | # Ignore weak gradients 56 | normal_gradients[(gradient_mags < 10)] = 0 57 | 58 | # Ratio of how often the absolute of the gradient is greater than the threshold 59 | # chessboard lines should have extremely high ~90% ratios 60 | normal_strong_gradient_ratio = np.sum(abs(normal_gradients) > normal_magnitude_threshold) / float(num_pts_on_line) 61 | 62 | # Ratio of mag near top vs mag near bot 63 | pos_edge_ratio = np.sum(normal_gradients > normal_magnitude_threshold) / float(num_pts_on_line) 64 | neg_edge_ratio = np.sum(-normal_gradients > normal_magnitude_threshold) / float(num_pts_on_line) 65 | 66 | # Calculate aspect ratio of positive to negative values, but only if they're both reasonably strong 67 | if (pos_edge_ratio < pos_neg_ratio_threshold or neg_edge_ratio < pos_neg_ratio_threshold): 68 | edge_ratio = 0 69 | else: 70 | if pos_edge_ratio > neg_edge_ratio: 71 | edge_ratio = neg_edge_ratio / pos_edge_ratio 72 | else: 73 | edge_ratio = pos_edge_ratio / neg_edge_ratio 74 | 75 | # if normal_strong_gradient_ratio > normal_strong_gradient_ratio_threshold and edge_ratio != 0: 76 | # print("%.2f / %.2f = %.2f %s | %.2f" % (pos_edge_ratio, neg_edge_ratio, edge_ratio, edge_ratio < 0.75, normal_strong_gradient_ratio)) 77 | 78 | 79 | # Calculate fft, since sampling rate is static, we can just use indices as a comparison method 80 | fft_result = np.abs(np.fft.rfft(normal_gradients).real) 81 | strongest_freq = np.argmax(fft_result) 82 | 83 | is_good = True 84 | 85 | # Sanity check normal gradients span positive and negative range 86 | if normal_gradients.min() > -normal_magnitude_threshold or normal_gradients.max() < normal_magnitude_threshold: 87 | is_good = False 88 | # Sanity check most of the normal gradients are maximized 89 | elif normal_strong_gradient_ratio < normal_strong_gradient_ratio_threshold: 90 | is_good = False 91 | # Check that ratio of positive to negative is somewhere near 50/50, 1.0 means perfect 50/50, 0.5 gives a lot of leeway 92 | elif edge_ratio < 0.5: 93 | is_good = False 94 | # Check that there is a low frequency signal in normal gradient 95 | elif strongest_freq < 2 or strongest_freq > 20: 96 | is_good = False 97 | 98 | # Recover potentially good lines 99 | if edge_ratio > 0.9 and normal_strong_gradient_ratio > normal_strong_gradient_ratio_threshold: 100 | is_good = True 101 | 102 | # if is_good: 103 | # print("%.2f : %.2f / %.2f = %.2f | %.2f | %d" % ( 104 | # avg_normal_gradient, pos_edge_ratio, neg_edge_ratio, edge_ratio, 105 | # normal_strong_gradient_ratio, strongest_freq)) 106 | 107 | return is_good, strongest_freq, normal_gradients, fft_result, edge_ratio, avg_normal_gradient 108 | 109 | def angleClose(a, b, angle_threshold=10*np.pi/180): 110 | d = np.abs(a - b) 111 | # Handle angles that are ~360 or ~180 degrees apart 112 | return d < angle_threshold or np.abs(np.pi-d) < angle_threshold or np.abs(2*np.pi-d) < angle_threshold 113 | 114 | def segmentAngles(angles, good_mask, angle_threshold=15*np.pi/180): 115 | # Partition lines based on similar angles int segments/groups 116 | good = np.zeros(len(angles),dtype=bool) 117 | segment_mask = np.zeros(angles.shape, dtype=int) 118 | 119 | segment_idx = 1 120 | for i in range(angles.size): 121 | # Skip if not a good line or line already grouped 122 | if not good_mask[i] or segment_mask[i] != 0: 123 | continue 124 | 125 | # Create new group 126 | segment_mask[i] = segment_idx 127 | for j in range(i+1, angles.size): 128 | # If good line, not yet grouped, and is close in angle, add to segment group 129 | if good_mask[j] and segment_mask[j] == 0 and angleClose(angles[i], angles[j], angle_threshold): 130 | segment_mask[j] = segment_idx 131 | # Iterate to next group 132 | segment_idx += 1 133 | 134 | return segment_mask # segments 135 | 136 | def chooseBestSegments(segments, line_mags): 137 | num_segments = segments.max() # 1-indexed, 0 is a masked/bad segment 138 | segment_mags = np.zeros(num_segments+1) 139 | for i in range(1, num_segments+1): 140 | num_in_segment = np.sum(segments == i) 141 | if num_in_segment < 4: 142 | # Need at least 4 lines in a segment 143 | segment_mags[i] = 0 144 | else: 145 | # Get average line gradient magnitude for that segment 146 | segment_mags[i] = np.sum(line_mags[segments == i])/num_in_segment 147 | 148 | # print("num:",num_segments) 149 | # print("mags:",segment_mags) 150 | order = np.argsort(segment_mags)[::-1] 151 | return order[:2] # Top two segments only -------------------------------------------------------------------------------- /hog.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from pylab import * 3 | import numpy as np 4 | import cv2 5 | import sys 6 | import IPython 7 | import matplotlib.pyplot as plt 8 | from skimage.feature import hog 9 | from skimage import color, exposure 10 | 11 | np.set_printoptions(suppress=True, precision=2) # Better printing of arrays 12 | # def getHog 13 | 14 | 15 | def main(filenames): 16 | for filename in filenames: 17 | print("Processing %s" % filename) 18 | img = cv2.imread(filename) 19 | if (img.size > 1000*1000): 20 | img = cv2.resize(img,None,fx=0.3, fy=0.3, interpolation = cv2.INTER_AREA) 21 | gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY) 22 | 23 | # Edges 24 | # edges = cv2.Canny(gray,200,500,apertureSize = 3, L2gradient=False) # Better thresholds 25 | 26 | # # Gradients 27 | # sobelx = cv2.Sobel(img,cv2.CV_64F,1,0,ksize=5) 28 | # sobely = cv2.Sobel(img,cv2.CV_64F,0,1,ksize=5) 29 | # grad_mag = np.sqrt(sobelx**2+sobely**2) 30 | # grad_phase = np.arctan2(sobely, sobelx) 31 | # grad_mag_norm = grad_mag + grad_mag[:].min() 32 | # grad_mag_norm = (grad_mag_norm / grad_mag_norm[:].max()) 33 | 34 | # grad_phase_norm = (grad_phase + np.pi)/(2*np.pi) * grad_mag_norm 35 | # w,h,_ = grad_phase_norm.shape 36 | # grad_phase_norm[grad_mag_norm < 0.05] = 0 37 | 38 | fd, hog_image = hog(gray, orientations=32, pixels_per_cell=(16, 16), 39 | cells_per_block=(1, 1), visualise=True) 40 | 41 | fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4), sharex=True, sharey=True) 42 | 43 | ax1.axis('off') 44 | ax1.imshow(gray, cmap=plt.cm.gray) 45 | ax1.grid('on') 46 | ax1.set_title('Input image') 47 | ax1.set_adjustable('box-forced') 48 | 49 | # Rescale histogram for better display 50 | hog_image_rescaled = exposure.rescale_intensity(hog_image, in_range=(0, 10.0)) 51 | 52 | ax2.axis('off') 53 | ax2.imshow(hog_image_rescaled, cmap=plt.cm.gray) 54 | ax2.grid('on') 55 | ax2.set_title('Histogram of Oriented Gradients') 56 | ax1.set_adjustable('box-forced') 57 | 58 | # pt = (139,304) 59 | # ax1.plot(pt[0], pt[1],'o') 60 | # ax2.plot(pt[0], pt[1],'o') 61 | 62 | plt.show() 63 | 64 | 65 | 66 | 67 | 68 | 69 | # Using opencv 70 | # cv2.imshow('image %dx%d' % (img.shape[1],img.shape[0]),img) 71 | # cv2.imshow('grad', grad_mag_norm) 72 | # cv2.imshow('phase', grad_phase_norm) 73 | # cv2.imshow('hist', h) 74 | # cv2.waitKey(0) 75 | # cv2.destroyAllWindows() 76 | 77 | if __name__ == '__main__': 78 | if len(sys.argv) > 1: 79 | filenames = sys.argv[1:] 80 | else: 81 | filenames = ['input/21.jpg'] 82 | main(filenames) 83 | -------------------------------------------------------------------------------- /hough_visualize.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from pylab import * 3 | import numpy as np 4 | import cv2 5 | import sys 6 | from board_detect import * 7 | from contour_detect import * 8 | from rectify_refine import * 9 | 10 | 11 | def getRhoTheta(line): 12 | x1,y1,x2,y2 = line 13 | theta = np.arctan2(y2-y1,x2-x1) 14 | rho = x1*np.cos(theta) + y1*sin(theta) 15 | return rho, theta 16 | 17 | 18 | def findAndDrawTile(img): 19 | contours, chosen_tile_idx, edges = findPotentialTiles(img) 20 | if not len(contours): 21 | return 22 | drawPotentialTiles(img, contours, chosen_tile_idx) 23 | 24 | tile_corners = getChosenTile(contours, chosen_tile_idx) 25 | 26 | hough_corners, corner_hough_lines, edges_roi = refineTile(img, edges, contours, chosen_tile_idx) 27 | drawBestHoughLines(img, hough_corners, corner_hough_lines) 28 | 29 | # Single tile warp 30 | ideal_tile = np.array([ 31 | [1,0], 32 | [1,1], 33 | [0,1], 34 | [0,0], 35 | ],dtype=np.float32) 36 | tile_res=32 37 | M = cv2.getPerspectiveTransform(hough_corners, 38 | (tile_res)*(ideal_tile+8+1)) 39 | side_len = tile_res*(8 + 1)*2 40 | out_img = cv2.warpPerspective(img, M, 41 | (side_len, side_len)) 42 | 43 | cv2.imshow('image %dx%d' % (img.shape[1],img.shape[0]),img) 44 | cv2.imshow('warp',out_img) 45 | 46 | 47 | def findAndDrawHough(img): 48 | img_diag_size = int(np.ceil(np.sqrt(img.shape[0]*img.shape[0] + img.shape[1]*img.shape[1]))) 49 | 50 | hough_img = np.zeros([2*img_diag_size/4, 180]) # -90 to 90 deg, -rho_max to rho_max 51 | 52 | gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY) 53 | edges = cv2.Canny(gray,100,650,apertureSize = 3) 54 | min_img_side = min(img.shape[:2]) 55 | minLineLength = min_img_side/4 56 | maxLineGap = min_img_side/10 57 | threshold = int(min_img_side/4) 58 | # print(minLineLength, maxLineGap, threshold) 59 | lines = cv2.HoughLinesP(edges,rho=1,theta=np.pi/180, 60 | threshold=threshold, minLineLength=minLineLength, maxLineGap=maxLineGap) 61 | 62 | if any(lines): 63 | rhothetas = np.zeros([lines.shape[0], 2]) 64 | for i, (x1,y1,x2,y2) in enumerate(lines[:,0,:]): 65 | cv2.line(img,(x1,y1),(x2,y2), (0,0,255),2) 66 | 67 | rho, theta = getRhoTheta((x1,y1,x2,y2)) 68 | rhothetas[i,:] = rho, theta 69 | img_rho, img_theta = (int(theta*180/np.pi + 90), int((rho+img_diag_size)/4)) 70 | cv2.circle(hough_img, (img_rho, img_theta), 3, (255,0,0),-1) 71 | 72 | 73 | 74 | 75 | # Using opencv 76 | cv2.imshow('image %dx%d' % (img.shape[1],img.shape[0]),img) 77 | cv2.imshow('edges',edges) 78 | cv2.imshow('hough',hough_img) 79 | 80 | def findAndDrawMask(img): 81 | # img = scaleImageIfNeeded(img, 600, 480) 82 | 83 | # Edges 84 | edges = cv2.Canny(img, 100, 550) 85 | mask, top_two_angles, min_area_rect, median_contour = getEstimatedChessboardMask(img, edges, iters=5) 86 | 87 | img_masked_full = cv2.bitwise_and(img,img,mask = (mask > 0.5).astype(np.uint8)) 88 | img_masked = cv2.addWeighted(img,0.2,img_masked_full,0.8,0) 89 | 90 | # Hough lines overlay 91 | edges_masked = cv2.bitwise_and(edges,edges,mask = (mask > 0.5).astype(np.uint8)) 92 | 93 | if top_two_angles is not None and len(top_two_angles) == 2: 94 | lines = getHoughLines(edges_masked, min_line_size=0.25*min(min_area_rect[1])) 95 | lines_a, lines_b = parseHoughLines(lines, top_two_angles, angle_threshold_deg=15) 96 | 97 | plotHoughLines(img_masked, lines, color=(255,255,255), line_thickness=1) 98 | plotHoughLines(img_masked, lines_a, color=(0,0,255)) 99 | plotHoughLines(img_masked, lines_b, color=(0,255,0)) 100 | 101 | if min_area_rect is not None: 102 | drawMinAreaRect(img_masked, min_area_rect) 103 | 104 | # cv2.imshow('Masked',img_masked) 105 | return img_masked 106 | # cv2.imshow('edges %s' % filename, edges_masked) 107 | # cv2.imshow('mask %s' % filename, mask) 108 | 109 | def findAndDrawChessboard(img): 110 | img_orig = img.copy() 111 | img_orig2 = img.copy() 112 | 113 | # Edges 114 | edges = cv2.Canny(img, 100, 550) 115 | 116 | # Get mask for where we think chessboard is 117 | mask, top_two_angles, min_area_rect, median_contour = getEstimatedChessboardMask(img, edges,iters=3) # More iters gives a finer mask 118 | if top_two_angles is None or len(top_two_angles) != 2 or min_area_rect is None: 119 | print('fail', top_two_angles) 120 | return img 121 | 122 | if mask.min() != 0: 123 | return img 124 | 125 | # Get hough lines of masked edges 126 | edges_masked = cv2.bitwise_and(edges,edges,mask = (mask > 0.5).astype(np.uint8)) 127 | img_orig = cv2.bitwise_and(img_orig,img_orig,mask = (mask > 0.5).astype(np.uint8)) 128 | 129 | lines = getHoughLines(edges_masked, min_line_size=0.25*min(min_area_rect[1])) 130 | 131 | lines_a, lines_b = parseHoughLines(lines, top_two_angles, angle_threshold_deg=35) 132 | if len(lines_a) < 2 or len(lines_b) < 2: 133 | print('fail2', lines_a, lines_b) 134 | return img 135 | 136 | # plotHoughLines(img, lines, color=(255,255,255), line_thickness=1) 137 | # plotHoughLines(img, lines_a, color=(0,0,255)) 138 | # plotHoughLines(img, lines_b, color=(0,255,0)) 139 | 140 | a = time() 141 | for i2 in range(2): 142 | for i in range(5): 143 | corners = chooseRandomGoodQuad(lines_a, lines_b, median_contour) 144 | 145 | # warp_img, M = getTileImage(img_orig, corners.astype(np.float32),tile_buffer=16, tile_res=16) 146 | M = getTileTransform(corners.astype(np.float32),tile_buffer=16, tile_res=16) 147 | 148 | # Warp lines and draw them on warped image 149 | all_lines = np.vstack([lines_a[:,:2], lines_a[:,2:], lines_b[:,:2], lines_b[:,2:]]).astype(np.float32) 150 | warp_pts = cv2.perspectiveTransform(all_lines[None,:,:], M) 151 | warp_pts = warp_pts[0,:,:] 152 | warp_lines_a = np.hstack([warp_pts[:len(lines_a),:], warp_pts[len(lines_a):2*len(lines_a),:]]) 153 | warp_lines_b = np.hstack([warp_pts[2*len(lines_a):2*len(lines_a)+len(lines_b),:], warp_pts[2*len(lines_a)+len(lines_b):,:]]) 154 | 155 | 156 | # Get thetas of warped lines 157 | thetas_a = np.array([getSegmentTheta(line) for line in warp_lines_a]) 158 | thetas_b = np.array([getSegmentTheta(line) for line in warp_lines_b]) 159 | median_theta_a = (np.median(thetas_a*180/np.pi)) 160 | median_theta_b = (np.median(thetas_b*180/np.pi)) 161 | 162 | # Gradually relax angle threshold over N iterations 163 | if i < 20: 164 | warp_angle_threshold = 0.03 165 | elif i < 30: 166 | warp_angle_threshold = 0.1 167 | elif i < 50: 168 | warp_angle_threshold = 0.3 169 | elif i < 70: 170 | warp_angle_threshold = 0.5 171 | elif i < 80: 172 | warp_angle_threshold = 1.0 173 | else: 174 | warp_angle_threshold = 2.0 175 | if ((angleCloseDeg(abs(median_theta_a), 0, warp_angle_threshold) and 176 | angleCloseDeg(abs(median_theta_b), 90, warp_angle_threshold)) or 177 | (angleCloseDeg(abs(median_theta_a), 90, warp_angle_threshold) and 178 | angleCloseDeg(abs(median_theta_b), 0, warp_angle_threshold))): 179 | break 180 | # else: 181 | # print('iter %d: %.2f %.2f' % (i, abs(median_theta_a), abs(median_theta_b))) 182 | 183 | warp_img, M = getTileImage(img_orig, corners.astype(np.float32),tile_buffer=16, tile_res=16) 184 | 185 | lines_x, lines_y, step_x, step_y = getWarpCheckerLines(warp_img) 186 | if len(lines_x) > 0: 187 | break 188 | 189 | warp_img, M = getTileImage(img_orig, corners.astype(np.float32),tile_buffer=16, tile_res=16) 190 | 191 | for corner in corners: 192 | cv2.circle(img, tuple(map(int,corner)), 5, (255,150,150),-1) 193 | 194 | if len(lines_x) > 0: 195 | warp_corners, all_warp_corners = getRectChessCorners(lines_x, lines_y) 196 | tile_centers = all_warp_corners + np.array([step_x/2.0, step_y/2.0]) # Offset from corner to tile centers 197 | M_inv = np.matrix(np.linalg.inv(M)) 198 | real_corners, all_real_tile_centers = getOrigChessCorners(warp_corners, tile_centers, M_inv) 199 | 200 | tile_res = 64 # Each tile has N pixels per side 201 | tile_buffer = 1 202 | warp_img, better_M = getTileImage(img_orig2, real_corners, tile_buffer=tile_buffer, tile_res=tile_res) 203 | # Further refine rectified image 204 | warp_img, was_rotated, refine_M = reRectifyImages(warp_img) 205 | # combined_M = better_M 206 | combined_M = np.matmul(refine_M,better_M) 207 | M_inv = np.matrix(np.linalg.inv(combined_M)) 208 | 209 | # Get better_M based corners 210 | hlines = vlines = (np.arange(8)+tile_buffer)*tile_res 211 | hcorner = (np.array([0,8,8,0])+tile_buffer)*tile_res 212 | vcorner = (np.array([0,0,8,8])+tile_buffer)*tile_res 213 | ideal_corners = np.vstack([hcorner,vcorner]).T 214 | ideal_all_corners = np.array(list(itertools.product(hlines, vlines))) 215 | ideal_tile_centers = ideal_all_corners + np.array([tile_res/2.0, tile_res/2.0]) # Offset from corner to tile centers 216 | 217 | real_corners, all_real_tile_centers = getOrigChessCorners(ideal_corners, ideal_tile_centers, M_inv) 218 | 219 | # Get final refined rectified warped image for saving 220 | warp_img, _ = getTileImage(img_orig2, real_corners, tile_buffer=tile_buffer, tile_res=tile_res) 221 | 222 | cv2.polylines(img, [real_corners.astype(np.int32)], True, (150,50,255), thickness=4) 223 | cv2.polylines(img, [all_real_tile_centers.astype(np.int32)], False, (0,50,255), thickness=1) 224 | 225 | 226 | img_masked_full = cv2.bitwise_and(img,img,mask = (mask > 0.5).astype(np.uint8)) 227 | img_masked = cv2.addWeighted(img,0.2,img_masked_full,0.8,0) 228 | 229 | drawMinAreaRect(img_masked, min_area_rect) 230 | 231 | return img_masked 232 | 233 | 234 | def processVideo(filename, func=findAndDrawHough, rate=1): 235 | 236 | # Define the codec and create VideoWriter object 237 | fourcc = cv2.VideoWriter_fourcc(*'XVID') 238 | cap = cv2.VideoCapture(filename) 239 | 240 | img_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) 241 | img_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) 242 | 243 | img_rescale_ratio = 1.0 244 | 245 | out_size = (int(img_width*img_rescale_ratio), int(img_height*img_rescale_ratio)) 246 | 247 | output_filename = 'output3_%s.avi' % (filename[:-4]) 248 | print("Writing to %s at scale %s" % (output_filename, out_size)) 249 | # out = cv2.VideoWriter(output_filename,fourcc, 20.0, (384,216)) # 0.2 250 | # out = cv2.VideoWriter(output_filename,fourcc, 20.0, (576,324)) # 0.3 251 | out = cv2.VideoWriter(output_filename,fourcc, 20.0, out_size) 252 | 253 | i = 0 254 | while(cap.isOpened()): 255 | ret, frame = cap.read() 256 | if not np.any(frame): 257 | break 258 | i+=1 259 | 260 | 261 | img = cv2.resize(frame,None,fx=img_rescale_ratio, fy=img_rescale_ratio, interpolation = cv2.INTER_AREA) 262 | 263 | if (i == 1): 264 | if img.shape[0] != out_size[1] or img.shape[1] != out_size[0]: 265 | print(img.shape, out_size) 266 | 267 | img_masked = func(img) 268 | 269 | out.write(img_masked) 270 | cv2.imshow('Masked',img_masked) 271 | 272 | if cv2.waitKey(rate) & 0xFF == ord('q'): 273 | break 274 | 275 | cap.release() 276 | out.release() 277 | cv2.destroyAllWindows() 278 | 279 | def main(filenames): 280 | for filename in filenames: 281 | print("Processing %s" % filename) 282 | img = cv2.imread(filename) 283 | 284 | img_diag_size = int(np.ceil(np.sqrt(img.shape[0]*img.shape[0] + img.shape[1]*img.shape[1]))) 285 | print(img_diag_size) 286 | 287 | gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY) 288 | edges = cv2.Canny(gray,100,650,apertureSize = 3) 289 | min_img_side = min(img.shape[:2]) 290 | minLineLength = min_img_side/8 291 | maxLineGap = min_img_side/10 292 | threshold = int(min_img_side/8) 293 | print(minLineLength, maxLineGap, threshold) 294 | lines = cv2.HoughLinesP(edges,rho=1,theta=np.pi/180, 295 | threshold=threshold, minLineLength=minLineLength, maxLineGap=maxLineGap) 296 | 297 | # colors = np.random.random([lines.shape[0],3])*255 298 | colors = [ 299 | [255,0,0], 300 | [0,255,0], 301 | [255,255,0], 302 | [0,0,255], 303 | [255,0,255], 304 | [0,255,255], 305 | [255,255,255], 306 | ] 307 | 308 | 309 | hough_img = np.zeros([2*img_diag_size/4, 180]) # -90 to 90 deg, -rho_max to rho_max 310 | 311 | if any(lines): 312 | rhothetas = np.zeros([lines.shape[0], 2]) 313 | for i, (x1,y1,x2,y2) in enumerate(lines[:,0,:]): 314 | color = list(map(int,colors[i%len(colors)])) # dtype needs to be int, not np.int32 315 | cv2.line(img,(x1,y1),(x2,y2), color,2) 316 | 317 | rho, theta = getRhoTheta((x1,y1,x2,y2)) 318 | rhothetas[i,:] = rho, theta 319 | img_rho, img_theta = (int(theta*180/np.pi + 90), int((rho+img_diag_size)/4)) 320 | cv2.circle(hough_img, (img_rho, img_theta), 3, (255,0,0),-1) 321 | 322 | plot(rhothetas[:,1]*180/np.pi, rhothetas[:,0], 'o') 323 | print(hough_img.shape) 324 | xlabel('theta (deg)') 325 | ylabel('rho') 326 | 327 | 328 | 329 | 330 | # Using matplotlib 331 | # imshow(img) 332 | # show() 333 | 334 | # Using opencv 335 | cv2.imshow('image %dx%d' % (img.shape[1],img.shape[0]),img) 336 | cv2.imshow('edges',edges) 337 | cv2.imshow('hough',hough_img) 338 | cv2.moveWindow('hough', 0,0) 339 | # cv2.waitKey(0) 340 | axis('equal') 341 | show() 342 | cv2.destroyAllWindows() 343 | 344 | if __name__ == '__main__': 345 | if len(sys.argv) > 1: 346 | filenames = sys.argv[1:] 347 | else: 348 | filenames = ['input2/27.jpg'] 349 | # filenames = ['input/2.jpg', 'input/6.jpg', 'input/17.jpg'] 350 | # filenames = ['input/1.jpg', 'input/2.jpg', 'input/3.jpg', 'input_fails/37.jpg', 'input_fails/38.jpg'] 351 | # filenames = ['input_fails/37.jpg', 'input_fails/38.jpg'] 352 | # main(filenames) 353 | # processVideo('chess1.mp4') 354 | # processVideo('chess2.mp4', func=findAndDrawTile) 355 | processVideo('chess1.mp4', func=findAndDrawMask) 356 | # processVideo('chess3.mp4', func=findAndDrawChessboard) 357 | print('Done.') -------------------------------------------------------------------------------- /image_segment.py: -------------------------------------------------------------------------------- 1 | # Image segmentation 2 | # Given rectified image with known tile boundaries 3 | # Segment image into background (black/white tiles?) 4 | # and dark or light pieces 5 | # 6 | # Some options include K-means clustering, watershed segmentation, texture segmentation, perhaps a combination 7 | 8 | import PIL.Image 9 | import matplotlib.pyplot as plt 10 | import cv2 11 | import numpy as np 12 | import itertools 13 | import os 14 | from skimage import color 15 | from sklearn.cluster import KMeans 16 | from skimage import exposure 17 | np.set_printoptions(precision=2, linewidth=400, suppress=True) # Better printing of arrays 18 | 19 | def getIdealCorners(tile_res, tile_buffer): 20 | hlines = vlines = (np.arange(9)+tile_buffer)*tile_res 21 | return np.array(list(itertools.product(hlines, vlines))) 22 | 23 | def getIdealCheckerboardPattern(tile_res, tile_buffer): 24 | side_len = tile_res*(8+2*tile_buffer) 25 | quadOne = np.ones([tile_res,tile_res], dtype=np.uint8) 26 | quadZero = np.zeros([tile_res,tile_res], dtype=np.uint8) 27 | kernel = np.vstack([np.hstack([quadOne,quadZero]), np.hstack([quadZero,quadOne])]) 28 | kernel = np.tile(kernel,(4,4)) # Becomes 8x8 alternating grid 29 | return kernel 30 | 31 | def getTile(img, i,j,tile_res): 32 | """Assumes no buffer in image""" 33 | return img[tile_res*i:tile_res*(i+1),tile_res*j:tile_res*(j+1)] 34 | 35 | 36 | if __name__ == '__main__': 37 | PLOT_RESULTS = True 38 | 39 | input_folder = "rectified" 40 | 41 | tile_res = 64 42 | tile_buffer = 1 43 | 44 | side_len = 8*tile_res 45 | buffer_size = tile_buffer*tile_res 46 | 47 | filename ="%d.png" % 7 48 | filepath = "%s/%s" % (input_folder,filename) 49 | print("Segmenting %s..." % filename) 50 | img_orig = np.array(PIL.Image.open(filepath).convert('RGB')) 51 | 52 | img_h, img_w, _ = img_orig.shape 53 | 54 | # Bilateral smooth image 55 | img = img_orig 56 | # img = cv2.blur(img,ksize=(5,5)) 57 | # img = cv2.bilateralFilter(img,int(tile_res/4),75,75) 58 | 59 | ideal_corners = getIdealCorners(tile_res, tile_buffer) 60 | ideal_checkerboard = getIdealCheckerboardPattern(tile_res, tile_buffer) 61 | ideal_checkerboard_corners = getIdealCorners(tile_res, 0) 62 | 63 | white_only_mask = ideal_checkerboard 64 | black_only_mask = (~white_only_mask.astype(bool)).astype(np.uint8) 65 | 66 | img_checkerboard = img[buffer_size:-buffer_size, buffer_size:-buffer_size] 67 | img_checkerboard = cv2.medianBlur(img_checkerboard,7) 68 | img_checkerboard = cv2.bilateralFilter(img_checkerboard,int(tile_res/4),75,75) 69 | 70 | # Local Histogram Equalization of checkerboard 71 | ycrcb = cv2.cvtColor(img_checkerboard, cv2.COLOR_RGB2YCR_CB) 72 | # ycrcb[:,:,0] = cv2.equalizeHist(ycrcb[:,:,0].astype(np.uint8)) 73 | clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(4,4)) 74 | # ycrcb[:,:,0] = clahe.apply(ycrcb[:,:,0].astype(np.uint8)) 75 | img_checkerboard = cv2.cvtColor(ycrcb, cv2.COLOR_YCR_CB2RGB) 76 | # ycrcb = cv2.medianBlur(ycrcb,7) 77 | # ycrcb = cv2.bilateralFilter(ycrcb,int(tile_res/4),75,75) 78 | 79 | 80 | responseA = cv2.bitwise_and(img_checkerboard, img_checkerboard, mask=white_only_mask) 81 | responseB = cv2.bitwise_and(img_checkerboard, img_checkerboard, mask=black_only_mask) 82 | 83 | img_checkerboard_gray = np.array(PIL.Image.fromarray(img_checkerboard).convert('L')) 84 | # img_checkerboard_gray = cv2.equalizeHist(img_checkerboard_gray) 85 | # clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(5,5)) 86 | # img_checkerboard_gray = clahe.apply(img_checkerboard_gray) 87 | 88 | # Get La*b* colorspace 89 | # lab = color.rgb2lab(img_checkerboard) 90 | # hsv = color.rgb2hsv(img_checkerboard) 91 | 92 | # Remove intensity changes 93 | # Local Histogram Equalization 94 | # lab[:,:,0] = cv2.equalizeHist(lab[:,:,0].astype(np.uint8)) 95 | # lab[:,:,0] = exposure.equalize_adapthist(lab[:,:,0]*0.01, clip_limit=0.3)*100 96 | # clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(4,4)) 97 | # lab[:,:,0] = clahe.apply(lab[:,:,0].astype(np.uint8)) 98 | 99 | # img_checkerboard_equalized = color.lab2rgb(lab) 100 | 101 | # edges = cv2.Canny(lab[:,:,0].astype(np.uint8),200,500,apertureSize = 3, L2gradient=False) # Better thresholds 102 | low_threshold = 30 103 | edges = cv2.Canny(img_checkerboard_gray,low_threshold,3*low_threshold,apertureSize = 3, L2gradient=False) # Better thresholds 104 | 105 | 106 | rgb_flat = img_checkerboard.reshape([img_checkerboard.shape[0]*img_checkerboard.shape[1],img_checkerboard.shape[2]]) 107 | # rgb_flat[white_only_mask.astype(bool).flatten()] = 0 108 | # lab_flat = lab.reshape([lab.shape[0]*lab.shape[1],lab.shape[2]]) 109 | # hsv_flat = hsv.reshape([hsv.shape[0]*hsv.shape[1],hsv.shape[2]]) 110 | 111 | # a_star_white = lab[:,:,1][white_only_mask.astype(bool)] 112 | # b_star_white = lab[:,:,2][white_only_mask.astype(bool)] 113 | 114 | # a_star_black = lab[:,:,1][black_only_mask.astype(bool)] 115 | # b_star_black = lab[:,:,2][black_only_mask.astype(bool)] 116 | 117 | # a_star = lab[:,:,1].flatten() 118 | # b_star = lab[:,:,2].flatten() 119 | 120 | # K-means cluster into 4 parts (black tile, white tile, black piece, white piece) 121 | 122 | print("Start K-means") 123 | random_state = 1 124 | clt = KMeans(n_clusters=4, random_state=random_state) 125 | clt.fit(ycrcb[:,:,1].reshape([ycrcb.shape[0]*ycrcb.shape[1],-1])) 126 | # clt.fit(ycrcb.reshape([ycrcb.shape[0]*ycrcb.shape[1],-1])) 127 | # clt.fit(img_checkerboard.reshape([img_checkerboard.shape[0]*img_checkerboard.shape[1],-1])) 128 | y_pred = clt.labels_ 129 | print("End K-means", y_pred.shape) 130 | 131 | deviations = np.zeros([8,8]) 132 | for i in range(8): 133 | for j in range(8): 134 | tile = getTile(img_checkerboard_gray,i,j,tile_res) 135 | edge_tile = getTile(edges,i,j,tile_res) 136 | inner_tile = tile[8:-8,8:-8] 137 | inner_edge_tile = edge_tile[16:-16,16:-16] 138 | if np.sum(inner_edge_tile) > 20: 139 | deviations[i,j] = np.std(inner_tile) 140 | 141 | 142 | 143 | if PLOT_RESULTS: 144 | print("Plotting") 145 | plt.figure(filename) 146 | 147 | plt.subplot(331) 148 | plt.imshow(img_orig) 149 | plt.plot(ideal_corners[:,0], ideal_corners[:,1], 'ro', ms=3) 150 | plt.title('Input rectified image') 151 | plt.axis([0,img_w,img_h, 0]) 152 | 153 | plt.subplot(332) 154 | plt.imshow(responseA) 155 | plt.plot(ideal_checkerboard_corners[:,0], ideal_checkerboard_corners[:,1], 'ro', ms=3) 156 | plt.title('White chessboard only') 157 | plt.axis([0,side_len, side_len, 0]) 158 | 159 | plt.subplot(333) 160 | plt.imshow(img_checkerboard_gray,cmap=plt.cm.gray) 161 | # plt.imshow(responseB) 162 | # plt.plot(ideal_checkerboard_corners[:,0], ideal_checkerboard_corners[:,1], 'ro', ms=3) 163 | # plt.title('Black chessboard only') 164 | plt.axis([0,side_len, side_len, 0]) 165 | 166 | plt.subplot(334) 167 | plt.imshow(img_checkerboard) 168 | plt.axis([0,side_len, side_len, 0]) 169 | 170 | # plt.hist2d(a_star, b_star, (100,100)) 171 | # plt.title('La*b* : all') 172 | # plt.xlabel('A*') 173 | # plt.ylabel('B*') 174 | # plt.colorbar() 175 | 176 | plt.subplot(335) 177 | plt.imshow(edges) 178 | plt.axis([0,side_len, side_len, 0]) 179 | 180 | # Plot only random N points 181 | # subset = np.random.choice(a_star.shape[0], 2000,replace=False) 182 | # x = np.hstack([np.ones([subset.shape[0],1])*50, clt.cluster_centers_[y_pred[subset]]]) 183 | # x = np.swapaxes(np.atleast_3d(x),1,2) 184 | # x = np.squeeze(color.lab2rgb(x)) 185 | # plt.scatter(a_star[subset], b_star[subset], c=y_pred[subset], edgecolor='') 186 | 187 | plt.subplot(336) 188 | plt.imshow(img_checkerboard) 189 | plt.plot(ideal_checkerboard_corners[:,0], ideal_checkerboard_corners[:,1], 'ro', ms=3) 190 | for i in range(8): 191 | for j in range(8): 192 | if deviations[i,j] > 8: 193 | plt.text(tile_res*(j+0.5)-20, tile_res*(i+0.5), 194 | '%.1f' % deviations[i,j], color='black', size=10, fontweight='heavy'); 195 | 196 | 197 | plt.title('Standard deviation of tiles') 198 | plt.axis('square') 199 | plt.axis([0,side_len, side_len, 0]) 200 | 201 | plt.subplot(337) 202 | kmeans_result = y_pred.reshape(ideal_checkerboard.shape) 203 | plt.imshow(kmeans_result) 204 | plt.colorbar() 205 | 206 | plt.subplot(338) 207 | kmeans_result_fix = kmeans_result.copy() 208 | kmeans_result_fix[white_only_mask.astype(bool)] = (~kmeans_result_fix[white_only_mask.astype(bool)].astype(bool)).astype(np.uint8) 209 | plt.imshow(kmeans_result_fix) 210 | plt.colorbar() 211 | 212 | plt.subplot(339) 213 | plt.imshow(ycrcb[:,:,1]) 214 | plt.colorbar() 215 | 216 | plt.show() -------------------------------------------------------------------------------- /inform_hough_on_pts.py: -------------------------------------------------------------------------------- 1 | # Given a set of ML pruned saddle points, remove outliers or keep only chessboard area 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | 5 | 6 | pts = np.array([[237, 332], 7 | [242, 287], 8 | [245, 263], 9 | [247, 360], 10 | [249, 337], 11 | [252, 314], 12 | [255, 290], 13 | [257, 389], 14 | [258, 266], 15 | [260, 366], 16 | [262, 342], 17 | [268, 294], 18 | [271, 269], 19 | [274, 372], 20 | [280, 323], 21 | [286, 272], 22 | [286, 404], 23 | [291, 353], 24 | [295, 328], 25 | [298, 302], 26 | [298, 437], 27 | [301, 411], 28 | [304, 386], 29 | [311, 333], 30 | [318, 279], 31 | [318, 420], 32 | [321, 393], 33 | [324, 366], 34 | [329, 339], 35 | [332, 311], 36 | [336, 283]]) 37 | 38 | outliers = np.array([ 39 | [50,23], 40 | [30,63], 41 | [400,203], 42 | [250,370], 43 | [500,303], 44 | ]) 45 | 46 | pts = np.vstack([pts, outliers]) 47 | 48 | 49 | # plt.plot(outliers[:,0],outliers[:,1],'rx') 50 | # plt.plot(pts[:,0],pts[:,1],'.') 51 | # plt.show() 52 | N = len(pts) 53 | 54 | new_order = np.arange(N) 55 | np.random.shuffle(new_order) 56 | pts = pts[new_order,:] 57 | 58 | print pts 59 | 60 | 61 | 62 | import time 63 | 64 | ta = time.time() 65 | 66 | outlier = calculateOutliers(pts) 67 | pred_outliers = pts[outlier,:] 68 | 69 | tb = time.time() 70 | 71 | proc_time = tb-ta 72 | print("Processed in %.2f ms" % (proc_time*1e3)) 73 | 74 | 75 | # plt.plot(outliers[:,0],outliers[:,1],'rx') 76 | plt.plot(pred_outliers[:,0],pred_outliers[:,1],'rx') 77 | plt.plot(pts[:,0],pts[:,1],'.') 78 | plt.show() -------------------------------------------------------------------------------- /line_intersection.py: -------------------------------------------------------------------------------- 1 | import PIL.Image 2 | import cv2 3 | import numpy as np 4 | import itertools 5 | 6 | # Calculate intersections 7 | def line_intersect(a1, a2, b1, b2): 8 | T = np.array([[0, -1], [1, 0]]) 9 | da = np.atleast_2d(a2 - a1) 10 | db = np.atleast_2d(b2 - b1) 11 | dp = np.atleast_2d(a1 - b1) 12 | dap = np.dot(da, T) 13 | denom = np.sum(dap * db, axis=1) 14 | num = np.sum(dap * dp, axis=1) 15 | return np.atleast_2d(num / denom).T * db + b1 16 | 17 | def getAllLineIntersections(linesA, linesB): 18 | # get all pairings of lines 19 | pairings = np.array(list(itertools.product(range(linesA.shape[0]),range(linesB.shape[0])))) 20 | return line_intersect(linesA[pairings[:,0],:2], linesA[pairings[:,0],2:], linesB[pairings[:,1],:2], linesB[pairings[:,1],2:]) 21 | 22 | def prunePoints(pts, max_dist2=5**2): 23 | # Prune points away that are close to each other 24 | # preferring points that come earlier in array order 25 | good_pts = np.ones(pts.shape[0], dtype=bool) 26 | for i in range(pts.shape[0]): 27 | if ~good_pts[i]: 28 | continue 29 | for j in range(i+1,pts.shape[0]): 30 | d2 = np.sum((pts[j] - pts[i])**2) 31 | if (d2 < max_dist2): # within (N pixels)**2 of another point 32 | good_pts[j] = False 33 | 34 | return pts[good_pts] 35 | 36 | 37 | if __name__ == '__main__': 38 | import matplotlib.pyplot as plt 39 | a = np.array( 40 | [[265, 192, 517, 389], 41 | [210, 219, 427, 418], 42 | [352, 164, 594, 318], 43 | [254, 219, 459, 391], 44 | [295, 182, 544, 363], 45 | [330, 176, 570, 341], 46 | [360, 142, 617, 297], 47 | [332, 178, 541, 322], 48 | [295, 183, 505, 336], 49 | [252, 217, 488, 415], 50 | [360, 168, 567, 300], 51 | [289, 291, 455, 443], 52 | [232, 240, 454, 444], 53 | [286, 209, 478, 359]]) 54 | b = np.array( 55 | [[253, 348, 605, 120], 56 | [274, 374, 604, 148], 57 | [318, 386, 558, 212], 58 | [230, 326, 578, 112], 59 | [348, 413, 650, 181], 60 | [207, 275, 440, 146], 61 | [209, 304, 531, 118], 62 | [319, 387, 560, 212], 63 | [256, 311, 579, 113], 64 | [260, 345, 606, 120], 65 | [191, 285, 387, 176], 66 | [234, 289, 530, 118], 67 | [210, 305, 531, 119], 68 | [275, 375, 604, 149]]) 69 | 70 | pts = getAllLineIntersections(a, b) 71 | print("Found %d points" % len(pts)) 72 | 73 | # Plot lines 74 | for line in a: 75 | x1, y1, x2, y2 = line 76 | plt.plot([x1,x2], [y1,y2],'b') 77 | for line in b: 78 | x1, y1, x2, y2 = line 79 | plt.plot([x1,x2], [y1,y2],'g') 80 | 81 | # Plot points 82 | plt.plot(pts[:,0], pts[:,1], 'ro',ms=8) 83 | 84 | plt.show() 85 | 86 | def getCorners(chess_pts, top_dirs): 87 | """top_dirs are the two top direction vectors for the chess board lines""" 88 | d_norm_a = top_dirs[0] 89 | vals = chess_pts.dot(d_norm_a) 90 | a = chess_pts[np.argmin(vals),:] 91 | b = chess_pts[np.argmax(vals),:] 92 | 93 | dist = (b-a) 94 | d_norm = np.array([-dist[1], dist[0]]) 95 | d_norm /= np.sqrt(np.sum(d_norm**2)) 96 | 97 | # print(d_norm) 98 | vals = chess_pts.dot(d_norm) 99 | # print(vals) 100 | c = chess_pts[np.argmin(vals),:] 101 | d = chess_pts[np.argmax(vals),:] 102 | 103 | corners = np.vstack([a,c,b,d]).astype(np.float32) 104 | return corners 105 | 106 | def getRectifiedChessLines(img): 107 | """Given a warped axis-aligned image of a chessboard, return internal line crossings""" 108 | # TODO: Fix awkward conversion 109 | # Convert RGB numpy array to image, then to grayscale image, then back to numpy array 110 | img_gray = np.array(PIL.Image.fromarray(img).convert('L')) 111 | img_gray = cv2.bilateralFilter(img_gray,15,75,75) 112 | 113 | # Find gradients 114 | sobelx = cv2.Sobel(img_gray,cv2.CV_64F,1,0,ksize=5) 115 | sobely = cv2.Sobel(img_gray,cv2.CV_64F,0,1,ksize=5) 116 | 117 | sobelx_pos = sobelx.copy() 118 | sobelx_pos[sobelx <= 0] = 0 119 | sobelx_neg = sobelx.copy() 120 | sobelx_neg[sobelx > 0] = 0 121 | 122 | sobely_pos = sobely.copy() 123 | sobely_pos[sobely <= 0] = 0 124 | sobely_neg = sobely.copy() 125 | sobely_neg[sobely > 0] = 0 126 | 127 | checker_x = np.sum(sobelx_pos, axis=0) * np.sum(-sobelx_neg, axis=0) 128 | checker_x = skeletonize_1d(checker_x) 129 | 130 | checker_y = np.sum(sobely_pos, axis=1) * np.sum(-sobely_neg, axis=1) 131 | checker_y = skeletonize_1d(checker_y) 132 | 133 | x_lines = np.argwhere(checker_x).flatten() 134 | y_lines = np.argwhere(checker_y).flatten() 135 | 136 | x_diff = np.diff(x_lines) 137 | y_diff = np.diff(y_lines) 138 | 139 | step_x_pred = np.median(x_diff) 140 | step_y_pred = np.median(y_diff) 141 | 142 | # Remove internal outlier lines that have the wrong step size 143 | x_good = np.ones(x_lines.shape,dtype=bool) 144 | y_good = np.ones(y_lines.shape,dtype=bool) 145 | 146 | x_good[1:] = abs(x_diff - step_x_pred) < 20 147 | y_good[1:] = abs(y_diff - step_y_pred) < 20 148 | 149 | x_keep = np.ones(x_lines.shape,dtype=bool) 150 | y_keep = np.ones(y_lines.shape,dtype=bool) 151 | 152 | for i in range(x_good.size-1): 153 | if ~np.any(x_good[i:i+2]): 154 | x_keep[i] = False 155 | for i in range(y_good.size-1): 156 | if ~np.any(y_good[i:i+2]): 157 | y_keep[i] = False 158 | 159 | x_lines = x_lines[x_keep] 160 | y_lines = y_lines[y_keep] 161 | 162 | if len(x_lines) < 7 or len(y_lines) < 7: 163 | return [], [], [], [] 164 | 165 | # Select set of 7 adjacent lines with max sum score 166 | x_scores = np.zeros(x_lines.shape[0]-7+1) 167 | for i in range(0,x_lines.shape[0]-7+1): 168 | x_scores[i] = np.sum(checker_x[x_lines[i:i+7]]) 169 | x_start = np.argmax(x_scores) 170 | strongest_x_lines = range(x_start,x_start+7) 171 | 172 | y_scores = np.zeros(y_lines.shape[0]-7+1) 173 | for i in range(0,y_lines.shape[0]-7+1): 174 | y_scores[i] = np.sum(checker_y[y_lines[i:i+7]]) 175 | y_start = np.argmax(y_scores) 176 | strongest_y_lines = range(y_start,y_start+7) 177 | 178 | # TODO: Sanity check areas between lines for consistent color when choosing? 179 | 180 | # Choose best internal 7 chessboard lines 181 | lines_x = x_lines[strongest_x_lines] 182 | lines_y = y_lines[strongest_y_lines] 183 | 184 | # Add outer chessboard edges assuming consistent step size 185 | step_x = np.median(np.diff(lines_x)) 186 | step_y = np.median(np.diff(lines_y)) 187 | 188 | lines_x = np.hstack([lines_x[0]-step_x, lines_x, lines_x[-1]+step_x]) 189 | lines_y = np.hstack([lines_y[0]-step_y, lines_y, lines_y[-1]+step_y]) 190 | 191 | return lines_x, lines_y, step_x, step_y 192 | 193 | def skeletonize_1d(arr, win=50): 194 | """return skeletonized 1d array (thin to single value, favor to the right)""" 195 | _arr = arr.copy() # create a copy of array to modify without destroying original 196 | # Go forwards 197 | for i in range(_arr.size-1): 198 | if _arr[i] == 0: 199 | continue 200 | # Will right-shift if they are the same 201 | if np.any(arr[i] <= arr[i+1:i+win+1]): 202 | _arr[i] = 0 203 | 204 | # Go reverse 205 | for i in np.arange(_arr.size-1, 0,-1): 206 | if _arr[i] == 0: 207 | continue 208 | 209 | if np.any(arr[max(0,i-win):i] > arr[i]): 210 | _arr[i] = 0 211 | return _arr 212 | 213 | def getRectChessCorners(lines_x, lines_y): 214 | pairs = np.array(list(itertools.product(range(8),range(8)))) 215 | all_warp_corners = np.vstack([lines_x[pairs[:,0]], lines_y[pairs[:,1]]]).T 216 | warp_corners = np.array([ 217 | [lines_x[0], lines_y[0]], 218 | [lines_x[-1], lines_y[0]], 219 | [lines_x[-1], lines_y[-1]], 220 | [lines_x[0], lines_y[-1]] 221 | ]) 222 | return warp_corners[:,:2].astype(np.float32), all_warp_corners[:,:2].astype(np.float32) 223 | 224 | def getOrigChessCorners(warp_corners, all_warp_corners, M_inv): 225 | all_stack = np.hstack([all_warp_corners, np.ones([all_warp_corners.shape[0],1])]).T 226 | all_real_corners = (M_inv * all_stack).T 227 | all_real_corners = all_real_corners / all_real_corners[:,2] 228 | 229 | stack = np.hstack([warp_corners, np.ones([4,1])]).T 230 | real_corners = (M_inv * stack).T 231 | real_corners = real_corners / real_corners[:,2] # Normalize by z 232 | return real_corners[:,:2].astype(np.float32), all_real_corners[:,:2].astype(np.float32) 233 | 234 | def getTileImage(input_img, quad_corners, tile_buffer=0, tile_res=64): 235 | # Add N tile worth buffer on outside edge, such that 236 | # CV/ML algorithms could potentially use this data for better predictions 237 | ideal_quad_corners = np.array([[0,0], [1,0], [1,1], [0,1]], dtype=np.float32) 238 | 239 | main_len = tile_res*(ideal_quad_corners*8+tile_buffer) 240 | side_len = tile_res*(8+2*tile_buffer) 241 | 242 | M = cv2.getPerspectiveTransform(quad_corners, main_len) 243 | out_img = cv2.warpPerspective(np.array(input_img), M, 244 | (side_len, side_len)) 245 | return out_img, M 246 | 247 | def getTileTransform(quad_corners, tile_buffer=0, tile_res=64): 248 | # Add N tile worth buffer on outside edge, such that 249 | # CV/ML algorithms could potentially use this data for better predictions 250 | ideal_quad_corners = np.array([[0,0], [1,0], [1,1], [0,1]], dtype=np.float32) 251 | 252 | main_len = tile_res*(ideal_quad_corners*8+tile_buffer) 253 | side_len = tile_res*(8+2*tile_buffer) 254 | 255 | M = cv2.getPerspectiveTransform(quad_corners, main_len) 256 | return M 257 | 258 | def getSegments(v, eps = 2): 259 | # Get segment mask given a vector v, segments are values 260 | # withing eps distance of each other 261 | n = len(v) 262 | segment_mask = np.zeros(n,dtype=np.uint16) 263 | k = 1 264 | for i in range(n): 265 | if segment_mask[i] != 0: 266 | continue 267 | segment_mask[i] = k 268 | for j in range(i+1,n): 269 | if abs(v[j] - v[i]) < eps: 270 | segment_mask[j] = k 271 | k += 1 272 | return segment_mask-1, k-1 273 | 274 | def mergePairs(pairs): 275 | if len(pairs) == 1: 276 | return pairs[0] 277 | 278 | vals = pairs[0] 279 | for i in range(1,len(pairs)): 280 | v_end = vals[-1] 281 | next_idx = np.argwhere(pairs[i] == v_end) 282 | if len(next_idx) > 0: 283 | vals = np.hstack([vals[:-1], pairs[i,next_idx[0][0]:]]) 284 | return vals 285 | 286 | def getBestEqualSpacing(vals, min_pts=7, eps=4, std_min=2): 287 | assert(min_pts>3) 288 | # Finds all combinations of triplets of points in vals where 289 | # the standard deviation is less than std_min, then merges 290 | # them into longer equally spaced sets and returns 291 | # the one with the largest equal spacing that has at least n_pts 292 | n_pts = 3 293 | pairs = np.array([k for k in itertools.combinations(vals, n_pts) if np.std(np.diff(k)) < std_min and np.mean(np.diff(k)) > 8]) 294 | 295 | spacings = np.array([np.mean(np.diff(k)) for k in pairs]) 296 | segments, num_segments = getSegments(spacings, eps) 297 | best_spacing = [] 298 | best_mean = 0 299 | for i in range(num_segments): 300 | merged = mergePairs(pairs[segments==i]) 301 | spacing_mean = np.mean(np.diff(merged)) 302 | 303 | # Keep the largest equally spaced set that has min number of points 304 | if len(merged) >= min_pts and spacing_mean > best_mean: 305 | best_spacing = merged 306 | best_mean = spacing_mean 307 | 308 | return best_spacing -------------------------------------------------------------------------------- /makeGif.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | filename=$1 3 | filename_out=$2 4 | 5 | ffmpeg -y -ss 5 -t 7 -i $1 \ 6 | -vf fps=10,scale=320:-1:flags=lanczos,palettegen palette.png && \ 7 | ffmpeg -ss 5 -t 7 -i $1 -i palette.png \ 8 | -filter_complex "fps=10,scale=640:-1:flags=lanczos[x];[x][1:v]paletteuse" $2 -y -------------------------------------------------------------------------------- /make_videos.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | filename=$1; 3 | filepath="results/${filename}_vidstream_frames" 4 | 5 | # Warped 6 | ffmpeg -i ${filepath}/ml_warp_frame_%03d.jpg -c:v libx264 -vf "fps=25,format=yuv420p" -t 30 ${filename}_ml_warp.avi -y 7 | 8 | # Better 9 | ffmpeg -i ${filepath}/ml_frame_%03d.jpg -c:v libx264 -vf "fps=25,format=yuv420p" -t 30 ${filename}_ml.avi -y 10 | 11 | 12 | # Both side by side 13 | ffmpeg -i ${filename}_ml.avi -i ${filename}_ml_warp.avi -filter_complex '[0:v]pad=iw*2:ih[int];[int][1:v]overlay=W/2:0[vid]' -map [vid] -c:v libx264 -crf 23 -preset veryfast ${filename}_both.mp4 -y 14 | -------------------------------------------------------------------------------- /oriented_convolve.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from pylab import * 3 | import numpy as np 4 | import scipy.ndimage 5 | import cv2 6 | import sys 7 | from board_detect import * 8 | 9 | def getCornerNormals(corners): 10 | # 4x2 array, rows are each point, columns are x and y 11 | dirs = np.zeros([4,2]) 12 | 13 | # Side lengths of rectangular contour 14 | dirs[0,:] = (corners[1,:] - corners[0,:]) / np.sqrt(np.sum((corners[1,:] - corners[0,:])**2)) 15 | dirs[1,:] = (corners[2,:] - corners[1,:]) / np.sqrt(np.sum((corners[2,:] - corners[1,:])**2)) 16 | dirs[2,:] = (corners[3,:] - corners[2,:]) / np.sqrt(np.sum((corners[3,:] - corners[2,:])**2)) 17 | dirs[3,:] = (corners[0,:] - corners[3,:]) / np.sqrt(np.sum((corners[0,:] - corners[3,:])**2)) 18 | 19 | # Rotate 90 deg to get normal vector via [-y,x] 20 | normals = -np.hstack([-dirs[:,[1]], dirs[:,[0]]]) 21 | # normals = dirs 22 | return normals, dirs 23 | 24 | 25 | def main(filenames): 26 | for filename in filenames: 27 | print("Processing %s" % filename) 28 | img = cv2.imread(filename) 29 | gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY) 30 | 31 | contours, chosen_tile_idx, edges = findPotentialTiles(img) 32 | if not len(contours): 33 | return 34 | drawPotentialTiles(img, contours, chosen_tile_idx) 35 | 36 | tile_corners = getChosenTile(contours, chosen_tile_idx) 37 | 38 | hough_corners, corner_hough_lines, edges_roi = refineTile(img, edges, contours, chosen_tile_idx) 39 | drawBestHoughLines(img, hough_corners, corner_hough_lines) 40 | 41 | corner_normals, _ = getCornerNormals(hough_corners) 42 | 43 | # print(hough_corners) 44 | # print(corner_normals) 45 | 46 | # Corner kernel 47 | # kX = np.array([[3,10,3],[0,0,0],[-3,-10,-3]]) 48 | # kY = np.array([[3,10,3],[0,0,0],[-3,-10,-3]]).T 49 | 50 | # Gradient of the combination of scharr X and Y operators 51 | # kX = np.array([[9,9,9],[0,0,0],[-9,-9,-9]]) 52 | # kY = np.array([[9,9,9],[0,0,0],[-9,-9,-9]]).T 53 | 54 | # kernel = kX * corner_normals[0,0] + kY * corner_normals[0,1] 55 | # print(kernel) 56 | 57 | # responseA = scipy.ndimage.filters.convolve(edges.astype(float), kernel, mode='constant') 58 | 59 | # Gradients 60 | sobelx = cv2.Sobel(gray,cv2.CV_64F,1,0,ksize=3) 61 | sobely = cv2.Sobel(gray,cv2.CV_64F,0,1,ksize=3) 62 | grad_mag = np.sqrt(sobelx**2+sobely**2) 63 | 64 | responseA = sobelx*corner_normals[0,0] + sobely*corner_normals[0,1] 65 | responseB = sobelx*corner_normals[1,0] + sobely*corner_normals[1,1] 66 | 67 | 68 | 69 | # # Normalize by absolute 70 | responseA = abs(responseA) 71 | responseB = abs(responseB) 72 | # # responseA /= responseA.max() 73 | 74 | # Normalize Response to 0-1 range 75 | # responseA = -responseA 76 | a,b = responseA.min(), responseA.max() 77 | responseA = ((responseA-a)/(b-a)) 78 | 79 | a,b = responseB.min(), responseB.max() 80 | responseB = ((responseB-a)/(b-a)) 81 | 82 | # responseA[grad_mag 1: 110 | filenames = sys.argv[1:] 111 | else: 112 | filenames = ['input2/27.jpg'] 113 | main(filenames) 114 | -------------------------------------------------------------------------------- /outlier_point_removal.py: -------------------------------------------------------------------------------- 1 | # Given a set of ML pruned saddle points, remove outliers or keep only chessboard area 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | 5 | 6 | pts = np.array([[237, 332], 7 | [242, 287], 8 | [245, 263], 9 | [247, 360], 10 | [249, 337], 11 | [252, 314], 12 | [255, 290], 13 | [257, 389], 14 | [258, 266], 15 | [260, 366], 16 | [262, 342], 17 | [268, 294], 18 | [271, 269], 19 | [274, 372], 20 | [280, 323], 21 | [286, 272], 22 | [286, 404], 23 | [291, 353], 24 | [295, 328], 25 | [298, 302], 26 | [298, 437], 27 | [301, 411], 28 | [304, 386], 29 | [311, 333], 30 | [318, 279], 31 | [318, 420], 32 | [321, 393], 33 | [324, 366], 34 | [329, 339], 35 | [332, 311], 36 | [336, 283]]) 37 | 38 | outliers = np.array([ 39 | [50,23], 40 | [30,63], 41 | [400,203], 42 | [250,370], 43 | [500,303], 44 | ]) 45 | 46 | pts = np.vstack([pts, outliers]) 47 | 48 | 49 | # plt.plot(outliers[:,0],outliers[:,1],'rx') 50 | # plt.plot(pts[:,0],pts[:,1],'.') 51 | # plt.show() 52 | N = len(pts) 53 | 54 | new_order = np.arange(N) 55 | np.random.shuffle(new_order) 56 | pts = pts[new_order,:] 57 | 58 | # TODO Calculate closest N points over threshold instead of just closest 59 | # because currently a pair of outliers next to each other are kept 60 | # def calculateOutliers(pts, threshold_mult = 3): 61 | # N = len(pts) 62 | # dists = np.zeros([N,N]) 63 | # best_dists = np.zeros(N) 64 | # for i in range(N): 65 | # dists[i,:] = np.linalg.norm(pts[:,:] - pts[i,:], axis=1) 66 | # x = np.linalg.norm(pts - pts[i,:], axis=1) 67 | # best_dists[i] = np.min(x[x!=0]) 68 | # print(best_dists) 69 | # # med = np.median(dists[dists!=0]) 70 | # med = np.median(best_dists) 71 | # print(med) 72 | # # return dists.min(axis=0) > med*threshold_mult 73 | # return best_dists > med*threshold_mult 74 | 75 | def calculateOutliers(pts, threshold_mult = 1.5): 76 | N = len(pts) 77 | std = np.std(pts, axis=0) 78 | ctr = np.mean(pts, axis=0) 79 | return (np.any(np.abs(pts-ctr) > threshold_mult * std, axis=1)) 80 | 81 | 82 | 83 | import time 84 | 85 | ta = time.time() 86 | 87 | outlier = calculateOutliers(pts) 88 | pred_outliers = pts[outlier,:] 89 | 90 | tb = time.time() 91 | 92 | proc_time = tb-ta 93 | print("Processed in %.2f ms" % (proc_time*1e3)) 94 | 95 | 96 | # plt.plot(outliers[:,0],outliers[:,1],'rx') 97 | plt.plot(pred_outliers[:,0],pred_outliers[:,1],'rx') 98 | plt.plot(pts[:,0],pts[:,1],'.') 99 | plt.show() -------------------------------------------------------------------------------- /output_ml.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/output_ml.gif -------------------------------------------------------------------------------- /processChessPoints.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import PIL.Image 3 | import matplotlib.image as mpimg 4 | import scipy.ndimage 5 | import cv2 # For Sobel etc 6 | import glob 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | import os 10 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 11 | 12 | def process(cpts, pts): 13 | # Given chess points (cpts) and all saddle points (pts) 14 | # Find the closest saddle point for each chess point and return the index 15 | closest_idxs = np.zeros(cpts.shape[0], dtype=np.int) 16 | best_dists = np.zeros(cpts.shape[0]) 17 | for i,cpt in enumerate(cpts): 18 | dists = np.sum((pts - cpt)**2, axis=1) 19 | closest_idxs[i] = np.argmin(dists) 20 | best_dists[i] = dists[closest_idxs[i]] 21 | return closest_idxs, best_dists 22 | 23 | def loadImage(filepath): 24 | img_orig = PIL.Image.open(filepath) 25 | img_width, img_height = img_orig.size 26 | 27 | # Resize 28 | aspect_ratio = min(500.0/img_width, 500.0/img_height) 29 | new_width, new_height = ((np.array(img_orig.size) * aspect_ratio)).astype(int) 30 | img = img_orig.resize((new_width,new_height), resample=PIL.Image.BILINEAR) 31 | # img = img.convert('L') # grayscale 32 | img = np.array(img) 33 | 34 | return img 35 | 36 | def makeProcessedImage(filename, chess_pts, all_pts, closest_idxs, best_dists): 37 | filename_img = 'input/%s.png' % filename[filename.find('/')+1:filename.find('.')] 38 | if not os.path.exists(filename_img): 39 | filename_img = 'input/%s.jpg' % filename[filename.find('/')+1:filename.find('.')] 40 | if not os.path.exists(filename_img): 41 | filename_img = 'input_yt/%s.jpg' % filename[filename.find('/')+1:filename.find('.')] 42 | if not os.path.exists(filename_img): 43 | filename_img = 'input_yt/%s.png' % filename[filename.find('/')+1:filename.find('.')] 44 | img = loadImage(filename_img) 45 | 46 | # Only show good updated saddle points 47 | for i, pt in enumerate(chess_pts): 48 | if (best_dists[i] <= 2): 49 | spt = tuple(all_pts[closest_idxs[i]].astype(np.int)) 50 | cv2.circle(img, spt, 2, (0,255,0), -1) 51 | for i, pt in enumerate(all_pts): 52 | if (i not in closest_idxs): 53 | cv2.circle(img, tuple(pt.astype(np.int)), 1, (255,0,0), -1) 54 | 55 | # Visualize all 56 | # for pt in all_pts: 57 | # cv2.circle(img, tuple(pt.astype(np.int)), 1, (255,0,0), -1) 58 | # for i, pt in enumerate(chess_pts): 59 | # if (best_dists[i] > 2): 60 | # cv2.circle(img, tuple(pt.astype(np.int)), 2, (0,100,0), -1) 61 | # else: 62 | # cv2.circle(img, tuple(pt.astype(np.int)), 2, (0,255,0), -1) 63 | 64 | # for i in range(len(chess_pts)): 65 | # pt_a = tuple(chess_pts[i].astype(np.int)) 66 | # pt_b = tuple(all_pts[closest_idxs[i]].astype(np.int)) 67 | # if (best_dists[i] > 2): 68 | # cv2.putText(img,'%.1f' % best_dists[i] ,pt_a, font, 0.5,(255,255,255),0,cv2.LINE_AA) 69 | # cv2.line(img, pt_a, pt_b, (255,0,255), 1) 70 | # else: 71 | # cv2.line(img, pt_a, pt_b, (0,0,255), 1) 72 | 73 | im = PIL.Image.fromarray(img).convert('RGB') 74 | processed_img_filename = filename[:filename.find('.')] 75 | im.save('%s_proc.png' % processed_img_filename) 76 | 77 | def main(): 78 | font = cv2.FONT_HERSHEY_PLAIN 79 | # all_pts 80 | filenames_chesspts = glob.glob('positions/*[!_all].txt') 81 | filenames_chesspts = sorted(filenames_chesspts) 82 | n_all = len(filenames_chesspts) 83 | 84 | to_skip = [5,7,16,27,28,36,37,38] 85 | 86 | all_good_pts = {} 87 | all_bad_pts = {} 88 | 89 | for i in range(n_all): 90 | filename = filenames_chesspts[i] 91 | filename_short = filename[filename.find('/')+1:filename.find('.')] 92 | if any('img_%02d' % skip_name in filename for skip_name in to_skip): 93 | print('Skipping %s' % filename) 94 | continue 95 | print ("Processing %d/%d : %s" % (i+1,n_all,filename)) 96 | filename_allpts = filename[:filename.find('.')] + '_all.txt' 97 | 98 | # Load chess points 99 | chess_pts = np.loadtxt(filename) 100 | 101 | # Load all saddle points 102 | all_pts = np.loadtxt(filename_allpts) 103 | # all_pts = all_pts[:,[1,0]] 104 | closest_idxs, best_dists = process(chess_pts, all_pts) 105 | # print(best_dists) 106 | 107 | makeProcessedImage(filename, chess_pts, all_pts, closest_idxs, best_dists) 108 | 109 | 110 | # good_pts = chess_pts[best_dists <= 2, :] 111 | good_pts = all_pts[closest_idxs[best_dists <= 2],:] 112 | bad_pts = all_pts.copy() 113 | bad_pts = np.delete(bad_pts, closest_idxs, 0) 114 | # print(len(good_pts), len(bad_pts)) 115 | all_good_pts[filename_short] = good_pts.astype(int) 116 | all_bad_pts[filename_short] = bad_pts.astype(int) 117 | print(len(good_pts), len(bad_pts)) 118 | 119 | 120 | # save all points to a file 121 | with open('pt_dataset2.txt', 'w') as f: 122 | for filename in sorted(all_good_pts.keys()): 123 | s0 = ' '.join(all_good_pts[filename][:,0].astype(str)) 124 | s1 = ' '.join(all_good_pts[filename][:,1].astype(str)) 125 | s2 = ' '.join(all_bad_pts[filename][:,0].astype(str)) 126 | s3 = ' '.join(all_bad_pts[filename][:,1].astype(str)) 127 | f.write('\n'.join([filename, s0, s1, s2, s3])+'\n') 128 | 129 | num_good = 0 130 | num_bad = 0 131 | for filename in all_good_pts: 132 | num_good += all_good_pts[filename].shape[0] 133 | num_bad += all_bad_pts[filename].shape[0] 134 | print('Collected %d true and %d false positives' % (num_good, num_bad)) 135 | if __name__ == '__main__': 136 | main() 137 | 138 | 139 | 140 | -------------------------------------------------------------------------------- /quickCheck_deleteeasily.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | np_dataset = np.load('dataset_5.npz') 4 | features = np_dataset['features'] 5 | labels = np_dataset['labels'] 6 | 7 | print(len(labels)) 8 | print(sum(labels)) -------------------------------------------------------------------------------- /quickFix.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import PIL.Image 3 | import matplotlib.image as mpimg 4 | import scipy.ndimage 5 | import cv2 # For Sobel etc 6 | import glob 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | import os 10 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 11 | 12 | # Reverse x and y in positions/*_all.txt once 13 | 14 | def main(): 15 | # all_pts 16 | filenames = glob.glob('positions/img_*_all.txt') 17 | filenames = sorted(filenames) 18 | n = len(filenames) 19 | 20 | for i in range(n): 21 | filename = filenames[i] 22 | print ("Processing %d/%d : %s" % (i+1,n,filename)) 23 | 24 | all_pts = np.loadtxt(filename) 25 | all_pts = all_pts[:,[1,0]] 26 | np.savetxt(filename, all_pts) 27 | 28 | if __name__ == '__main__': 29 | main() 30 | 31 | 32 | 33 | -------------------------------------------------------------------------------- /readme_find_warp_example.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/readme_find_warp_example.png -------------------------------------------------------------------------------- /readme_labeled.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/readme_labeled.png -------------------------------------------------------------------------------- /readme_output.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/readme_output.png -------------------------------------------------------------------------------- /readme_rectified.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/readme_rectified.gif -------------------------------------------------------------------------------- /rectify_refine.py: -------------------------------------------------------------------------------- 1 | import PIL.Image 2 | import matplotlib.pyplot as plt 3 | import cv2 4 | import numpy as np 5 | import itertools 6 | import os 7 | np.set_printoptions(precision=2, linewidth=400, suppress=True) # Better printing of arrays 8 | 9 | def nonMaxSupress2D(img, win=10): 10 | out_img = img.copy() 11 | w, h = img.shape 12 | for i in range(w): 13 | for j in range(h): 14 | # Skip if already suppressed 15 | # if out_img[i,j] == 0: 16 | # continue 17 | # Get neigborhood 18 | ta=max(0,i-win) 19 | tb=min(w,i+win+1) 20 | tc=max(0,j-win) 21 | td=min(h,j+win+1) 22 | cell = img[ta:tb,tc:td] 23 | # Blank out all non-max values in window 24 | out_img[ta:tb,tc:td] = (cell == cell.max()) * out_img[ta:tb,tc:td] 25 | return out_img 26 | 27 | def reRectifyImages(color_img, tile_res=64, tile_buffer=1): 28 | """Recenter off-by-N-tiles chessboards, then use 29 | corner subpixel with ransac findHomography to further optimize rectification of image""" 30 | if type(color_img).__module__ == np.__name__: 31 | img = np.array(PIL.Image.fromarray(color_img).convert('L')) # grayscale uint8 numpy array 32 | else: 33 | img = np.array(color_img.convert('L')) # grayscale uint8 numpy array 34 | # img = cv2.equalizeHist(img) # Global histogram equalization 35 | clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) 36 | img = clahe.apply(img) 37 | # Make sure our tile resolutions and buffer expectations match image size 38 | assert( (img.shape[0] / tile_res) == (8+2*tile_buffer) ) 39 | 40 | k = tile_res 41 | side_len = tile_res*(8+2*tile_buffer) 42 | quad = np.ones([k,k]) 43 | # kernel = np.vstack([np.hstack([quad,-quad]), np.hstack([-quad,quad])]) 44 | kernel = np.vstack([np.hstack([quad,-quad]), np.hstack([-quad,quad])]) 45 | kernel = np.tile(kernel,(4,4)) # Becomes 8x8 alternating grid 46 | kernel = kernel/np.linalg.norm(kernel) 47 | 48 | response = cv2.filter2D(img, cv2.CV_32F, kernel) 49 | corners = abs(response) 50 | expected_ctr_pt = np.array([tile_res*(4+tile_buffer), tile_res*(4+tile_buffer)]) 51 | best_pt = np.array(np.unravel_index(corners.argmax(), corners.shape))[::-1] # Flipped due to image y/x 52 | center_offset = best_pt - expected_ctr_pt 53 | 54 | # print(best_pt) 55 | should_rotate = response[best_pt[1],best_pt[0]] < 0 56 | # print(should_rotate) 57 | 58 | hlines = best_pt[0] + (np.arange(9)-4)*tile_res 59 | vlines = best_pt[1] + (np.arange(9)-4)*tile_res 60 | 61 | all_corners = np.array(list(itertools.product(hlines, vlines))) 62 | criteria = (cv2.TERM_CRITERIA_MAX_ITER + cv2.TERM_CRITERIA_COUNT, 30, 0.01) 63 | better_corners = cv2.cornerSubPix(img, all_corners.astype(np.float32), 64 | (20,20), (-1,-1), criteria) 65 | 66 | M, good_pts = cv2.findHomography(better_corners.astype(np.float32) + center_offset, all_corners.astype(np.float32), cv2.RANSAC) 67 | 68 | # should_rotate = checkChessboardAlignment(img) 69 | 70 | color_img_arr = np.array(color_img) 71 | 72 | if (should_rotate): 73 | # color_img_arr = np.rot90(color_img_arr) 74 | # rot90 = cv2.getRotationMatrix2D((side_len/2,side_len/2),90,1) 75 | rot90 = np.matrix([[0, 1, 0], [-1, 0, side_len], [0, 0, 1]]) 76 | M = np.matmul(rot90, M) 77 | 78 | return cv2.warpPerspective(color_img_arr, M, color_img_arr.shape[:2]), should_rotate, M 79 | 80 | # def checkChessboardAlignment(img, tile_res=64, tile_buffer=1): 81 | # Guess if chessboard black/white corners are correct, rotate 90 deg otherwise 82 | 83 | if __name__ == '__main__': 84 | PLOT_RESULTS = True 85 | 86 | input_folder = "rectified" 87 | 88 | tile_res = 64 89 | tile_buffer = 1 90 | 91 | filename ="%d.png" % 7 92 | filepath = "%s/%s" % (input_folder,filename) 93 | print("Refining %s" % filename) 94 | img_orig = PIL.Image.open(filepath) 95 | 96 | # Grayscale 97 | better_img, was_rotated, refine_M = reRectifyImages(img_orig) 98 | # print(refine_M) 99 | if was_rotated: 100 | print("Was rotated") 101 | 102 | if PLOT_RESULTS: 103 | hlines = vlines = (np.arange(9)+tile_buffer)*tile_res 104 | ideal_corners = np.array(list(itertools.product(hlines, vlines))) 105 | 106 | plt.subplot(121) 107 | plt.imshow(img_orig) 108 | plt.plot(ideal_corners[:,0], ideal_corners[:,1], 'ro', ms=3) 109 | plt.title('Input rectified image') 110 | 111 | plt.subplot(122) 112 | plt.imshow(better_img) 113 | plt.plot(ideal_corners[:,0], ideal_corners[:,1], 'ro', ms=3) 114 | plt.title('Refined output rectified image') 115 | 116 | plt.show() -------------------------------------------------------------------------------- /result.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/result.png -------------------------------------------------------------------------------- /run_xcorner_model_on_img.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | import PIL.Image 4 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 5 | 6 | 7 | featureA = tf.feature_column.numeric_column("x", shape=[11,11], dtype=tf.uint8) 8 | 9 | estimator = tf.estimator.DNNClassifier( 10 | feature_columns=[featureA], 11 | hidden_units=[256, 32], 12 | n_classes=2, 13 | dropout=0.1, 14 | model_dir='./xcorner_model2', 15 | ) 16 | 17 | 18 | # filepath = 'input/img_01.jpg' 19 | filepath = 'testB.jpg' 20 | img = PIL.Image.open(filepath).convert('L') # Grayscale full image with chessboard 21 | img_width, img_height = img.size 22 | aspect_ratio = min(200.0/img_width, 200.0/img_height) 23 | new_width, new_height = ((np.array(img.size) * aspect_ratio)).astype(int) 24 | img = img.resize((new_width,new_height), resample=PIL.Image.BILINEAR) 25 | img = np.array(img) 26 | 27 | winsize = 5 28 | 29 | new_size = [img.shape[0]-2*winsize, img.shape[1]-2*winsize] 30 | print("out_size : " , new_size) 31 | 32 | tiles = [] 33 | 34 | for r in range(winsize+1,img.shape[0]-winsize): 35 | for c in range(winsize+1,img.shape[1]-winsize): 36 | win = img[r-winsize:r+winsize+1,c-winsize:c+winsize+1] 37 | tiles.append(win) 38 | 39 | img_features = np.array(tiles) 40 | print(img_features.shape) 41 | 42 | # filepath = 'dataset_gray_5/bad/img_01_002.png' 43 | # img = PIL.Image.open(filepath) 44 | # img_features = np.array([np.array(img)]) 45 | 46 | def input_fn_predict(): # returns x, None 47 | dataset = tf.data.Dataset.from_tensor_slices( 48 | { 49 | 'x':img_features 50 | } 51 | ) 52 | # return dataset.make_one_shot_iterator().get_next(), tf.one_hot(labels,2,dtype=tf.int32) 53 | dataset = dataset.batch(25) 54 | iterator = dataset.make_one_shot_iterator() 55 | k = iterator.get_next() 56 | return k, None 57 | 58 | # Do prediction 59 | predictions = estimator.predict(input_fn=input_fn_predict) 60 | 61 | heatmap = np.zeros(img_features.shape[0]) 62 | for i,prediction in enumerate(predictions): 63 | if (i % 1000 == 0): 64 | print ('%d / %d ' % (i, len(img_features))) 65 | if (prediction['probabilities'].argmax() == 1): 66 | heatmap[i] = prediction['probabilities'][1] # Probability is saddle 67 | # heatmap[i] = prediction['probabilities'].argmax() # Final answer 68 | 69 | np.save('heatmap', heatmap) 70 | print(heatmap) -------------------------------------------------------------------------------- /run_xcorner_on_saddle_tiles.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import PIL.Image 3 | import matplotlib.image as mpimg 4 | import scipy.ndimage 5 | import cv2 # For Sobel etc 6 | import glob 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | import os 10 | import tensorflow as tf 11 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 12 | 13 | 14 | featureA = tf.feature_column.numeric_column("x", shape=[11,11], dtype=tf.uint8) 15 | 16 | estimator = tf.estimator.DNNClassifier( 17 | feature_columns=[featureA], 18 | hidden_units=[256, 32], 19 | n_classes=2, 20 | dropout=0.1, 21 | model_dir='./xcorner_model_7k', 22 | ) 23 | 24 | 25 | # Load pt_dataset.txt and generate the windowed tiles for all the good and bad 26 | # points in folders dataset/good dataset/bad 27 | 28 | 29 | def loadImage(filepath, doGrayscale=False): 30 | img_orig = PIL.Image.open(filepath) 31 | img_width, img_height = img_orig.size 32 | 33 | # Resize 34 | aspect_ratio = min(500.0/img_width, 500.0/img_height) 35 | new_width, new_height = ((np.array(img_orig.size) * aspect_ratio)).astype(int) 36 | img = img_orig.resize((new_width,new_height), resample=PIL.Image.BILINEAR) 37 | # if (doGrayscale): 38 | img_gray = img.convert('L') # grayscale 39 | img = np.array(img) 40 | img_gray = np.array(img_gray) 41 | 42 | return img, img_gray 43 | 44 | import errno 45 | def mkdir_p(path): 46 | try: 47 | os.makedirs(path) 48 | except OSError as exc: # Python >2.5 49 | if os.path.isdir(path): 50 | pass 51 | else: 52 | raise 53 | 54 | def main(): 55 | input_data = 'pt_dataset.txt' 56 | 57 | results_folder = 'ml_xcorner_results' 58 | mkdir_p(results_folder) 59 | 60 | WINSIZE = 5 61 | 62 | DO_BINARIZATION = False 63 | DO_OPENING = False 64 | 65 | if (DO_BINARIZATION): 66 | raise('Error, must be grayscale if doing binarization.') 67 | 68 | count_good = 0 69 | count_bad = 0 70 | 71 | # save all points to a file 72 | with open('pt_dataset.txt', 'r') as f: 73 | lines = [x.strip() for x in f.readlines()] 74 | 75 | n = len(lines)/5 76 | # n = 1 77 | 78 | def input_fn_predict(img_data): # returns x, None 79 | def ret_func(): 80 | dataset = tf.data.Dataset.from_tensor_slices( 81 | { 82 | 'x':img_data 83 | } 84 | ) 85 | # return dataset.make_one_shot_iterator().get_next(), tf.one_hot(labels,2,dtype=tf.int32) 86 | dataset = dataset.batch(25) 87 | iterator = dataset.make_one_shot_iterator() 88 | k = iterator.get_next() 89 | return k, None 90 | return ret_func 91 | 92 | 93 | for i in range(n): 94 | print("On %d/%d" % (i+1, n)) 95 | filename = lines[i*5] 96 | s0 = lines[i*5+1].split() 97 | s1 = lines[i*5+2].split() 98 | s2 = lines[i*5+3].split() 99 | s3 = lines[i*5+4].split() 100 | good_pts = np.array([s1, s0], dtype=np.int).T 101 | bad_pts = np.array([s3, s2], dtype=np.int).T 102 | 103 | img_filepath = 'input/%s.png' % filename 104 | if not os.path.exists(img_filepath): 105 | img_filepath = 'input/%s.jpg' % filename 106 | img, img_gray = loadImage(img_filepath) 107 | 108 | kernel = np.ones((3,3),np.uint8) 109 | 110 | tiles = [] 111 | all_pts = [] 112 | 113 | # Good points 114 | for i in range(good_pts.shape[0]): 115 | pt = good_pts[i,:] 116 | if (np.any(pt <= WINSIZE) or np.any(pt >= np.array(img_gray.shape[:2]) - WINSIZE)): 117 | # print("Skipping point %s" % pt) 118 | continue 119 | else: 120 | tile = img_gray[pt[0]-WINSIZE:pt[0]+WINSIZE+1, pt[1]-WINSIZE:pt[1]+WINSIZE+1] 121 | if DO_BINARIZATION: 122 | tile = cv2.adaptiveThreshold(tile,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2) 123 | 124 | if DO_OPENING: 125 | tile = cv2.morphologyEx(tile, cv2.MORPH_OPEN, kernel) 126 | 127 | tiles.append(tile) 128 | all_pts.append(pt) 129 | count_good += 1 130 | 131 | # Bad points 132 | for i in range(bad_pts.shape[0]): 133 | pt = bad_pts[i,:] 134 | if (np.any(pt <= WINSIZE) or np.any(pt >= np.array(img_gray.shape[:2]) - WINSIZE)): 135 | continue 136 | else: 137 | tile = img_gray[pt[0]-WINSIZE:pt[0]+WINSIZE+1, pt[1]-WINSIZE:pt[1]+WINSIZE+1] 138 | if DO_BINARIZATION: 139 | tile = cv2.adaptiveThreshold(tile,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2) 140 | if DO_OPENING: 141 | tile = cv2.morphologyEx(tile, cv2.MORPH_OPEN, kernel) 142 | 143 | 144 | tiles.append(tile) 145 | all_pts.append(pt) 146 | count_bad += 1 147 | 148 | tiles = np.array(tiles) 149 | all_pts = np.array(all_pts) 150 | 151 | # Do prediction 152 | import time 153 | a = time.time() 154 | predictions = estimator.predict(input_fn=input_fn_predict(tiles)) 155 | 156 | for i, prediction in enumerate(predictions): 157 | c = prediction['probabilities'].argmax() 158 | pt = all_pts[i] 159 | if (c == 1): 160 | cv2.circle(img, tuple(pt[::-1]), 4, (0,255,0), -1) 161 | else: 162 | cv2.circle(img, tuple(pt[::-1]), 3, (255,0,0), -1) 163 | b = time.time() 164 | print("predict() took %.2f ms" % ((b-a)*1e3)) 165 | 166 | 167 | im = PIL.Image.fromarray(img).convert('RGB') 168 | im.save('%s/%s_xcorner_7k.png' % (results_folder, filename)) 169 | 170 | print ("Finished %d good and %d bad tiles" % (count_good, count_bad)) 171 | 172 | 173 | 174 | if __name__ == '__main__': 175 | main() 176 | 177 | 178 | 179 | -------------------------------------------------------------------------------- /saddle/0.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/saddle/0.jpg -------------------------------------------------------------------------------- /saddle/1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/saddle/1.jpg -------------------------------------------------------------------------------- /saddle/2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/saddle/2.jpg -------------------------------------------------------------------------------- /sam2_composite.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/sam2_composite.gif -------------------------------------------------------------------------------- /speedchess1_composite.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/speedchess1_composite.gif -------------------------------------------------------------------------------- /speedchess1_ml.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/speedchess1_ml.gif -------------------------------------------------------------------------------- /tile_segment.py: -------------------------------------------------------------------------------- 1 | # Image segmentation 2 | # Given rectified image with known tile boundaries 3 | # Segment image into background (black/white tiles?) 4 | # and dark or light pieces 5 | # 6 | # Some options include K-means clustering, watershed segmentation, texture segmentation, perhaps a combination 7 | 8 | import PIL.Image 9 | import matplotlib.pyplot as plt 10 | import cv2 11 | import numpy as np 12 | import itertools 13 | import os 14 | from skimage import color 15 | from sklearn.cluster import KMeans 16 | from skimage import exposure 17 | np.set_printoptions(precision=2, linewidth=400, suppress=True) # Better printing of arrays 18 | 19 | def getIdealCorners(tile_res, tile_buffer): 20 | hlines = vlines = (np.arange(9)+tile_buffer)*tile_res 21 | return np.array(list(itertools.product(hlines, vlines))) 22 | 23 | def getIdealCheckerboardPattern(tile_res, tile_buffer): 24 | side_len = tile_res*(8+2*tile_buffer) 25 | quadOne = np.ones([tile_res,tile_res], dtype=np.uint8) 26 | quadZero = np.zeros([tile_res,tile_res], dtype=np.uint8) 27 | kernel = np.vstack([np.hstack([quadOne,quadZero]), np.hstack([quadZero,quadOne])]) 28 | kernel = np.tile(kernel,(4,4)) # Becomes 8x8 alternating grid 29 | return kernel 30 | 31 | def getTile(img, i,j,tile_res): 32 | """Assumes no buffer in image""" 33 | return img[tile_res*i:tile_res*(i+1),tile_res*j:tile_res*(j+1)] 34 | 35 | 36 | if __name__ == '__main__': 37 | PLOT_RESULTS = True 38 | 39 | input_folder = "rectified" 40 | 41 | tile_res = 64 42 | tile_buffer = 1 43 | 44 | side_len = 8*tile_res 45 | buffer_size = tile_buffer*tile_res 46 | 47 | filename ="%d.png" % 31 48 | filepath = "%s/%s" % (input_folder,filename) 49 | print("Segmenting %s..." % filename) 50 | img_orig = np.array(PIL.Image.open(filepath).convert('RGB')) 51 | 52 | img_h, img_w, _ = img_orig.shape 53 | 54 | # Bilateral smooth image 55 | img = img_orig 56 | bg_illum = cv2.blur(img,ksize=(tile_res*4,tile_res*4)) 57 | img = (bg_illum.mean() + (img.astype(np.float64) - bg_illum)).astype(np.uint8) 58 | # img = cv2.blur(img,ksize=(5,5)) 59 | # img = cv2.medianBlur(img,7) 60 | # img = cv2.bilateralFilter(img,int(tile_res/4),75,75) 61 | 62 | ideal_corners = getIdealCorners(tile_res, tile_buffer) 63 | 64 | img_checkerboard = img[buffer_size:-buffer_size, buffer_size:-buffer_size] 65 | img_checkerboard_before = img_checkerboard.copy() 66 | 67 | ycrcb = cv2.cvtColor(img_checkerboard, cv2.COLOR_RGB2YCR_CB) 68 | # ycrcb[:,:,0] = cv2.equalizeHist(ycrcb[:,:,0].astype(np.uint8)) 69 | clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(4,4)) 70 | ycrcb[:,:,0] = clahe.apply(ycrcb[:,:,0].astype(np.uint8)) 71 | img_checkerboard = cv2.cvtColor(ycrcb, cv2.COLOR_YCR_CB2RGB) 72 | img_checkerboard = cv2.medianBlur(img_checkerboard,7) 73 | 74 | img_checkerboard_gray = np.array(PIL.Image.fromarray(img_checkerboard).convert('L')) 75 | # img_checkerboard_gray = cv2.equalizeHist(img_checkerboard_gray) 76 | 77 | img_draw = np.zeros(img_checkerboard_gray.shape) 78 | 79 | # Watershed 80 | markers = np.zeros(img_checkerboard_gray.shape, dtype=np.int32) 81 | teb = int(64/2-2) 82 | N = 8 83 | marker_counter = 2 84 | # watershed_mask = np.zeros(markers.shape, dtype=bool) 85 | for i in range(N): 86 | # watershed_mask[tile_res*i,:] = False 87 | # watershed_mask[:,tile_res*i] = False 88 | for j in range(N): 89 | markers[tile_res*i+teb:tile_res*(i+1)-teb,tile_res*j+teb:tile_res*(j+1)-teb] = marker_counter 90 | marker_counter += 1 91 | 92 | markers_init = markers.copy() 93 | img_watershed = cv2.watershed(img_checkerboard, markers) 94 | # from skimage.morphology import watershed 95 | # img_watershed = watershed(img_checkerboard_gray, markers, mask=~watershed_mask) 96 | 97 | 98 | if PLOT_RESULTS: 99 | print("Plotting") 100 | plt.figure(filename, figsize=(20,8)) 101 | 102 | teb = 10 103 | N = 8 104 | for i in range(N): 105 | for j in range(N): 106 | k = i*N + j + 1 107 | # plt.subplot(N,N,k) 108 | 109 | tile = getTile(img_checkerboard,i,j,tile_res) 110 | tile_gray = getTile(img_checkerboard_gray,i,j,tile_res)[teb:-teb,teb:-teb] 111 | _,tile_thresh = cv2.threshold(tile_gray,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU) 112 | tile_draw = np.ma.masked_where(tile_thresh, tile_gray) 113 | if np.std(tile_gray) > 10: 114 | img_draw[tile_res*i+teb:tile_res*(i+1)-teb,tile_res*j+teb:tile_res*(j+1)-teb] = tile_thresh 115 | # plt.imshow(tile_draw) 116 | # plt.axis([0,tile_res,tile_res,0]) 117 | # plt.axis('off') 118 | plt.subplot(241) 119 | plt.imshow(img_orig) 120 | plt.subplot(242) 121 | plt.imshow(img_checkerboard_before) 122 | plt.title('Checkerboard Before') 123 | plt.subplot(243) 124 | plt.imshow(img_draw) 125 | plt.title('Thresholded') 126 | plt.subplot(244) 127 | plt.imshow(img_watershed) 128 | plt.title('Watershed') 129 | plt.subplot(245) 130 | plt.imshow(img_checkerboard_gray, cmap='Greys_r') 131 | plt.title('Gray') 132 | plt.subplot(246) 133 | plt.imshow(img_checkerboard) 134 | plt.title('Checkerboard After') 135 | 136 | plt.subplot(247) 137 | plt.imshow(markers_init) 138 | plt.title('Markers init') 139 | 140 | plt.subplot(248) 141 | plt.imshow(bg_illum) 142 | plt.title('Background Illumination') 143 | 144 | 145 | plt.show() -------------------------------------------------------------------------------- /trainML_model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 4 | # import tensorflow.contrib.eager as tfe 5 | # tfe.enable_eager_execution() 6 | 7 | with np.load('dataset2_5.npz') as np_dataset: 8 | # Get equal good/bad 9 | features = np_dataset['features']#[:(6227*2)] 10 | labels = np_dataset['labels']#[:(6227*2)] 11 | 12 | print(len(features), len(labels)) 13 | 14 | # Shuffle data so good/bad are mixed 15 | shuffle_order = np.arange(len(labels)) 16 | np.random.shuffle(shuffle_order) 17 | features = features[shuffle_order,:,:] 18 | labels = labels[shuffle_order] 19 | 20 | # Split into train and test datasets 21 | split = np.int(len(labels) * 0.8) 22 | train_features = features[:split] 23 | train_labels = labels[:split] 24 | 25 | test_features = features[split:] 26 | test_labels = labels[split:] 27 | 28 | 29 | # Assume that each row of `features` corresponds to the same row as `labels`. 30 | assert features.shape[0] == labels.shape[0] 31 | 32 | # dataset = tf.data.Dataset.from_tensor_slices({'features':features, 'labels':labels}).repeat() 33 | # Todo add placeholders to avoid repeat copies of data 34 | # dataset = tf.data.Dataset.from_tensor_slices((features, labels)) 35 | 36 | # iterator = tfe.Iterator(dataset) 37 | # for i,value in enumerate(iterator): 38 | # print(i) 39 | # print(value) 40 | # if (i > 10): 41 | # break; 42 | 43 | def input_fn(feats, labs, do_shuffle=True): 44 | def return_fn(): 45 | dataset = tf.data.Dataset.from_tensor_slices( 46 | { 47 | 'x':feats, 48 | 'label':labs 49 | } 50 | ) 51 | if (do_shuffle): 52 | dataset = dataset.shuffle(len(feats)*2) 53 | dataset = dataset.batch(25).repeat() 54 | 55 | k = dataset.make_one_shot_iterator().get_next() 56 | return k, k['label'] 57 | return return_fn 58 | 59 | featureA = tf.feature_column.numeric_column("x", shape=[11,11], dtype=tf.uint8) 60 | 61 | estimator = tf.estimator.DNNClassifier( 62 | feature_columns=[featureA], 63 | hidden_units=[256, 32], 64 | n_classes=2, 65 | dropout=0.1, 66 | model_dir='./xcorner_model_all', 67 | ) 68 | 69 | n = 20 70 | for i in range(n): 71 | print("Training %d-%d/%d" % (i*1000,(i+1)*1000,n*1000)) 72 | estimator.train(input_fn=input_fn(train_features, train_labels), steps=1000) 73 | 74 | print("Evaluation #%d" % (i+1)) 75 | metrics = estimator.evaluate(input_fn=input_fn(test_features, test_labels), steps=1000) 76 | 77 | def input_fn_predict(): # returns x, None 78 | dataset = tf.data.Dataset.from_tensor_slices( 79 | { 80 | 'x':test_features 81 | } 82 | ) 83 | # return dataset.make_one_shot_iterator().get_next(), tf.one_hot(labels,2,dtype=tf.int32) 84 | dataset = dataset.batch(25) 85 | iterator = dataset.make_one_shot_iterator() 86 | k = iterator.get_next() 87 | return k, None 88 | 89 | predictions = estimator.predict(input_fn=input_fn_predict) 90 | 91 | count_good = 0 92 | for i,(prediction,true_answer) in enumerate(zip(predictions, test_labels.astype(np.int))): 93 | if prediction['probabilities'].argmax() != true_answer: 94 | print("Failure on %d: %s" % (i, prediction['probabilities'])) 95 | else: 96 | count_good += 1 97 | 98 | success_rate = float(count_good) / len(test_labels) 99 | print("Total %d/%d right ~ %.2f%% success rate" % (count_good, len(test_labels), 100*success_rate)) -------------------------------------------------------------------------------- /training_pipeline/doTestSweep.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | steps=10 4 | 5 | for (( i = 0; i < 100; i++ )); do 6 | echo "On ${i}/100" 7 | # Run test with just original set 5k/5k of each 8 | # folders="../datasets/dataset_gray_10/" 9 | # python trainML_pipeline.py -s ${steps} --name orig_5k -m 5000 ${folders} 10 | 11 | # # Run test with 3 more tiles, 1k/1k of each 12 | # folders="../datasets/dataset_gray_10/ ../results/bro_1_vidstream_frames/tiles ../results/chess_beer_vidstream_frames/tiles ../results/gm_magnus_1_vidstream_frames/tiles" 13 | # python trainML_pipeline.py -s ${steps} --name orig_plus_3_1k -m 1000 ${folders} 14 | 15 | # # Run test with 3 more tiles, 6k/6k of each 16 | # folders="../datasets/dataset_gray_10/ ../results/bro_1_vidstream_frames/tiles ../results/chess_beer_vidstream_frames/tiles ../results/gm_magnus_1_vidstream_frames/tiles" 17 | # python trainML_pipeline.py -s ${steps} --name orig_plus_3_6k -m 6000 ${folders} 18 | 19 | # # Run test with all tiles, 100/100 of each 20 | # # folders="../datasets/dataset_gray_10/ ../results/bro_1_vidstream_frames/tiles ../results/chess_beer_vidstream_frames/tiles ../results/gm_magnus_1_vidstream_frames/tiles ../results/john1_vidstream_frames/tiles ../results/john2_vidstream_frames/tiles ../results/match2_vidstream_frames/tiles ../results/output2_vidstream_frames/tiles ../results/output_vidstream_frames/tiles ../results/sam1_vidstream_frames/tiles ../results/sam2_vidstream_frames/tiles ../results/speedchess1_vidstream_frames/tiles ../results/swivel_vidstream_frames/tiles ../results/wgm_1_vidstream_frames/tiles" 21 | # # python trainML_pipeline.py -s ${steps} --name all_100 -m 100 ${folders} 22 | 23 | # # # Run test with all tiles, 1k/1k of each 24 | # # python trainML_pipeline.py -s ${steps} --name all_1k -m 1000 ${folders} 25 | 26 | # # Run test with all tiles, 6k/6k of each 27 | # python trainML_pipeline.py -s ${steps} --name all_6k -m 6000 ${folders} 28 | 29 | # # Run test with all tiles, 15k/15k of each 30 | # python trainML_pipeline.py -s ${steps} --name all_15k -m 15000 ${folders} 31 | 32 | 33 | folders="../datasets/dataset_gray_10/" 34 | python trainML_CNN_pipeline.py -s ${steps} --name orig_5k -m 5000 ${folders} 35 | done -------------------------------------------------------------------------------- /training_pipeline/preprocess.py: -------------------------------------------------------------------------------- 1 | # Methods to load datasets from folders, preprocess them, 2 | # and build input functions for the estimators. 3 | import tensorflow as tf 4 | import glob 5 | import numpy as np 6 | from tensorflow.contrib.data import Dataset 7 | import random 8 | 9 | def loadMultipleDatapaths(parent_folder_list, max_count_each_entries=None, pre_shuffle=False, do_shuffle=True, make_equal=False): 10 | # Pass list of folder paths 11 | filepaths_good = [] 12 | filepaths_bad = [] 13 | for parent_folder in parent_folder_list: 14 | folder_filepaths_good = glob.glob("%s/good/*.png" % parent_folder) 15 | folder_filepaths_bad = glob.glob("%s/bad/*.png" % parent_folder) 16 | 17 | if pre_shuffle: 18 | # Shuffle individual folder paths 19 | random.shuffle(folder_filepaths_good) 20 | random.shuffle(folder_filepaths_bad) 21 | 22 | if max_count_each_entries: 23 | folder_filepaths_good = folder_filepaths_good[:max_count_each_entries] 24 | folder_filepaths_bad = folder_filepaths_bad[:max_count_each_entries] 25 | 26 | filepaths_good.extend(folder_filepaths_good) 27 | filepaths_bad.extend(folder_filepaths_bad) 28 | 29 | # Make count of good and bad equal 30 | if make_equal: 31 | n_each = min(len(filepaths_good), len(filepaths_bad)) 32 | filepaths_good = filepaths_good[:n_each] 33 | filepaths_bad = filepaths_bad[:n_each] 34 | 35 | N_good, N_bad = len(filepaths_good), len(filepaths_bad) 36 | 37 | # Set up labels 38 | labels = np.array([1]*N_good + [0]*N_bad, dtype=np.float64) 39 | 40 | # Shuffle all entries keeping labels and paths together. 41 | entries = zip(filepaths_good + filepaths_bad, labels) 42 | if do_shuffle: 43 | random.shuffle(entries) 44 | # Separate back into imgs / labels and return. 45 | imgs, labels = zip(*entries) 46 | return tf.constant(imgs), tf.constant(labels), len(labels), sum(labels) 47 | 48 | def buildBothDatasets(img_paths, labels, train_test_split_percentage=0.8): 49 | # Split into training and test 50 | split = int(len(img_paths) * train_test_split_percentage) 51 | tr_imgs = tf.constant(img_paths[:split]) 52 | tr_labels = tf.constant(labels[:split]) 53 | val_imgs = tf.constant(img_paths[split:]) 54 | val_labels = tf.constant(labels[split:]) 55 | 56 | return tr_imgs, tr_labels, val_imgs, val_labels 57 | 58 | def input_parser(img_path, label): 59 | # Read the img from file. 60 | img_file = tf.read_file(img_path) 61 | img_decoded = tf.image.decode_image(img_file, channels=1) 62 | 63 | return img_decoded, label 64 | 65 | def randomize_image(img, contrast_range=[0.2,1.8], brightness_max=0.5): 66 | # Apply random flips/rotations and contrast/brightness changes to image 67 | img = tf.image.random_flip_left_right(img) 68 | img = tf.image.random_flip_up_down(img) 69 | img = tf.image.random_contrast(img, lower=contrast_range[0], upper=contrast_range[1]) 70 | img = tf.image.random_brightness(img, max_delta=brightness_max) 71 | img = tf.contrib.image.rotate(img, tf.random_uniform([1], minval=-np.pi, maxval=np.pi)) 72 | return img 73 | 74 | def preprocessor(dataset, batch_size, dataset_length=None, is_training=False): 75 | if is_training and dataset_length: 76 | # Shuffle dataset. 77 | dataset = dataset.shuffle(dataset_length*2) 78 | 79 | # Load images from image paths. 80 | dataset = dataset.map(input_parser) 81 | 82 | if is_training: 83 | # Slightly randomize images. 84 | dataset = dataset.map(lambda img, label: (randomize_image(img), label)) 85 | 86 | # Zero mean and unit normalize images, float image output. 87 | # TODO : Check if this needs to be applied to the predict function also 88 | # TODO : Does this cancel out random_brightness? 89 | # dataset = dataset.map(lambda img, label: (tf.image.per_image_standardization(img), label)) 90 | 91 | # Bring down to 15x15 from 21x21 92 | dataset = dataset.map(lambda img, label: (tf.image.central_crop(img, 0.666666), label)) 93 | 94 | # Batch and repeat. 95 | dataset = dataset.batch(batch_size) 96 | if is_training: 97 | dataset = dataset.repeat() 98 | 99 | return dataset 100 | 101 | def input_fn(imgs, labels, dataset_length=None, is_training=False, batch_size=50): 102 | # Returns an appropriate input function for training/evaluation. 103 | def sub_input_fn(): 104 | dataset = Dataset.from_tensor_slices((imgs, labels)) 105 | # Pre-process dataset into correct form/batching/shuffle etc. 106 | dataset = preprocessor(dataset, batch_size, dataset_length, is_training) 107 | 108 | # Build iterator and return 109 | one_shot_iterator = dataset.make_one_shot_iterator() 110 | next_element = one_shot_iterator.get_next() 111 | 112 | # Return in a dict so the premade estimators can use it. 113 | return {"x": next_element[0]}, next_element[1] 114 | return sub_input_fn -------------------------------------------------------------------------------- /training_pipeline/saveModel.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import tensorflow as tf 3 | import os 4 | from model import cnn_model_fn 5 | from argparse import ArgumentParser 6 | 7 | def main(args): 8 | export_dir = args.export_dir 9 | if not export_dir: 10 | export_dir = 'exported_models/%s' % (args.model_dir[args.model_dir.rstrip('/').rfind('/')+1:]) 11 | 12 | print('Loading from %s ...' % args.model_dir) 13 | print('Exporting to %s ...' % export_dir) 14 | 15 | with tf.Session(graph=tf.Graph()) as sess: 16 | # Load CNN estimator from model_dir 17 | estimator = tf.estimator.Estimator(model_fn=cnn_model_fn, model_dir=args.model_dir, 18 | params={ 19 | # CNN filter layers 1 and 2, and then dense layer # of units 20 | 'filter_sizes':args.filter_sizes 21 | }) 22 | 23 | # Build input function. 24 | image = tf.placeholder(tf.uint8, [None, args.winsize, args.winsize, args.channels]) 25 | input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({ 26 | 'x': image, 27 | }) 28 | 29 | # Export model. 30 | estimator.export_savedmodel(export_dir, input_fn) 31 | 32 | if __name__ == '__main__': 33 | parser = ArgumentParser() 34 | parser.add_argument("--model_dir", dest="model_dir", required=True, 35 | help="Directory to load model from (Ex. 'ml/model/cnn1').") 36 | parser.add_argument("--export_dir", dest="export_dir", 37 | help="Directory to export model to, default uses model_dir name.") 38 | parser.add_argument("-fs", "--filter_sizes", dest="filter_sizes", nargs='+', type=int, 39 | default=[32, 64, 1024], help="CNN model filter sizes") 40 | parser.add_argument("-ws", "--winsize", dest="winsize", default=10, type=int, 41 | help="Half window size (full kernel = 2*winsize + 1)") 42 | parser.add_argument("-c", "--channels", dest="channels", default=3, type=int, 43 | help="# of channels") 44 | 45 | args = parser.parse_args() 46 | main(args) -------------------------------------------------------------------------------- /training_pipeline/testModelOnVideo.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | import sys 3 | import os 4 | sys.path.append(os.path.dirname(os.getcwd())) 5 | from argparse import ArgumentParser 6 | import cv2 7 | import PIL.Image 8 | import skvideo.io 9 | import numpy as np 10 | # import Brutesac 11 | import SaddlePoints 12 | from functools import wraps 13 | import time 14 | from scipy.spatial import ConvexHull 15 | 16 | def predictOnTiles(tiles, predict_fn): 17 | predictions = predict_fn( 18 | {"x": tiles}) 19 | 20 | # Return array of probability of tile being an xcorner. 21 | # return np.array([p[1] for p in predictions['probabilities']]) 22 | return np.array([p[1] for p in predictions['probabilities']]) 23 | 24 | def predictOnImage(pts, img, gx, gy, predict_fn, WINSIZE = 10): 25 | # Build tiles to run classifier on. (23 ms) 26 | tiles = getTilesFromImage(pts, img, WINSIZE=WINSIZE) 27 | # tiles = getTilesFromGradients(pts, gx, gy, WINSIZE=WINSIZE) 28 | 29 | # Classify tiles. (~137ms) 30 | probs = predictOnTiles(tiles, predict_fn) 31 | 32 | return probs 33 | 34 | def getTilesFromImage(pts, img, WINSIZE=10): 35 | # NOTE : Assumes no point is within WINSIZE of an edge! 36 | 37 | # Points Nx2, columns should be x and y, not r and c. 38 | # Build tiles 39 | img_shape = np.array([img.shape[1], img.shape[0]]) 40 | tiles = np.zeros([len(pts), WINSIZE*2+1, WINSIZE*2+1, 3]) 41 | for i, pt in enumerate(np.round(pts).astype(np.int64)): 42 | tiles[i,:,:,:] = img[ 43 | pt[1]-WINSIZE:pt[1]+WINSIZE+1, pt[0]-WINSIZE:pt[0]+WINSIZE+1, :] 44 | 45 | return tiles 46 | 47 | def getTilesFromGradients(pts, gx, gy, WINSIZE=10): 48 | # NOTE : Assumes no point is within WINSIZE of an edge! 49 | 50 | # Points Nx2, columns should be x and y, not r and c. 51 | # Build tiles 52 | tiles = np.zeros([len(pts), WINSIZE*2+1, WINSIZE*2+1, 2]) 53 | for i, pt in enumerate(np.round(pts).astype(np.int64)): 54 | tiles[i,:,:,0] = gx[ 55 | pt[1]-WINSIZE:pt[1]+WINSIZE+1, pt[0]-WINSIZE:pt[0]+WINSIZE+1] 56 | tiles[i,:,:,1] = gy[ 57 | pt[1]-WINSIZE:pt[1]+WINSIZE+1, pt[0]-WINSIZE:pt[0]+WINSIZE+1] 58 | 59 | return tiles 60 | 61 | def classifyFrame(frame, gray, predict_fn, WINSIZE=10): 62 | # All saddle points 63 | spts, gx, gy = SaddlePoints.getFinalSaddlePoints(gray, WINSIZE) 64 | 65 | # Saddle points classified as Chessboard X-corners 66 | probabilities = predictOnImage(spts, frame, gx, gy, predict_fn, WINSIZE=WINSIZE) 67 | 68 | return spts, probabilities 69 | 70 | def processFrame(frame, gray, predict_fn, probability_threshold=0.9,WINSIZE=10): 71 | overlay_frame = frame.copy() 72 | # Overlay good and bad points onto the frame 73 | spts, probabilities = classifyFrame(frame, gray, predict_fn, WINSIZE=WINSIZE) 74 | 75 | # 10ms for the rest of this 76 | 77 | # Threshold over 50% probability as xpts 78 | xpts = spts[probabilities > probability_threshold,:] 79 | not_xpts = spts[probabilities <= probability_threshold,:] 80 | 81 | # Draw xcorner points 82 | for pt in np.round(xpts).astype(np.int64): 83 | cv2.rectangle(overlay_frame, tuple(pt-2),tuple(pt+2), (0,255,0), -1) 84 | 85 | # Draw rejects 86 | for pt in np.round(not_xpts).astype(np.int64): 87 | cv2.rectangle(overlay_frame, tuple(pt-0),tuple(pt+0), (0,0,255), -1) 88 | 89 | return overlay_frame, spts, probabilities 90 | 91 | 92 | def videostream(predict_fn, filepath='carlsen_match.mp4', 93 | output_folder_prefix='', SAVE_FRAME=True, MAX_FRAME=None, 94 | DO_VISUALS=True, EVERY_N_FRAMES=1): 95 | print("Loading video %s" % filepath) 96 | 97 | # Load frame-by-frame 98 | vidstream = skvideo.io.vreader(filepath) 99 | filename = os.path.basename(filepath) 100 | 101 | output_folder = "%s/%s_vidstream_frames" % (output_folder_prefix, filename[:-4]) 102 | if SAVE_FRAME: 103 | if not os.path.exists(output_folder): 104 | os.mkdir(output_folder) 105 | 106 | # # Set up pts.txt, first line is the video filename 107 | # # Following lines is the frame number and the flattened M matrix for the chessboard 108 | # output_filepath_pts = '%s/xpts.txt' % (output_folder) 109 | # with open(output_filepath_pts, 'w') as f: 110 | # f.write('%s\n' % filepath) 111 | 112 | for i, frame in enumerate(vidstream): 113 | if i >= MAX_FRAME: 114 | print('Reached max frame %d >= %d' % (i, MAX_FRAME)) 115 | break 116 | print("Frame %d" % i) 117 | if (i%EVERY_N_FRAMES!=0): 118 | continue 119 | 120 | # # Resize to 960x720 121 | frame = cv2.resize(frame, (480, 360), interpolation = cv2.INTER_CUBIC) 122 | 123 | # Our operations on the frame come here 124 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) 125 | gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) 126 | 127 | 128 | a = time.time() 129 | overlay_frame, spts, probabilities = processFrame(frame, gray, predict_fn, WINSIZE=args.winsize) 130 | t_proc = time.time() - a 131 | 132 | # Add frame counter 133 | cv2.putText(overlay_frame, 'Frame % 4d (Processed in % 6.1f ms)' % (i, t_proc*1e3), (5,15), cv2.FONT_HERSHEY_PLAIN, 1.0,(255,255,255),0) 134 | 135 | if DO_VISUALS: 136 | # Display the resulting frame 137 | cv2.imshow('overlayFrame',overlay_frame) 138 | 139 | 140 | output_orig_filepath = '%s/frame_%03d.jpg' % (output_folder, i) 141 | output_filepath = '%s/ml_frame_%03d.jpg' % (output_folder, i) 142 | 143 | if SAVE_FRAME: 144 | cv2.imwrite(output_orig_filepath, frame) 145 | cv2.imwrite(output_filepath, overlay_frame) 146 | 147 | # Append line of frame index and chessboard_corners matrix 148 | # if chessboard_corners is not None: 149 | # with open(output_filepath_pts, 'a') as f: 150 | # chessboard_corners_str = ','.join(map(str,spts.flatten())) 151 | # # M_str = M.tostring() # binary 152 | # f.write(u'%d,%s\n' % (i, chessboard_corners_str)) 153 | 154 | if DO_VISUALS: 155 | if cv2.waitKey(1) & 0xFF == ord('q'): 156 | break 157 | 158 | # When everything done, release the capture 159 | if DO_VISUALS: 160 | cv2.destroyAllWindows() 161 | 162 | 163 | def main(): 164 | filename = 'chess_out1.png' 165 | 166 | print ("Processing %s" % (filename)) 167 | img = PIL.Image.open(filename).resize([600,400]) 168 | rgb = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) 169 | gray = np.array(img.convert('L')) 170 | 171 | ### 172 | overlay_frame, spts, probabilities = processFrame(rgb, gray) 173 | ### 174 | 175 | cv2.imshow('frame',overlay_frame) 176 | cv2.waitKey() 177 | 178 | print('Finished') 179 | 180 | def getModel(export_dir='ml/model/006/1528565066'): 181 | from tensorflow.contrib import predictor 182 | return predictor.from_saved_model(export_dir, signature_def_key='predict') 183 | 184 | if __name__ == '__main__': 185 | parser = ArgumentParser() 186 | parser.add_argument("--model", dest="model", default='', 187 | help="Path to exported model to use.") 188 | parser.add_argument("video_inputs", nargs='+', 189 | help="filepaths to videos to process") 190 | parser.add_argument("-ws", "--winsize", dest="winsize", default=10, type=int, 191 | help="Half window size (full kernel = 2*winsize + 1)") 192 | 193 | 194 | args = parser.parse_args() 195 | print("Arguments passed: \n\t%s\n" % args) 196 | # main() 197 | # filename = 'output2.avi' # Slow low rez 198 | # filename = 'random1.mp4' # Long video wait for 1k frames or so 199 | # filename = 'match2.mp4' # difficult 200 | # filename = 'output.avi' # Hard low rez 201 | # filename = 'output.mp4' # Hard 202 | # filename = 'speedchess1.mp4' # Great example 203 | # filename = 'wgm_1.mp4' # Lots of motion blur, slow 204 | # filename = 'gm_magnus_1.mp4' # Hard lots of scene transitions and blurry (init state with all pieces in a row not so good). 205 | # filename = 'bro_1.mp4' # Little movement, easy. 206 | # filename = 'chess_beer.mp4' # Reasonably easy, some off-by-N errors 207 | # filename = 'john1.mp4' # Simple clean 208 | # filename = 'john2.mp4' # Slight motion, clean but slow 209 | # filename = 'swivel.mp4' # Moving around a fancy gold board 210 | 211 | allfiles = ['chess_beer.mp4', 'random1.mp4', 'match2.mp4','output.avi','output.mp4', 212 | 'speedchess1.mp4','wgm_1.mp4','gm_magnus_1.mp4', 213 | 'bro_1.mp4','output2.avi','john1.mp4','john2.mp4','swivel.mp4', 'sam2.mp4'] 214 | 215 | # for filename in allfiles: 216 | # for filename in ['match2.mp4']: 217 | for fullpath in args.video_inputs: 218 | # fullpath = 'datasets/raw/videos/%s' % filename 219 | output_folder_prefix = '../results' 220 | print('\n\n - ON %s\n\n' % fullpath) 221 | # predict_fn = RunExportedMLOnImage.getModel() 222 | # predict_fn = RunExportedMLOnImage.getModel('ml/model/run97pct/1528942225') 223 | predict_fn = getModel(args.model) 224 | videostream(predict_fn, fullpath, output_folder_prefix, 225 | SAVE_FRAME=False, MAX_FRAME=1000, DO_VISUALS=True) -------------------------------------------------------------------------------- /training_pipeline/trainCNN_tfrecords.py: -------------------------------------------------------------------------------- 1 | # Training pipeline loading images from tfrecords 2 | import tensorflow as tf 3 | import numpy as np 4 | import random 5 | import glob 6 | import preprocess 7 | from argparse import ArgumentParser 8 | from model import cnn_model_fn 9 | from time import time 10 | 11 | def randomize_image(img, contrast_range=[0.9,1.1], brightness_max=0.5): 12 | # Apply random flips/rotations and contrast/brightness changes to image 13 | img = tf.image.random_flip_left_right(img) 14 | img = tf.image.random_flip_up_down(img) 15 | img = tf.image.random_contrast(img, lower=contrast_range[0], upper=contrast_range[1]) 16 | img = tf.image.random_brightness(img, max_delta=brightness_max) 17 | img = tf.image.random_hue(img, max_delta=0.1) 18 | img = tf.image.random_saturation(img, lower=0.9, upper=1.1) 19 | img = tf.contrib.image.rotate(img, tf.random_uniform([1], minval=-np.pi, maxval=np.pi)) 20 | return img 21 | 22 | def add_image_gradient(img): 23 | # Add image gradient x and y to the 3rd dimension 24 | # Input is 21x21x3, output should be 21x21x5 25 | # gx 21x21x3 for each color channel 26 | # gx = tf.concat([tf.zeros([1,img.shape[1],img.shape[2]], dtype=img.dtype), img[1:,:,:] - img[:-1,:,:]], axis=0) 27 | # # gy ditto 28 | # gy = tf.concat([tf.zeros([img.shape[1],1,img.shape[2]], dtype=img.dtype), img[:,1:,:] - img[:,:-1,:]], axis=1) 29 | 30 | sobel_x = tf.constant([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], tf.float32) 31 | sobel_x_filter = tf.reshape(sobel_x, [3, 3, 1, 1]) 32 | sobel_y_filter = tf.transpose(sobel_x_filter, [1, 0, 2, 3]) 33 | 34 | # Shape = 1 x height x width x 1. 35 | gx = tf.nn.conv2d(tf.expand_dims(img, -1), sobel_x_filter, 36 | strides=[1, 1, 1, 1], padding='SAME') 37 | gy = tf.nn.conv2d(tf.expand_dims(img, -1), sobel_y_filter, 38 | strides=[1, 1, 1, 1], padding='SAME') 39 | 40 | gx = tf.squeeze(gx, axis=3) 41 | gy = tf.squeeze(gy, axis=3) 42 | 43 | # Reduce max gx and gy across their 3 channels 44 | gx = tf.reduce_max(gx, reduction_indices=[2], keepdims=True) 45 | gy = tf.reduce_max(gy, reduction_indices=[2], keepdims=True) 46 | 47 | # img = tf.concat([img, gx, gy], axis=2) 48 | img = tf.concat([gx, gy], axis=2) # Only keep image gradients 49 | return img 50 | 51 | def input_fn(inp_dataset, is_training=False, batch_size=50, buffer_size=1000000, prefetch_buffer_size=5000): 52 | # Returns an appropriate input function for training/evaluation. 53 | def sub_input_fn(): 54 | dataset = inp_dataset 55 | 56 | # Shuffle if training 57 | if is_training: 58 | dataset = dataset.shuffle(buffer_size=buffer_size) 59 | 60 | # Make a float image 61 | dataset = dataset.map(lambda img, label: (tf.to_float(img), label), num_parallel_calls=4) 62 | 63 | if is_training: 64 | # Slightly randomize images. 65 | dataset = dataset.map(lambda img, label: (randomize_image(img), label), num_parallel_calls=4) 66 | 67 | 68 | # Get image gradient only, turning into 2 channels gx and gy 69 | # dataset = dataset.map(lambda img, label: (add_image_gradient(tf.to_float(img)), label), num_parallel_calls=4) 70 | 71 | # Bring down to 45x45 to 31x31 72 | dataset = dataset.map(lambda img, label: (tf.reshape( 73 | tf.image.central_crop(img, 0.6888888888), [15,15,3], name='ReshapeAfterCentalCrop') 74 | , label), num_parallel_calls=4) 75 | # dataset = dataset.map(lambda img, label: (tf.image.central_crop(img, 0.66666666666), label), num_parallel_calls=4) 76 | 77 | 78 | # Shuffle/Batch/prefetch should happen after size changes 79 | 80 | # Batch and repeat. 81 | if is_training: 82 | dataset = dataset.repeat() 83 | dataset = dataset.batch(batch_size) 84 | dataset = dataset.prefetch(buffer_size=prefetch_buffer_size) # Doesn't necessarily need to be same as shuffle. 85 | 86 | # Build iterator and return 87 | one_shot_iterator = dataset.make_one_shot_iterator() 88 | next_element = one_shot_iterator.get_next() 89 | 90 | # Return in a dict so the premade estimators can use it. 91 | return {"x": next_element[0]}, next_element[1] 92 | return sub_input_fn 93 | 94 | def _parse_function(example_proto): 95 | features = {"image": tf.FixedLenFeature((), tf.string), 96 | "label": tf.FixedLenFeature((), tf.int64)} 97 | parsed_features = tf.parse_single_example(example_proto, features) 98 | image = parsed_features['image'] 99 | image = tf.decode_raw(image, tf.uint8, name='InitDecodeRaw') 100 | # image = tf.reshape(image,[45, 45, 3]) 101 | # Note: Make sure input dataset isn't corrupted with different sized tensors (via say RGBA instead of RGB) 102 | image = tf.reshape(image,[21, 21, 3], name='InitReshape21x21x3') 103 | return image, parsed_features["label"] 104 | 105 | def main(args): 106 | random.seed(100) 107 | 108 | # filenames = glob.glob('%s/*.tfrecords' % args.tfrecords_path) 109 | # train_filenames = filenames[:-2] 110 | # test_filenames = filenames[2:] 111 | # winsize22 112 | # train_filenames = ['../datasets/tfrecords/winsize_22_color/wgm_1_ws22.tfrecords', 113 | # '../datasets/tfrecords/winsize_22_color/swivel_ws22.tfrecords', 114 | # '../datasets/tfrecords/winsize_22_color/sam2_ws22.tfrecords', 115 | # '../datasets/tfrecords/winsize_22_color/speedchess1_ws22.tfrecords', 116 | # '../datasets/tfrecords/winsize_22_color/carlsen_match2_ws22.tfrecords', 117 | # '../datasets/tfrecords/winsize_22_color/output_ws22.tfrecords', 118 | # '../datasets/tfrecords/winsize_22_color/john2_ws22.tfrecords', 119 | # '../datasets/tfrecords/winsize_22_color/match1_ws22.tfrecords', 120 | # '../datasets/tfrecords/winsize_22_color/john1_ws22.tfrecords', 121 | # '../datasets/tfrecords/winsize_22_color/random1_ws22.tfrecords', 122 | # '../datasets/tfrecords/winsize_22_color/gm_magnus_1_ws22.tfrecords', 123 | # '../datasets/tfrecords/winsize_22_color/chess_beer_ws22.tfrecords', 124 | # '../datasets/tfrecords/winsize_22_color/match2_ws22.tfrecords', 125 | # '../datasets/tfrecords/winsize_22_color/bro_1_ws22.tfrecords'] 126 | # test_filenames = ['../datasets/tfrecords/winsize_22_color/sam1_ws22.tfrecords'] 127 | train_filenames = [ 128 | # '../datasets/tfrecords/winsize_10_color/wgm_1_ws10.tfrecords', 129 | '../datasets/tfrecords/winsize_10_color/input_images_ws10.tfrecords', 130 | # '../datasets/tfrecords/winsize_10_color/swivel_ws10.tfrecords', 131 | '../datasets/tfrecords/winsize_10_color/sam2_ws10.tfrecords', 132 | '../datasets/tfrecords/winsize_10_color/speedchess1_ws10.tfrecords', 133 | # '../datasets/tfrecords/winsize_10_color/carlsen_match2_ws10.tfrecords', 134 | # '../datasets/tfrecords/winsize_10_color/output_ws10.tfrecords', 135 | '../datasets/tfrecords/winsize_10_color/john2_ws10.tfrecords', 136 | # '../datasets/tfrecords/winsize_10_color/match1_ws10.tfrecords', 137 | '../datasets/tfrecords/winsize_10_color/john1_ws10.tfrecords', 138 | # '../datasets/tfrecords/winsize_10_color/random1_ws10.tfrecords', 139 | # '../datasets/tfrecords/winsize_10_color/gm_magnus_1_ws10.tfrecords', 140 | '../datasets/tfrecords/winsize_10_color/chess_beer_ws10.tfrecords', 141 | # '../datasets/tfrecords/winsize_10_color/match2_ws10.tfrecords', 142 | '../datasets/tfrecords/winsize_10_color/bro_1_ws10.tfrecords' 143 | ] 144 | # train_filenames = ['../datasets/tfrecords/winsize_10_color/input_images_ws10.tfrecords',] 145 | # train_filenames = ['../datasets/tfrecords/winsize_10_color/sam2_ws10.tfrecords',] 146 | test_filenames = ['../datasets/tfrecords/winsize_10_color/sam1_ws10.tfrecords'] # Used for ultrasmall v3 96% 147 | # test_filenames = ['../datasets/tfrecords/winsize_10_color/chess_beer_ws10.tfrecords'] 148 | 149 | 150 | with tf.Session() as sess: 151 | print("In session") 152 | train_dataset = tf.data.TFRecordDataset(train_filenames) 153 | # Convert img byte str back to 21x21 np.uint8 array. 154 | train_dataset = train_dataset.map(_parse_function) 155 | 156 | test_dataset = tf.data.TFRecordDataset(test_filenames) 157 | test_dataset = test_dataset.map(_parse_function) 158 | 159 | # Build model. 160 | model_dir = './training_models/cnn_tfrecord_%s' % (args.run_name) 161 | 162 | ### 163 | # one_shot_iterator = train_dataset.make_one_shot_iterator() 164 | # next_element = one_shot_iterator.get_next() 165 | # a,b = next_element[0], next_element[1] 166 | # for i in range(100): 167 | # x = sess.run(a) 168 | # if x.shape != (21,21,3): 169 | # print(i, x.shape, x.dtype) 170 | # exit() 171 | ### 172 | 173 | ### 174 | # a, b = input_fn(train_dataset, is_training=True)() 175 | # x = sess.run(a)['x'] 176 | # print(x.shape) 177 | # print(x.dtype) 178 | # exit() 179 | ### 180 | 181 | estimator = tf.estimator.Estimator(model_fn=cnn_model_fn, model_dir=model_dir, 182 | params={ 183 | # CNN filter layers 1 and 2, and then dense layer # of units 184 | 'filter_sizes':args.filter_sizes 185 | }) 186 | 187 | print("Using CNN ChessXCorner Classifier, output to %s" % (model_dir)) 188 | ############################ 189 | # Train (For steps * train_steps = total steps) 190 | train_steps = 100 191 | for i in range(args.steps): 192 | # Test 193 | print("\n\tTraining %d-%d/%d" % (i*train_steps,(i+1)*train_steps,args.steps*train_steps)) 194 | ta = time() 195 | estimator.train( 196 | input_fn=input_fn(train_dataset, is_training=True, batch_size=args.batch_size), 197 | steps=train_steps) 198 | tb = time() 199 | 200 | # Evaluate 201 | print('\n\tEvaluating...') 202 | metrics = estimator.evaluate(input_fn=input_fn(test_dataset)) 203 | accuracy_score = metrics["accuracy"] 204 | print("-- Test Accuracy: {0:f}\n".format(accuracy_score)) 205 | tc = time() 206 | print("Train took %g, Evaluate took %g" % (tb-ta, tc-tb)) 207 | 208 | if __name__ == '__main__': 209 | parser = ArgumentParser() 210 | # parser.add_argument("-m", "--max_each", dest="max_count_each_entries", type=int, 211 | # help="Maximum count of good or bad tiles per each folder") 212 | parser.add_argument("--name", dest="run_name", required=True, 213 | help="Name of model run name") 214 | parser.add_argument("--tfrecords_path", default='datasets/tfrecords/winsize10', 215 | help="Folder to load tfrecord output") 216 | parser.add_argument("-s", "--steps", dest="steps", required=True, type=int, 217 | help="Number of eval steps to run (x1000 train steps).") 218 | parser.add_argument("-bs", "--batch_size", type=int, default=50, help="Batch size.") 219 | parser.add_argument("-fs", "--filter_sizes", dest="filter_sizes", nargs='+', type=int, 220 | default=[32, 64, 1024], help="CNN model filter sizes") 221 | 222 | 223 | args = parser.parse_args() 224 | print("Arguments passed: \n\t%s\n" % args) 225 | main(args) -------------------------------------------------------------------------------- /training_pipeline/trainML_CNN_pipeline.py: -------------------------------------------------------------------------------- 1 | # Training pipeline loading images from filepaths. 2 | import tensorflow as tf 3 | import numpy as np 4 | import random 5 | import preprocess 6 | from argparse import ArgumentParser 7 | from model import cnn_model_fn 8 | from time import time 9 | 10 | def main(args): 11 | random.seed(100) 12 | make_equal = not args.use_all_entries 13 | # Load test dataset, used for evaluation 14 | test_folders = ['../datasets/test_datasetA',] 15 | test_imgs, test_labels, test_dataset_length, n_test_good = preprocess.loadMultipleDatapaths(test_folders, make_equal=make_equal, pre_shuffle=True) 16 | if args.max_count_each_entries: 17 | print('Using %d max_count_each_entries' % args.max_count_each_entries) 18 | print("Loaded test dataset '%s' : %d entries (%d good, %d bad)" % (test_folders, test_dataset_length, n_test_good, test_dataset_length - n_test_good )) 19 | 20 | # Load training dataset(s). 21 | if args.training_folders: 22 | train_folders = args.training_folders 23 | else: 24 | # train_folders = ['../datasets/dataset_gray_10'] 25 | allfiles = ['chess_beer.mp4', 'random1.mp4', 'match2.mp4','output.avi','output.mp4', 26 | 'speedchess1.mp4','wgm_1.mp4','gm_magnus_1.mp4', 27 | 'bro_1.mp4','output2.avi','john1.mp4','john2.mp4','swivel.mp4', 28 | 'sam1.mp4', 'sam2.mp4',] 29 | train_folders = ['../datasets/dataset_gray_10'] 30 | train_folders.extend(map(lambda x: '../results/%sam1_vidstream_frames/tiles' % x[:-4], allfiles )) 31 | 32 | train_imgs, train_labels, train_dataset_length, n_train_good = preprocess.loadMultipleDatapaths(train_folders, make_equal=make_equal, pre_shuffle=True, max_count_each_entries=args.max_count_each_entries) 33 | print("Loaded train dataset '%s' : %d entries (%d good, %d bad)" % (train_folders, train_dataset_length, n_train_good, train_dataset_length - n_train_good )) 34 | 35 | # Build model. 36 | model_dir = './training_models/cnn_%s_n%d' % (args.run_name, train_dataset_length) 37 | 38 | estimator = tf.estimator.Estimator(model_fn=cnn_model_fn, model_dir=model_dir, 39 | params={ 40 | # CNN filter layers 1 and 2, and then dense layer # of units 41 | 'filter_sizes':args.filter_sizes 42 | }) 43 | 44 | print("Using CNN ChessXCorner Classifier, output to %s" % (model_dir)) 45 | ############################ 46 | # Train (For steps * train_steps = total steps) 47 | train_steps = 100 48 | for i in range(args.steps): 49 | # Test 50 | print("Training %d-%d/%d" % (i*train_steps,(i+1)*train_steps,args.steps*train_steps)) 51 | ta = time() 52 | estimator.train( 53 | input_fn=preprocess.input_fn(train_imgs, train_labels, train_dataset_length, is_training=True), 54 | steps=train_steps) 55 | tb = time() 56 | 57 | # Evaluate 58 | metrics = estimator.evaluate(input_fn=preprocess.input_fn(test_imgs, test_labels)) 59 | accuracy_score = metrics["accuracy"] 60 | print("-- Test Accuracy: {0:f}\n".format(accuracy_score)) 61 | tc = time() 62 | print("Train took %g, Evaluate took %g" % (tb-ta, tc-tb)) 63 | 64 | 65 | ############################ 66 | # Validate 67 | # predictions = estimator.predict(input_fn=preprocess.input_fn(test_imgs, test_labels)) 68 | 69 | # test_labels = np.array(f_labels[split:]) 70 | 71 | # count_good = 0 72 | # for i,(prediction,true_answer) in enumerate(zip(predictions, test_labels.astype(np.int))): 73 | # if prediction['probabilities'].argmax() != true_answer: 74 | # # print("Failure on %d: %s" % (i, prediction['probabilities'])) 75 | # continue 76 | # else: 77 | # count_good += 1 78 | 79 | # success_rate = float(count_good) / len(test_labels) 80 | # print("Total %d/%d right ~ %.2f%% success rate" % (count_good, len(test_labels), 100*success_rate)) 81 | 82 | if __name__ == '__main__': 83 | parser = ArgumentParser() 84 | parser.add_argument("-m", "--max_each", dest="max_count_each_entries", type=int, 85 | help="Maximum count of good or bad tiles per each folder") 86 | parser.add_argument("--name", dest="run_name", required=True, 87 | help="Name of model run name") 88 | parser.add_argument("-s", "--steps", dest="steps", required=True, type=int, 89 | help="Number of eval steps to run (x1000 train steps).") 90 | parser.add_argument("training_folders", nargs='+', 91 | help="Training folder tile paths used for training") 92 | parser.add_argument("-fs", "--filter_sizes", dest="filter_sizes", nargs='+', type=int, 93 | default=[32, 64, 1024], help="CNN model filter sizes") 94 | parser.add_argument("-use_all_entries", 95 | action='store_true', help="Whether to make input datasets equal good/bad.") 96 | 97 | 98 | args = parser.parse_args() 99 | print("Arguments passed: \n\t%s\n" % args) 100 | main(args) -------------------------------------------------------------------------------- /training_pipeline/trainML_pipeline.py: -------------------------------------------------------------------------------- 1 | # Training pipeline loading images from filepaths. 2 | import tensorflow as tf 3 | import numpy as np 4 | import random 5 | import preprocess 6 | from argparse import ArgumentParser 7 | 8 | def main(args): 9 | random.seed(100) 10 | # Load test dataset, used for evaluation 11 | test_folders = ['../datasets/test_datasetA',] 12 | test_imgs, test_labels, test_dataset_length, n_test_good = preprocess.loadMultipleDatapaths(test_folders, make_equal=True, pre_shuffle=True) 13 | if args.max_count_each_entries: 14 | print('Using %d max_count_each_entries' % args.max_count_each_entries) 15 | print("Loaded test dataset '%s' : %d entries (%d good, %d bad)" % (test_folders, test_dataset_length, n_test_good, test_dataset_length - n_test_good )) 16 | 17 | # Load training dataset(s). 18 | if args.training_folders: 19 | train_folders = args.training_folders 20 | else: 21 | # train_folders = ['../datasets/dataset_gray_10'] 22 | allfiles = ['chess_beer.mp4', 'random1.mp4', 'match2.mp4','output.avi','output.mp4', 23 | 'speedchess1.mp4','wgm_1.mp4','gm_magnus_1.mp4', 24 | 'bro_1.mp4','output2.avi','john1.mp4','john2.mp4','swivel.mp4', 25 | 'sam1.mp4', 'sam2.mp4',] 26 | train_folders = ['../datasets/dataset_gray_10'] 27 | train_folders.extend(map(lambda x: '../results/%sam1_vidstream_frames/tiles' % x[:-4], allfiles )) 28 | 29 | train_imgs, train_labels, train_dataset_length, n_train_good = preprocess.loadMultipleDatapaths(train_folders, make_equal=True, pre_shuffle=True, max_count_each_entries=args.max_count_each_entries) 30 | print("Loaded train dataset '%s' : %d entries (%d good, %d bad)" % (train_folders, train_dataset_length, n_train_good, train_dataset_length - n_train_good )) 31 | 32 | # Build model. 33 | feature_img = tf.feature_column.numeric_column("x", shape=[21,21], dtype=tf.uint8) 34 | # units = [512,256,128] 35 | units = [256,32,8] 36 | model_str = "_".join(map(str,units)) 37 | model_dir = './training_models/run1_%s_%s_dataset_%d' % (model_str, args.run_name, train_dataset_length) 38 | 39 | print("Using DNNClassifier %s : Output to %s" % (units, model_dir)) 40 | estimator = tf.estimator.DNNClassifier( 41 | feature_columns=[feature_img], 42 | hidden_units=units, 43 | n_classes=2, 44 | dropout=0.1, 45 | model_dir=model_dir, 46 | # model_dir='./training_models/lin_win%s' % (dataset_folder[dataset_folder.rfind('_')+1:]), 47 | ) 48 | ############################ 49 | # Train (For steps * train_steps = total steps) 50 | train_steps = 1000 51 | for i in range(args.steps): 52 | # Test 53 | print("Training %d-%d/%d" % (i*train_steps,(i+1)*train_steps,args.steps*train_steps)) 54 | estimator.train(input_fn=preprocess.input_fn(train_imgs, train_labels, train_dataset_length, is_training=True), steps=train_steps) 55 | 56 | # Evaluate 57 | metrics = estimator.evaluate(input_fn=preprocess.input_fn(test_imgs, test_labels)) 58 | accuracy_score = metrics["accuracy"] 59 | print("-- Test Accuracy: {0:f}\n".format(accuracy_score)) 60 | 61 | 62 | ############################ 63 | # Validate 64 | # predictions = estimator.predict(input_fn=preprocess.input_fn(test_imgs, test_labels)) 65 | 66 | # test_labels = np.array(f_labels[split:]) 67 | 68 | # count_good = 0 69 | # for i,(prediction,true_answer) in enumerate(zip(predictions, test_labels.astype(np.int))): 70 | # if prediction['probabilities'].argmax() != true_answer: 71 | # # print("Failure on %d: %s" % (i, prediction['probabilities'])) 72 | # continue 73 | # else: 74 | # count_good += 1 75 | 76 | # success_rate = float(count_good) / len(test_labels) 77 | # print("Total %d/%d right ~ %.2f%% success rate" % (count_good, len(test_labels), 100*success_rate)) 78 | 79 | if __name__ == '__main__': 80 | parser = ArgumentParser() 81 | parser.add_argument("-m", "--max_each", dest="max_count_each_entries", type=int, 82 | help="Maximum count of good or bad tiles per each folder") 83 | parser.add_argument("--name", dest="run_name", required=True, 84 | help="Name of model run name") 85 | parser.add_argument("-s", "--steps", dest="steps", required=True, type=int, 86 | help="Number of eval steps to run (x1000 train steps).") 87 | parser.add_argument("training_folders", nargs='+', 88 | help="Training folder tile paths used for training") 89 | 90 | 91 | args = parser.parse_args() 92 | print(args) 93 | main(args) -------------------------------------------------------------------------------- /triangle_mesh.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Elucidation/ChessboardDetect/ece9f9dafd1376e37b1462570e3b49d82c5e8ffc/triangle_mesh.png -------------------------------------------------------------------------------- /view_xcorner_heatmap.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | import PIL.Image 4 | np.set_printoptions(suppress=True, linewidth=200) # Better printing of arrays 5 | 6 | # filepath = 'input/img_01.jpg' 7 | filepath = 'testB.jpg' 8 | img = PIL.Image.open(filepath).convert('L') # Grayscale full image with chessboard 9 | img_width, img_height = img.size 10 | aspect_ratio = min(200.0/img_width, 200.0/img_height) 11 | new_width, new_height = ((np.array(img.size) * aspect_ratio)).astype(int) 12 | img = img.resize((new_width,new_height), resample=PIL.Image.BILINEAR) 13 | img = np.array(img) 14 | heatmap[i] = prediction['probabilities'][1] # Probability is saddle 15 | winsize=5 16 | 17 | print(img.shape) 18 | 19 | 20 | print(img.shape[0]-winsize-1 - (winsize+1)) 21 | print(img.shape[1]-winsize-1 - (winsize+1)) 22 | 23 | new_size = np.array([img.shape[0] - 2*winsize - 1, img.shape[1] - 2*winsize - 1]) 24 | 25 | print(new_size) 26 | 27 | heatmap = np.load('heatmap.npy') 28 | heatmap = np.reshape(heatmap, new_size) 29 | print(heatmap) 30 | # heatmap[heatmap<0.5] = 0 31 | 32 | plt.imshow(heatmap) 33 | plt.show() --------------------------------------------------------------------------------