Polygon tracking feature with XMem tracker

cvat-ui/src/components/annotation-page/standard-workspace/controls-side-bar/tools-control.tsx

@@ -146,6 +146,7 @@ interface State {
  activeLabelID: number | null;
  activeTracker: MLModel | null;
  convertMasksToPolygons: boolean;
+ selectedObjectType: ObjectType;
  trackedShapes: TrackedShape[];
  fetching: boolean;
  pointsReceived: boolean;
@@ -235,6 +236,7 @@ export class ToolsControlComponent extends React.PureComponent<Props, State> {
  super(props);
  this.state = {
  convertMasksToPolygons: false,
+ selectedObjectType: ObjectType.SHAPE,
  activeInteractor: props.interactors.length ? props.interactors[0] : null,
  activeTracker: props.trackers.length ? props.trackers[0] : null,
  activeLabelID: props.labels.length ? props.labels[0].id as number : null,
@@ -593,7 +595,7 @@ export class ToolsControlComponent extends React.PureComponent<Props, State> {
  const portals = !activeTracker ?
  [] :
  states
- .filter((objectState) => objectState.objectType === 'track' && objectState.shapeType === 'rectangle')
+ .filter((objectState) => objectState.objectType === 'track' && (objectState.shapeType === 'rectangle' || objectState.shapeType === 'polygon'))
  .map((objectState: any): React.ReactPortal | null => {
  const { clientID } = objectState;
  const selectorID = `#cvat-objects-sidebar-state-item-${clientID}`;
@@ -822,7 +824,10 @@ export class ToolsControlComponent extends React.PureComponent<Props, State> {
  job: jobInstance.id,
  }) as TrackerResults;
 
- response.shapes = response.shapes.map(trackedRectangleMapper);
+ // If shape type is rectangle, keep same approach
+ if (response.shapes[0].length === 4) {
+ response.shapes = response.shapes.map(trackedRectangleMapper);
+ }
  for (let i = 0; i < trackableObjects.clientIDs.length; i++) {
  const clientID = trackableObjects.clientIDs[i];
  const shape = response.shapes[i];
@@ -859,15 +864,15 @@ export class ToolsControlComponent extends React.PureComponent<Props, State> {
  }
 
  private async constructFromPoints(): Promise<void> {
- const { convertMasksToPolygons } = this.state;
+ const { convertMasksToPolygons, selectedObjectType } = this.state;
  const {
  frame, labels, curZOrder, activeLabelID, createAnnotations,
  } = this.props;
 
  if (convertMasksToPolygons) {
  const object = new core.classes.ObjectState({
  frame,
- objectType: ObjectType.SHAPE,
+ objectType: selectedObjectType,
  source: core.enums.Source.SEMI_AUTO,
  label: labels.find((label) => label.id === activeLabelID as number) as Label,
  shapeType: ShapeType.POLYGON,
@@ -958,6 +963,29 @@ export class ToolsControlComponent extends React.PureComponent<Props, State> {
  );
  }
 
+ private renderObjectTypeBlock(): JSX.Element {
+ const { selectedObjectType } = this.state;
+ const objectTypes = Object.values(ObjectType);
+ objectTypes.splice(objectTypes.indexOf(ObjectType.TAG), 1);
+ return (
+ <Row className='cvat-interactors-setups-container'>
+ <Select
+ value={selectedObjectType}
+ onChange={(value: ObjectType) => {
+ this.setState({ selectedObjectType: value });
+ }}
+ >
+ {objectTypes.map((type) => (
+ <Select.Option key={type} value={type}>
+ {type}
+ </Select.Option>
+ ))}
+ </Select>
+ <Text>Object Type</Text>
+ </Row>
+ );
+ }
+
  private renderLabelBlock(): JSX.Element {
  const { labels } = this.props;
  const { activeLabelID } = this.state;
@@ -1346,6 +1374,7 @@ export class ToolsControlComponent extends React.PureComponent<Props, State> {
  }
 
  private renderPopoverContent(): JSX.Element {
+ const { convertMasksToPolygons } = this.state;
  return (
  <div className='cvat-tools-control-popover-content'>
  <Row justify='start'>
@@ -1358,6 +1387,7 @@ export class ToolsControlComponent extends React.PureComponent<Props, State> {
  <Tabs type='card' tabBarGutter={8}>
  <Tabs.TabPane key='interactors' tab='Interactors'>
  {this.renderMasksConvertingBlock()}
+ {convertMasksToPolygons ? this.renderObjectTypeBlock() : null}
  {this.renderLabelBlock()}
  {this.renderInteractorBlock()}
  </Tabs.TabPane>

serverless/pytorch/omerferhatt/xmem/nuclio/function-gpu.yaml

@@ -0,0 +1,67 @@
+metadata:
+ name: pth-omerferhatt-xmem
+ namespace: cvat
+ annotations:
+ name: XMem
+ type: tracker
+ spec:
+ framework: pytorch
+
+spec:
+ description: Long-Term Object Segmentation with an Atkinson-Shiffrin Memory Model
+ runtime: 'python:3.9'
+ handler: main:handler
+ eventTimeout: 30s
+
+ env:
+ - name: PYTHONPATH
+ value: /opt/nuclio/xmem
+
+ build:
+ image: cvat.pth.omerferhatt.xmem:latest-gpu
+ baseImage: nvidia/cuda:12.1.0-runtime-ubuntu22.04
+
+ directives:
+ preCopy:
+ - kind: RUN
+ value: |-
+ apt update \
+ && apt install -y --no-install-recommends \
+ wget \
+ git \
+ ca-certificates \
+ python-is-python3 \
+ python3 \
+ python3-pip \
+ && rm -rf /var/lib/apt/lists/*
+ - kind: WORKDIR
+ value: /opt/nuclio
+ - kind: RUN
+ value: git clone --branch main https://github.com/omerferhatt/XMem xmem
+ - kind: RUN
+ value: pip install opencv-python-headless jsonpickle
+ - kind: RUN
+ value: |-
+ pip install torch torchvision
+ - kind: RUN
+ value: wget 'https://www.dropbox.com/scl/fi/5m1l747p15qzgq023e0q9/xmem.pth?rlkey=ss2kjaq4qlvvk5juucyvtmrh8&dl=0' -O '/xmem.pth'
+
+ triggers:
+ myHttpTrigger:
+ maxWorkers: 1
+ kind: 'http'
+ workerAvailabilityTimeoutMilliseconds: 10000
+ attributes:
+ # Set value from the calculation of tracking of 100 objects at the same time on a 4k image
+ maxRequestBodySize: 1073741824 # 1GB
+
+ resources:
+ limits:
+ nvidia.com/gpu: 1
+
+ platform:
+ attributes:
+ restartPolicy:
+ name: always
+ maximumRetryCount: 3
+ mountMode: volume

serverless/pytorch/omerferhatt/xmem/nuclio/function.yaml

@@ -0,0 +1,63 @@
+metadata:
+ name: pth-omerferhatt-xmem
+ namespace: cvat
+ annotations:
+ name: XMem
+ type: tracker
+ spec:
+ framework: pytorch
+
+spec:
+ description: Long-Term Object Segmentation with an Atkinson-Shiffrin Memory Model
+ runtime: 'python:3.9'
+ handler: main:handler
+ eventTimeout: 30s
+
+ env:
+ - name: PYTHONPATH
+ value: /opt/nuclio/xmem
+
+ build:
+ image: cvat.pth.omerferhatt.xmem
+ baseImage: ubuntu:22.04
+
+ directives:
+ preCopy:
+ - kind: RUN
+ value: |-
+ apt update \
+ && apt install -y --no-install-recommends \
+ wget \
+ git \
+ ca-certificates \
+ python-is-python3 \
+ python3 \
+ python3-pip \
+ && rm -rf /var/lib/apt/lists/*
+ - kind: WORKDIR
+ value: /opt/nuclio
+ - kind: RUN
+ value: git clone --branch main https://github.com/omerferhatt/XMem xmem
+ - kind: RUN
+ value: pip install opencv-python-headless jsonpickle
+ - kind: RUN
+ value: |-
+ pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu
+ - kind: RUN
+ value: wget 'https://www.dropbox.com/scl/fi/5m1l747p15qzgq023e0q9/xmem.pth?rlkey=ss2kjaq4qlvvk5juucyvtmrh8&dl=0' -O '/xmem.pth'
+
+ triggers:
+ myHttpTrigger:
+ maxWorkers: 1
+ kind: 'http'
+ workerAvailabilityTimeoutMilliseconds: 10000
+ attributes:
+ # Set value from the calculation of tracking of 100 objects at the same time on a 4k image
+ maxRequestBodySize: 1073741824 # 1GB
+
+ platform:
+ attributes:
+ restartPolicy:
+ name: always
+ maximumRetryCount: 3
+ mountMode: volume

serverless/pytorch/omerferhatt/xmem/nuclio/main.py

@@ -0,0 +1,42 @@
+import base64
+import io
+import json
+
+import numpy as np
+from model_handler import ModelHandler
+from PIL import Image
+
+
+def init_context(context):
+ context.logger.info("Init context... 0%")
+ model = ModelHandler()
+ context.user_data.model = model
+ context.logger.info("Init context...100%")
+
+
+def handler(context, event):
+ context.logger.info("Run XMem model")
+ data = event.body
+ buf = io.BytesIO(base64.b64decode(data["image"]))
+ shapes = data.get("shapes")
+ states = data.get("states")
+
+ image = Image.open(buf).convert("RGB")
+ image = np.array(image)[:, :, ::-1].copy()
+ results = {"shapes": [], "states": []}
+ for i, shape in enumerate(shapes):
+ context.logger.info(f"Inference [{i}] started")
+
+ shape, state = context.user_data.model.infer(
+ image, shape, states[i] if i < len(states) else None
+ )
+ results["shapes"].append(shape)
+ results["states"].append(state)
+ context.logger.info(f"Inference [{i}] finised")
+
+ return context.Response(
+ body=json.dumps(results),
+ headers={},
+ content_type="application/json",
+ status_code=200,
+ )