Update image_to_fen/fen.py

image_to_fen/fen.py

@@ -7,6 +7,8 @@ import torch
 import torchvision
 import numpy as np
 import re
+import chess
+from fentoimage.board import BoardImage
 
 import image_to_fen.util as util
 
@@ -94,6 +96,14 @@ def fen_from_onehot(one_hot):
 
     return output
 
+def fen_and_image(input):
+  itf = ImageToFen()
+  output = itf.predict(input)
+  fen = output.replace('-', '/')
+  renderer = BoardImage(fen)
+  image = renderer.render()
+  return fen, image
+    
 def main():
     """Run prediction on image."""
     parser = argparse.ArgumentParser(description="Predict FEN string for image of chess board.")