Rubiks-Cube-Face-Detection-Model.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "423cd7bb",
      "metadata": {
        "id": "423cd7bb"
      },
      "outputs": [],
      "source": [
        "PATH_TO_CKPT = \"/home2/testdev/TensorFlow/workspace/training_demo/exported-models/my_model/checkpoint/\"\n",
        "PATH_TO_CFG = \"/home2/testdev/TensorFlow/workspace/training_demo/exported-models/my_model/pipeline.config\"\n",
        "PATH_TO_LABELS = \"/home2/testdev/TensorFlow/workspace/training_demo/annotations/label_map.pbtxt\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dd6ee7d2",
      "metadata": {
        "id": "dd6ee7d2"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'    # Suppress TensorFlow logging\n",
        "import tensorflow as tf\n",
        "from object_detection.utils import label_map_util\n",
        "from object_detection.utils import config_util\n",
        "from object_detection.utils import visualization_utils as viz_utils\n",
        "from object_detection.builders import model_builder\n",
        "\n",
        "tf.get_logger().setLevel('ERROR')           # Suppress TensorFlow logging (2)\n",
        "\n",
        "# Enable GPU dynamic memory allocation\n",
        "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
        "for gpu in gpus:\n",
        "    tf.config.experimental.set_memory_growth(gpu, True)\n",
        "\n",
        "# Load pipeline config and build a detection model\n",
        "configs = config_util.get_configs_from_pipeline_file(PATH_TO_CFG)\n",
        "model_config = configs['model']\n",
        "detection_model = model_builder.build(model_config=model_config, is_training=False)\n",
        "\n",
        "# Restore checkpoint\n",
        "ckpt = tf.compat.v2.train.Checkpoint(model=detection_model)\n",
        "ckpt.restore(os.path.join(PATH_TO_CKPT, 'ckpt-0')).expect_partial()\n",
        "\n",
        "@tf.function\n",
        "def detect_fn(image):\n",
        "    \"\"\"Detect objects in image.\"\"\"\n",
        "\n",
        "    image, shapes = detection_model.preprocess(image)\n",
        "    prediction_dict = detection_model.predict(image, shapes)\n",
        "    detections = detection_model.postprocess(prediction_dict, shapes)\n",
        "\n",
        "    return detections, prediction_dict, tf.reshape(shapes, [-1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cccdc059",
      "metadata": {
        "id": "cccdc059"
      },
      "outputs": [],
      "source": [
        "category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS,use_display_name=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dd1d42bb",
      "metadata": {
        "id": "dd1d42bb"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "cap = cv2.VideoCapture(\"/home2/testdev/basics/test_video_1.mp4\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c2547a9e",
      "metadata": {
        "id": "c2547a9e"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "while True:\n",
        "    # Read frame from camera\n",
        "    ret, image_np = cap.read()\n",
        "    image_rgb = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)\n",
        "    # Expand dimensions since the model expects images to have shape: [1, None, None, 3]\n",
        "    image_np_expanded = np.expand_dims(image_rgb, axis=0)\n",
        "\n",
        "    # Things to try:\n",
        "    # Flip horizontally\n",
        "    # image_np = np.fliplr(image_np).copy()\n",
        "\n",
        "    # Convert image to grayscale\n",
        "    # image_np = np.tile(\n",
        "    #     np.mean(image_np, 2, keepdims=True), (1, 1, 3)).astype(np.uint8)\n",
        "\n",
        "    input_tensor = tf.convert_to_tensor(np.expand_dims(image_rgb, 0), dtype=tf.float32)\n",
        "    detections, predictions_dict, shapes = detect_fn(input_tensor)\n",
        "\n",
        "    label_id_offset = 1\n",
        "    image_np_with_detections = image_np.copy()\n",
        "\n",
        "    viz_utils.visualize_boxes_and_labels_on_image_array(\n",
        "          image_np_with_detections,\n",
        "          detections['detection_boxes'][0].numpy(),\n",
        "          (detections['detection_classes'][0].numpy() + label_id_offset).astype(int),\n",
        "          detections['detection_scores'][0].numpy(),\n",
        "          category_index,\n",
        "          use_normalized_coordinates=True,\n",
        "          max_boxes_to_draw=200,\n",
        "          min_score_thresh=.50,\n",
        "          line_thickness = 4,\n",
        "          agnostic_mode=False)\n",
        "\n",
        "    # Display output\n",
        "    cv2.imshow('object detection', cv2.resize(image_np_with_detections, (1280, 1024)))\n",
        "\n",
        "    if cv2.waitKey(25) & 0xFF == ord('q'):\n",
        "        break\n",
        "\n",
        "cap.release()\n",
        "cv2.destroyAllWindows()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d7d0af29",
      "metadata": {
        "id": "d7d0af29"
      },
      "outputs": [],
      "source": []
    }
  ],
  "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.7.13"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}