make-VIA-project-from-images.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "ac1b54da",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "# Root directory of the project\n",
    "from PIL import Image\n",
    "\n",
    "ROOT_DIR = os.path.abspath(\".\")\n",
    "\n",
    "from mrcnn import utils\n",
    "import mrcnn.model as modellib\n",
    "from mrcnn import visualize\n",
    "import skimage.io\n",
    "\n",
    "# Directory to save logs and trained model\n",
    "MODEL_DIR = os.path.join(ROOT_DIR, \"logs\")\n",
    "\n",
    "# Directory of images to run detection on\n",
    "IMAGE_DIR = os.path.join(ROOT_DIR, \"images\")\n",
    "\n",
    "import samples.coco.coco as coco\n",
    "import samples.police2.backlash as backlash\n",
    "\n",
    "\n",
    "class MaskRCNNModel():\n",
    "    MODEL_FILE_NAME = \"mask_rcnn_coco.h5\"\n",
    "    CLASSES = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',\n",
    "               'bus', 'train', 'truck', 'boat', 'traffic light',\n",
    "               'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',\n",
    "               'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',\n",
    "               'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',\n",
    "               'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',\n",
    "               'kite', 'baseball bat', 'baseball glove', 'skateboard',\n",
    "               'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',\n",
    "               'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',\n",
    "               'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',\n",
    "               'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',\n",
    "               'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',\n",
    "               'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',\n",
    "               'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',\n",
    "               'teddy bear', 'hair drier', 'toothbrush']\n",
    "\n",
    "    CONFIG_CLASS = coco.CocoConfig\n",
    "\n",
    "    def __init__(self):\n",
    "        # Local path to trained weights file\n",
    "        COCO_MODEL_PATH = os.path.join(ROOT_DIR, self.MODEL_FILE_NAME)\n",
    "\n",
    "        # Download COCO trained weights from Releases if needed\n",
    "        if not os.path.exists(COCO_MODEL_PATH):\n",
    "            utils.download_trained_weights(COCO_MODEL_PATH)\n",
    "\n",
    "        class InferenceConfig(self.CONFIG_CLASS):\n",
    "            # Set batch size to 1 since we'll be running inference on\n",
    "            # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU\n",
    "            GPU_COUNT = 1\n",
    "            IMAGES_PER_GPU = 1\n",
    "            NUM_CLASSES = len(self.CLASSES)  # COCO has 80 classes\n",
    "\n",
    "        self.config = InferenceConfig()\n",
    "        #         config.display()\n",
    "\n",
    "        # Create model object in inference mode.\n",
    "        model = modellib.MaskRCNN(mode=\"inference\", model_dir=MODEL_DIR, config=self.config)\n",
    "\n",
    "        # Load weights\n",
    "        model.load_weights(COCO_MODEL_PATH, by_name=True)\n",
    "\n",
    "        # COCO Class names\n",
    "        # Index of the class in the list is its ID. For example, to get ID of\n",
    "        # the teddy bear class, use: class_names.index('teddy bear')\n",
    "        class_names = self.CLASSES\n",
    "\n",
    "        self.model = model\n",
    "\n",
    "model_full = MaskRCNNModel()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e6e21e10",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 664/664 [1:06:44<00:00,  6.03s/it]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "from imantics import Polygons, Mask\n",
    "\n",
    "import json\n",
    "project = json.load(open(\"via-project-template.json\"))\n",
    "\n",
    "import glob\n",
    "images_path = \"/Users/skabbit/Projects/backlash/backlash/datasets/usa-police\"\n",
    "files = sorted(glob.glob(images_path + \"/*.jpg\"))\n",
    "\n",
    "data = {}\n",
    "\n",
    "for filename in tqdm(files):\n",
    "#     print(filename)\n",
    "    image = skimage.io.imread(filename)\n",
    "    results = model_full.model.detect([image], verbose=0)\n",
    "    \n",
    "    masks = results[0]['masks'][:,:,results[0]['class_ids'] == 1] #  & results[0]['score'] > 0.9 where?\n",
    "    regions = []\n",
    "    \n",
    "#     print(\"Regions:\", masks.shape[2])\n",
    "    for i in range(masks.shape[2]):\n",
    "        mask = masks[:,:,i]\n",
    "        polygons = Mask(mask).polygons()\n",
    "        polygons.simplify()\n",
    "        shape_attributes = {\n",
    "            \"name\": \"polygon\",\n",
    "            \"all_points_x\": [],\n",
    "            \"all_points_y\": []\n",
    "        }\n",
    "        for point in polygons.points[0]:\n",
    "            shape_attributes['all_points_x'].append(int(point[0]))\n",
    "            shape_attributes['all_points_y'].append(int(point[1]))        \n",
    "        regions.append({\n",
    "            \"shape_attributes\": shape_attributes,\n",
    "            \"region_attributes\": {},\n",
    "        })\n",
    "\n",
    "    fn = os.path.split(filename)[1]\n",
    "    size = os.path.getsize(filename)\n",
    "    data[fn + str(size)] = {\n",
    "        \"filename\": fn,\n",
    "        \"size\": size,\n",
    "        \"file_attributes\": {},\n",
    "        \"regions\": regions\n",
    "    }\n",
    "\n",
    "project['_via_img_metadata'] = data\n",
    "project[\"_via_image_id_list\"] = list(data.keys())\n",
    "project['_via_settings']['core']['default_filepath'] = images_path+\"/\"\n",
    "project['_via_settings']['core']['filepath'] = {images_path+\"/\": 1}\n",
    "\n",
    "json.dump(project, open(\"via-project.json\", \"w\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61102cd3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}