Merge pull request #437 from wearepal/segmentation-component

Segmentation component

Caddyfile

@@ -5,6 +5,8 @@ $LMTADDRESS {
   reverse_proxy web:3000 {
     lb_try_duration 30s
   }
+
+  reverse_proxy /api* landscapes-services:5001
 
   encode gzip
   log

app/javascript/projects/modelling/components/index.ts

@@ -37,6 +37,7 @@ import { IMDComponent } from "./imd_component"
 import { HedgerowComponent } from "./hedgerow_component"
 import { ProjectPermissions } from "../../project_editor"
 import { SoilComponent } from "./soil_component"
+import { SegmentComponent } from "./segment_component"
 
 export interface ProjectProperties {
   extent: Extent
@@ -75,6 +76,7 @@ export function createDefaultComponents(saveMapLayer: SaveMapLayer, saveModel: S
       new ATIComponent(projectProps),
       new DesignationsComponent(projectProps),
       new SoilComponent(projectProps),
+      new SegmentComponent(projectProps),
 
       // Outputs
       new MapLayerComponent(saveMapLayer),

app/javascript/projects/modelling/components/segment_component.ts

@@ -0,0 +1,136 @@
+import { BaseComponent } from "./base_component"
+import { NodeData, WorkerInputs, WorkerOutputs } from 'rete/types/core/data'
+import { Input, Node, Output, Socket } from 'rete'
+import { booleanDataSocket, numericDataSocket } from "../socket_types"
+import { ProjectProperties } from "."
+import { TextControl } from "../controls/text"
+import { BooleanTileGrid, NumericTileGrid } from "../tile_grid"
+import { createXYZ } from "ol/tilegrid"
+import { Point, Polygon } from "ol/geom"
+import { Coordinate } from "ol/coordinate"
+
+async function retrieveSegmentationMasks(prompts: string, threshold: string, projectProps: ProjectProperties) : Promise<any[]>{
+
+    const tileGrid = createXYZ()
+
+    const outputTileRange = tileGrid.getTileRangeForExtentAndZ(projectProps.extent, projectProps.zoom)
+
+    const segs = await fetch("http://landscapes.wearepal.ai/api/v1/segment?" + new URLSearchParams(
+        {
+            labels: prompts,
+            threshold,
+            bbox: projectProps.extent.join(","),
+            layer: "rgb:full_mosaic_3857",
+            height: outputTileRange.getHeight().toString(),
+            width: outputTileRange.getWidth().toString(),
+        }
+    ))
+
+    const segsJson = await segs.json()
+
+    const preds = segsJson.predictions
+
+    const result  = new BooleanTileGrid(
+        projectProps.zoom,
+        outputTileRange.minX,
+        outputTileRange.minY,
+        outputTileRange.getWidth(),
+        outputTileRange.getHeight()
+    )
+
+    const box = new BooleanTileGrid(
+        projectProps.zoom,
+        outputTileRange.minX,
+        outputTileRange.minY,
+        outputTileRange.getWidth(),
+        outputTileRange.getHeight()
+    )
+
+    const confBox = new NumericTileGrid(
+        projectProps.zoom,
+        outputTileRange.minX,
+        outputTileRange.minY,
+        outputTileRange.getWidth(),
+        outputTileRange.getHeight()
+    )
+
+    preds.forEach((pred: any) => {
+
+        const predMask = pred.mask
+
+        predMask.forEach((coord : Coordinate) => {
+
+            const p = new Point(coord)
+
+            const featureTileRange = tileGrid.getTileRangeForExtentAndZ(
+                p.getExtent(),
+                projectProps.zoom
+            )
+
+            result.set(featureTileRange.maxX, featureTileRange.minY, true)
+            confBox.set(featureTileRange.maxX, featureTileRange.minY, pred.score)
+        })
+
+        const predBox = pred.box
+        const predExtent = [Math.min(predBox.xmin, predBox.xmin), Math.min(predBox.ymin, predBox.ymax), Math.max(predBox.xmin, predBox.xmax), Math.max(predBox.ymin, predBox.ymax)]
+
+        const featureTileRange = tileGrid.getTileRangeForExtentAndZ(
+            predExtent,
+            projectProps.zoom
+        )
+
+        box.iterate((x, y) => {
+            if (featureTileRange.containsXY(x, y)) {
+                box.set(x, y, true)
+            }
+        })
+
+    })
+
+    return [result, box, confBox]
+}
+
+export class SegmentComponent extends BaseComponent {
+    projectProps: ProjectProperties
+
+    constructor(projectProps: ProjectProperties) {
+        super("Segmentation Model")
+        this.category = "Inputs"
+        this.projectProps = projectProps
+    }
+
+    async builder(node: Node) {
+
+        if (!('threshold' in node.data)) {
+            node.data.threshold = "0.1"
+        }
+
+        if (!('prompt' in node.data)) {
+            node.data.prompt = "trees"
+        }
+
+        node.addOutput(new Output('mask', 'Segmentation Mask', booleanDataSocket))
+        node.addOutput(new Output('conf', 'Segmentation Mask (Confidence)', numericDataSocket))
+        node.addOutput(new Output('box', 'Segmentation Box', booleanDataSocket))
+
+        node.addControl(new TextControl(this.editor, 'prompt', 'Prompt'))
+        node.addControl(new TextControl(this.editor, 'threshold', 'Threshold'))
+
+    }
+
+    async worker(node: NodeData, inputs: WorkerInputs, outputs: WorkerOutputs, ...args: unknown[]) {
+
+        const editorNode = this.editor?.nodes.find(n => n.id === node.id)
+        if (editorNode === undefined) { return }
+
+        const prompts = node.data.prompt as string
+        const threshold = node.data.threshold as string
+
+        const result = await retrieveSegmentationMasks(prompts, threshold, this.projectProps)
+
+        outputs['mask'] = result[0]
+        outputs['box'] = result[1]
+        outputs['conf'] = result[2]
+
+    }
+}

app/javascript/projects/modelling/controls/text.tsx

@@ -6,29 +6,35 @@ import { EventsTypes } from "rete/types/events"
 interface TextFieldProps {
   getValue: () => string
   setValue: (value: string) => void
+  title: string
 }
-const TextField = ({ getValue, setValue }: TextFieldProps) => {
+const TextField = ({ getValue, setValue, title }: TextFieldProps) => {
   // https://reactjs.org/docs/hooks-faq.html#is-there-something-like-forceupdate
   const [, forceUpdate] = React.useReducer(x => x + 1, 0)
 
-  return <input
-    type="text"
-    className="form-control"
-    value={getValue()}
-    onChange={e => {
-      setValue(e.target.value)
-      forceUpdate()
-    }}
-    onPointerDown={e => e.stopPropagation()}
-    onDoubleClick={e => e.stopPropagation()}
-  />
+  return <>
+    <label>
+      {title}
+    </label>
+    <input
+      type="text"
+      className="form-control"
+      value={getValue()}
+      onChange={e => {
+        setValue(e.target.value)
+        forceUpdate()
+      }}
+      onPointerDown={e => e.stopPropagation()}
+      onDoubleClick={e => e.stopPropagation()}
+    />
+  </>
 }
 
 export class TextControl extends Control {
   props: TextFieldProps
   component: (props: TextFieldProps) => JSX.Element
 
-  constructor(emitter: Emitter<EventsTypes> | null, key: string) {
+  constructor(emitter: Emitter<EventsTypes> | null, key: string, title: string = "") {
     super(key)
 
     const process = debounce(() => emitter?.trigger("process"), 1000)
@@ -41,6 +47,7 @@ export class TextControl extends Control {
         this.putData(key, value)
         process()
       },
+      title: title
     }
     this.component = TextField
   }

docker-compose.yml

@@ -23,6 +23,7 @@ services:
       - caddy_traefik
       - caddy_astute
       - default
+      - landscapes_services
 
   db:
     image: postgres:14
@@ -102,3 +103,5 @@ networks:
     external: true
   caddy_astute:
     external: true
+  landscapes_services:
+    external: true