diff --git a/app/api/models.py b/app/api/models.py index 5b0c4a13..d701e278 100644 --- a/app/api/models.py +++ b/app/api/models.py @@ -1,25 +1,66 @@ from flask_restplus import Namespace, Resource, reqparse from werkzeug.datastructures import FileStorage -from imantics import Image as ImanticsImage +from imantics import Mask from flask_login import login_required from ..config import Config from PIL import Image +from ..models import ImageModel import os -MODEL_LOADED = len(Config.MASK_RCNN_FILE) != 0 and os.path.isfile(Config.MASK_RCNN_FILE) - -if MODEL_LOADED: +MASKRCNN_LOADED = os.path.isfile(Config.MASK_RCNN_FILE) +if MASKRCNN_LOADED: from ..util.mask_rcnn import model as maskrcnn else: print("MaskRCNN model is disabled.", flush=True) +DEXTR_LOADED = os.path.isfile(Config.DEXTR_FILE) +if DEXTR_LOADED: + from ..util.dextr import model as dextr +else: + print("DEXTR model is disabled.", flush=True) + api = Namespace('model', description='Model related operations') image_upload = reqparse.RequestParser() image_upload.add_argument('image', location='files', type=FileStorage, required=True, help='Image') +dextr_args = reqparse.RequestParser() +dextr_args.add_argument('points', location='json', type=list, required=True) +dextr_args.add_argument('padding', location='json', type=int, default=50) +dextr_args.add_argument('threshold', location='json', type=int, default=80) + + +@api.route('/dextr/') +class MaskRCNN(Resource): + + @login_required + @api.expect(dextr_args) + def post(self, image_id): + """ COCO data test """ + + if not DEXTR_LOADED: + return {"disabled": True, "message": "DEXTR is disabled"}, 400 + + args = dextr_args.parse_args() + points = args.get('points') + padding = args.get('padding') + threshold = args.get('threshold') + + if len(points) != 4: + return {"message": "Invalid points entered"}, 400 + + image_model = ImageModel.objects(id=image_id).first() + if not image_model: + return {"message": "Invalid image ID"}, 400 + + image = Image.open(image_model.path) + result = dextr.predict_mask(image, points) + + return { "segmentaiton": Mask(result).polygons().segmentation } + + @api.route('/maskrcnn') class MaskRCNN(Resource): @@ -27,7 +68,7 @@ class MaskRCNN(Resource): @api.expect(image_upload) def post(self): """ COCO data test """ - if not MODEL_LOADED: + if not MASKRCNN_LOADED: return {"disabled": True, "coco": {}} args = image_upload.parse_args() diff --git a/app/config.py b/app/config.py index e1a4917d..bf4e50b4 100644 --- a/app/config.py +++ b/app/config.py @@ -31,3 +31,5 @@ class Config: # Models MASK_RCNN_FILE = os.getenv("MASK_RCNN_FILE", "") MASK_RCNN_CLASSES = os.getenv("MASK_RCNN_CLASSES", "BG") + + DEXTR_FILE = os.getenv("DEXTR_FILE", "/models/dextr_pascal-sbd.h5") diff --git a/app/util/dextr.py b/app/util/dextr.py new file mode 100644 index 00000000..14e2e7e2 --- /dev/null +++ b/app/util/dextr.py @@ -0,0 +1,7 @@ +from dextr import DEXTR +from ..config import Config +import os + + +model = DEXTR(nb_classes=1, resnet_layers=101, input_shape=(512, 512), weights_path=Config.DEXTR_FILE, + num_input_channels=4, classifier='psp', sigmoid=True) \ No newline at end of file diff --git a/client/src/components/annotator/panels/DEXTRPanel.vue b/client/src/components/annotator/panels/DEXTRPanel.vue new file mode 100755 index 00000000..ab185705 --- /dev/null +++ b/client/src/components/annotator/panels/DEXTRPanel.vue @@ -0,0 +1,33 @@ + + + diff --git a/client/src/components/annotator/tools/DEXTRTool.vue b/client/src/components/annotator/tools/DEXTRTool.vue new file mode 100644 index 00000000..76d0b34a --- /dev/null +++ b/client/src/components/annotator/tools/DEXTRTool.vue @@ -0,0 +1,91 @@ + diff --git a/client/src/views/Annotator.vue b/client/src/views/Annotator.vue index 50d497ef..3d14968f 100755 --- a/client/src/views/Annotator.vue +++ b/client/src/views/Annotator.vue @@ -45,6 +45,12 @@ @setcursor="setCursor" ref="keypoint" /> +
@@ -171,6 +177,11 @@ :current-annotation="currentAnnotation" /> +
+ +
@@ -205,6 +216,7 @@ import MagicWandTool from "@/components/annotator/tools/MagicWandTool"; import EraserTool from "@/components/annotator/tools/EraserTool"; import BrushTool from "@/components/annotator/tools/BrushTool"; import KeypointTool from "@/components/annotator/tools/KeypointTool"; +import DEXTRTool from "@/components/annotator/tools/DEXTRTool"; import CopyAnnotationsButton from "@/components/annotator/tools/CopyAnnotationsButton"; import CenterButton from "@/components/annotator/tools/CenterButton"; @@ -224,6 +236,7 @@ import MagicWandPanel from "@/components/annotator/panels/MagicWandPanel"; import BrushPanel from "@/components/annotator/panels/BrushPanel"; import EraserPanel from "@/components/annotator/panels/EraserPanel"; import KeypointPanel from "@/components/annotator/panels/KeypointPanel"; +import DEXTRPanel from "@/components/annotator/panels/DEXTRPanel"; import { mapMutations } from "vuex"; @@ -255,7 +268,9 @@ export default { HideAllButton, ShowAllButton, KeypointPanel, - AnnotateButton + AnnotateButton, + DEXTRTool, + DEXTRPanel }, mixins: [toastrs, shortcuts], props: { diff --git a/docker/api/Dockerfile b/docker/api/Dockerfile index c2aa1727..b3e12ac6 100644 --- a/docker/api/Dockerfile +++ b/docker/api/Dockerfile @@ -5,14 +5,20 @@ WORKDIR /workspace/ # Install python package dependices COPY requirements.txt requirements.txt RUN pip install -r requirements.txt && \ - pip install gunicorn[eventlet]==19.9.0 && \ - pip install pycocotools + pip install gunicorn[eventlet]==19.9.0 && \ + pip install pycocotools # Install maskrcnn -RUN git clone --single-branch --depth 1 https://github.com/matterport/Mask_RCNN.git $TEMP_MRCNN_DIR /tmp/maskrcnn && \ - cd /tmp/maskrcnn && \ - pip install -r requirements.txt && \ - python3 setup.py install +RUN git clone --single-branch --depth 1 https://github.com/matterport/Mask_RCNN.git /tmp/maskrcnn && \ + cd /tmp/maskrcnn && \ + pip install -r requirements.txt && \ + python3 setup.py install + +# Install DEXTR +RUN git clone --single --depth 1 https://github.com/jsbroks/dextr-keras.git /tmp/dextr && \ + cd /tmp/dextr && \ + pip install -r requirements.txt && \ + python setup.py install COPY ./.git /workspace/.git diff --git a/docker/production/Dockerfile b/docker/production/Dockerfile index 21ffc78d..32fc519c 100644 --- a/docker/production/Dockerfile +++ b/docker/production/Dockerfile @@ -37,6 +37,12 @@ RUN git clone --single-branch --depth 1 https://github.com/matterport/Mask_RCNN. pip install -r requirements.txt && \ python3 setup.py install +# Install DEXTR +RUN git clone --single --depth 1 https://github.com/jsbroks/dextr-keras.git /tmp/dextr && \ + cd /tmp/dextr && \ + pip install -r requirements.txt && \ + python setup.py install + COPY ./.git /workspace/.git COPY ./app /workspace/app diff --git a/models/dextr_pascal_sbd.sh b/models/dextr_pascal_sbd.sh new file mode 100644 index 00000000..66a76b9a --- /dev/null +++ b/models/dextr_pascal_sbd.sh @@ -0,0 +1,3 @@ +#!/bin/bash + +wget https://github.com/jsbroks/dextr-keras/releases/download/v1.0.0/dextr_pascal-sbd.h5 \ No newline at end of file