eye_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Collect Darkchess Image or Video Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "import cv2\n",
    "import IPython.display as display\n",
    "from datasets import Dataset\n",
    "from datasets.features import Features, Image, ClassLabel\n",
    "\n",
    "from config import CHESS\n",
    "from eye.imgproc import (\n",
    "    rotate,\n",
    "    get_board_frame,\n",
    "    get_form_frame,\n",
    "    get_chess_frame\n",
    ")\n",
    "\n",
    "# ========================================\n",
    "# Custom Settings\n",
    "# ========================================\n",
    "CAP_SRC = None  # set your IP camera URL, or video file path\n",
    "SAVE_VIDEO_PATH = None  # set your video file path to save, or None to disable\n",
    "SAVE_IMAGE_DIR = None # set your image directory to save, or None to disable\n",
    "DISPLAY_CAP = False # display captured frame\n",
    "\n",
    "SPECIFIC_CHESS_LABEL = None\n",
    "CHESS_IMAGE_NUM = 144000 # 50 x 360 x 8 = 144000\n",
    "image_data = {\n",
    "    \"image\": [],\n",
    "    \"label\": []\n",
    "}\n",
    "\n",
    "# Example: process board frame to chess frame and save\n",
    "def process_image_to_save(frame, frame_count) -> None:\n",
    "    if frame_count % 5 != 0 or len(image_data[\"image\"]) >= CHESS_IMAGE_NUM:\n",
    "        return\n",
    "\n",
    "    for angle in range(360):\n",
    "        for shift in range(8):\n",
    "            img = get_board_frame(frame)\n",
    "            img = get_form_frame(img, 20)\n",
    "            img = get_chess_frame(img, shift)\n",
    "            img = rotate(img, angle)\n",
    "            image_data[\"image\"].append(img)\n",
    "            image_data[\"label\"].append(SPECIFIC_CHESS_LABEL)\n",
    "\n",
    "    print(f\"Now image data number: {len(image_data['image'])}\", end=\"\\r\")\n",
    "\n",
    "# ========================================\n",
    "# End of Custom Settings\n",
    "# ========================================\n",
    "\n",
    "cap = cv2.VideoCapture(CAP_SRC)\n",
    "display_handle = display.display(None, display_id=True)\n",
    "out = cv2.VideoWriter(\n",
    "    SAVE_VIDEO_PATH,\n",
    "    cv2.VideoWriter_fourcc(*\"mp4v\"), \n",
    "    10,\n",
    "    (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))\n",
    ")\n",
    "\n",
    "while True:\n",
    "    try:\n",
    "        frame = cap.read()[1]\n",
    "        if frame is None:\n",
    "            print(\"\\nEnd of video\")\n",
    "            break\n",
    "\n",
    "        if SAVE_IMAGE_DIR:\n",
    "            os.makedirs(SAVE_IMAGE_DIR, exist_ok=True)\n",
    "            process_image_to_save(frame, int(cap.get(cv2.CAP_PROP_POS_FRAMES)))\n",
    "\n",
    "        if SAVE_VIDEO_PATH:\n",
    "            os.makedirs(os.path.dirname(SAVE_VIDEO_PATH), exist_ok=True)\n",
    "            out.write(frame)\n",
    "        \n",
    "        if DISPLAY_CAP:\n",
    "            time.sleep(0.5)  # 2 fps, buffer time\n",
    "            encoded_frame = cv2.imencode(\".jpg\", frame)[1]\n",
    "            display_handle.update(display.Image(data=encoded_frame.tobytes()))\n",
    "\n",
    "    except KeyboardInterrupt:\n",
    "        break\n",
    "\n",
    "cap.release()\n",
    "out.release()\n",
    "display_handle.update(None)\n",
    "if SAVE_IMAGE_DIR:\n",
    "    print(f\"Saving {len(image_data['image'])} images to {SAVE_IMAGE_DIR}\")\n",
    "    image_datasets = Dataset.from_dict(\n",
    "        mapping=image_data,\n",
    "        features=Features({\"image\": Image(\"L\"), \"label\": ClassLabel(num_classes=len(CHESS[:-2]), names=[item[\"code\"] for item in CHESS[:-2]])})\n",
    "    ).train_test_split(train_size=0.8, test_size=0.2, shuffle=False)\n",
    "    image_datasets.save_to_disk(SAVE_IMAGE_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Login to HuggingFace (just login once)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import interpreter_login\n",
    "interpreter_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Push Board Dataset to Hugging Face\n",
    "> manually write `metadata.jsonl` is required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"datasets/darkchess-board\")\n",
    "dataset.push_to_hub(\"ryanlinjui/darkchess-board\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Combine All Chess Datasets and Push it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import DatasetDict\n",
    "from datasets import load_from_disk, concatenate_datasets\n",
    "\n",
    "from config import CHESS\n",
    "\n",
    "train_datasets = []\n",
    "test_datasets = []\n",
    "\n",
    "for i in range(len(CHESS[:-2])):\n",
    "    print(f\"Loading dataset {i}\", end=\"\\r\")\n",
    "    dataset = load_from_disk(f\"datasets/darkchess-single/{i}\")\n",
    "    train_datasets.append(dataset[\"train\"])\n",
    "    test_datasets.append(dataset[\"test\"])\n",
    "\n",
    "datasets = DatasetDict({\n",
    "    \"train\": concatenate_datasets(train_datasets),\n",
    "    \"test\": concatenate_datasets(test_datasets)\n",
    "})\n",
    "datasets.push_to_hub(\"ryanlinjui/darkchess-single\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Dataset from HuggingFace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from datasets import load_dataset\n",
    "from tensorflow.keras.utils import img_to_array\n",
    "\n",
    "datasets = load_dataset(\"ryanlinjui/darkchess-single\")\n",
    "\n",
    "train_images = np.array([img_to_array(image[\"image\"]) for image in datasets[\"train\"]]).astype(\"float32\") / 255.0\n",
    "train_labels = np.array([image[\"label\"] for image in datasets[\"train\"]])\n",
    "\n",
    "test_images = np.array([img_to_array(image[\"image\"]) for image in datasets[\"test\"]]).astype(\"float32\") / 255.0\n",
    "test_labels = np.array([image[\"label\"] for image in datasets[\"test\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# List devices, check for CPU, GPU, MPU availability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.python.client import device_lib \n",
    "print(device_lib.list_local_devices())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Construct the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras import Sequential\n",
    "from keras.optimizers import Adam\n",
    "from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout\n",
    "\n",
    "model = Sequential()\n",
    "model.add(Conv2D(filters=64, kernel_size=(3, 3), strides=(1, 1), input_shape=(56, 56, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(Conv2D(filters=64, kernel_size=(3, 3), strides=(1, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(MaxPooling2D(pool_size=(2, 2)))\n",
    "model.add(Conv2D(filters=128, kernel_size=(3, 2), strides=(1, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(Conv2D(filters=128, kernel_size=(3, 3), strides=(1, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(MaxPooling2D(pool_size=(2,2)))\n",
    "model.add(Conv2D(filters=256, kernel_size=(3, 3), strides=(1, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(Conv2D(filters=256, kernel_size=(3, 3), strides=(1, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(Conv2D(filters=256, kernel_size=(3, 3), strides=(1, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(MaxPooling2D(pool_size=(2,2)))\n",
    "model.add(Conv2D(filters=512, kernel_size=(3, 3), strides=(1, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(Conv2D(filters=512, kernel_size=(3, 3), strides=(1, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(Conv2D(filters=512, kernel_size=(3, 3), strides=(1, 1), padding=\"same\", activation=\"relu\"))\n",
    "model.add(MaxPooling2D(pool_size=(2, 2)))\n",
    "model.add(Flatten())\n",
    "model.add(Dense(units=4096, activation=\"relu\"))\n",
    "model.add(Dropout(rate=0.5))\n",
    "model.add(Dense(units=4096, activation=\"relu\"))\n",
    "model.add(Dropout(rate=0.5))\n",
    "model.add(Dense(units=1000, activation=\"relu\"))\n",
    "model.add(Dropout(rate=0.5))\n",
    "model.add(Dense(units=14, activation=\"softmax\"))\n",
    "\n",
    "model.compile(\n",
    "    optimizer=Adam(lr=0.00001, beta_1=0.9, beta_2=0.999, epsilon=1e-08),\n",
    "    loss = \"categorical_crossentropy\",\n",
    "    metrics=[\"accuracy\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gray_frame1 = np.reshape(gray_frame,[-1, 56, 56, 1])\n",
    "gray_label1 = np_utils.to_categorical(gray_label, num_classes=14)\n",
    "history = model.fit(x=gray_frame1, y=gray_label1, epochs=15, batch_size=360)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.evaluate(x=test_images, y=test_labels, batch_size=360)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check predictions on test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,10))\n",
    "predict_index = 0\n",
    "for i in range(100):\n",
    "    plt.subplot(10,10,i+1)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    plt.grid(False)\n",
    "    a = gray_frame[predict_index+i].reshape([-1,56,56,1])\n",
    "    m = model.predict(a)\n",
    "    plt.imshow(gray_frame[predict_index+i])\n",
    "    plt.xlabel(m.argmax())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save(\"cnn_darkchess_model.h5\")\n",
    "model.save_weights(\"cnn_darkchess_weights.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Push Model to Hugging Face"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from tensorflow.keras.utils import plot_model\n",
    "from tensorflow.keras.models import load_model\n",
    "from huggingface_hub import push_to_hub_keras\n",
    "\n",
    "from config import DEFAULT_MODEL_NAME, DEFAULT_WEIGHTS_NAME\n",
    "\n",
    "model_dir = \"\"\n",
    "\n",
    "model = load_model(os.path.join(model_dir, DEFAULT_MODEL_NAME))\n",
    "model.load_weights(os.path.join(model_dir, DEFAULT_WEIGHTS_NAME))\n",
    "plot_model(model, show_shapes=True, show_dtype=True, show_layer_activations=True, expand_nested=True)\n",
    "model.summary()\n",
    "\n",
    "push_to_hub_keras(model, \"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "predictions = model.predict_classes(gray_frame1) \n",
    "pd.crosstab(gray_label, predictions, rownames=[\"實際值\"], colnames=[\"預測值\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(history.history[\"accuracy\"], label=\"accuracy\")\n",
    "plt.plot(history.history[\"loss\"], label = \"loss\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.ylim([-0.05, 1.05])\n",
    "plt.xlim([0, 5])\n",
    "plt.legend(loc=\"lower right\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}