Allow click prompt for SAM 2

CHANGELOG.unreleased.md

@@ -26,6 +26,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - In the proofreading mode, you can enable/disable that only the active segment and the hovered segment are rendered. [#7654](https://github.com/scalableminds/webknossos/pull/7654)
 - Upgraded s3 client for improved performance when loading remote datasets. [#7936](https://github.com/scalableminds/webknossos/pull/7936)
 - The AI-based Quick Select can now be run on multiple sections at once. This can be configured in the tool settings. Also, the underlying model now uses Segment Anything 2. [#7965](https://github.com/scalableminds/webknossos/pull/7965)
+- The AI-based Quick Select can now be triggered with a single click. Drawing a rectangle is still supported. [#7993](https://github.com/scalableminds/webknossos/pull/7993)
 - To improve performance, only the visible bounding boxes are rendered in the bounding box tab (so-called virtualization). [#7974](https://github.com/scalableminds/webknossos/pull/7974)
 - Added support for reading zstd-compressed zarr2 datasets [#7964](https://github.com/scalableminds/webknossos/pull/7964)
 - The alignment job is in a separate tab of the "AI Tools" now. The "Align Sections" AI job now supports including manually created matches between adjacent section given as skeletons. [#7967](https://github.com/scalableminds/webknossos/pull/7967)  

app/controllers/DatasetController.scala

@@ -1,33 +1,34 @@
 package controllers
 
-import play.silhouette.api.Silhouette
 import com.scalableminds.util.accesscontext.{DBAccessContext, GlobalAccessContext}
+import com.scalableminds.util.enumeration.ExtendedEnumeration
 import com.scalableminds.util.geometry.{BoundingBox, Vec3Int}
 import com.scalableminds.util.time.Instant
 import com.scalableminds.util.tools.{Fox, TristateOptionJsonHelper}
+import com.scalableminds.webknossos.datastore.models.AdditionalCoordinate
 import com.scalableminds.webknossos.datastore.models.datasource.ElementClass
+import mail.{MailchimpClient, MailchimpTag}
 import models.analytics.{AnalyticsService, ChangeDatasetSettingsEvent, OpenDatasetEvent}
 import models.dataset._
 import models.dataset.explore.{
   ExploreAndAddRemoteDatasetParameters,
   WKExploreRemoteLayerParameters,
   WKExploreRemoteLayerService
 }
+import models.folder.FolderService
 import models.organization.OrganizationDAO
 import models.team.{TeamDAO, TeamService}
 import models.user.{User, UserDAO, UserService}
 import play.api.i18n.{Messages, MessagesProvider}
 import play.api.libs.functional.syntax._
 import play.api.libs.json._
 import play.api.mvc.{Action, AnyContent, PlayBodyParsers}
+import play.silhouette.api.Silhouette
+import security.{URLSharing, WkEnv}
 import utils.{ObjectId, WkConf}
 
 import javax.inject.Inject
 import scala.concurrent.{ExecutionContext, Future}
-import com.scalableminds.webknossos.datastore.models.AdditionalCoordinate
-import mail.{MailchimpClient, MailchimpTag}
-import models.folder.FolderService
-import security.{URLSharing, WkEnv}
 
 case class DatasetUpdateParameters(
     description: Option[Option[String]] = Some(None),
@@ -43,14 +44,24 @@ object DatasetUpdateParameters extends TristateOptionJsonHelper {
     Json.configured(tristateOptionParsing).format[DatasetUpdateParameters]
 }
 
+object SAMInteractionType extends ExtendedEnumeration {
+  type SAMInteractionType = Value
+  val BOUNDING_BOX, POINT = Value
+}
+
 case class SegmentAnythingMaskParameters(
     mag: Vec3Int,
     surroundingBoundingBox: BoundingBox, // in mag1 (when converted to target mag, size must be 1024×1024×depth with depth <= 12)
     additionalCoordinates: Option[Seq[AdditionalCoordinate]] = None,
-    selectionTopLeftX: Int, // in target-mag, relative to paddedBoundingBox topleft
-    selectionTopLeftY: Int,
-    selectionBottomRightX: Int,
-    selectionBottomRightY: Int,
+    interactionType: SAMInteractionType.SAMInteractionType,
+    // selectionTopLeft and selectionBottomRight are required as input in case of bounding box interaction type.
+    // Else pointX and pointY are required.
+    selectionTopLeftX: Option[Int], // in target-mag, relative to paddedBoundingBox topleft
+    selectionTopLeftY: Option[Int],
+    selectionBottomRightX: Option[Int],
+    selectionBottomRightY: Option[Int],
+    pointX: Option[Int], // in target-mag, relative to paddedBoundingBox topleft
+    pointY: Option[Int],
 )
 
 object SegmentAnythingMaskParameters {
@@ -406,9 +417,14 @@ class DatasetController @Inject()(userService: UserService,
           dataLayer <- usableDataSource.dataLayers.find(_.name == dataLayerName) ?~> "dataset.noLayers"
           datastoreClient <- datasetService.clientFor(dataset)(GlobalAccessContext)
           targetMagSelectedBbox: BoundingBox = request.body.surroundingBoundingBox / request.body.mag
-          _ <- bool2Fox(targetMagSelectedBbox.size.sorted.z == 1024) ?~> s"Target-mag selected bbox must be sized 1024×1024×depth (or transposed), got ${targetMagSelectedBbox.size}"
-          _ <- bool2Fox(targetMagSelectedBbox.size.sorted.y == 1024) ?~> s"Target-mag selected bbox must be sized 1024×1024×depth (or transposed), got ${targetMagSelectedBbox.size}"
+          _ <- bool2Fox(targetMagSelectedBbox.size.sorted.z <= 1024 && targetMagSelectedBbox.size.sorted.y <= 1024) ?~> s"Target-mag selected bbox must be smaller than 1024×1024×depth (or transposed), got ${targetMagSelectedBbox.size}"
           _ <- bool2Fox(targetMagSelectedBbox.size.sorted.x <= 12) ?~> s"Target-mag selected bbox depth must be at most 12"
+          _ <- bool2Fox(targetMagSelectedBbox.size.sorted.z == targetMagSelectedBbox.size.sorted.y) ?~> s"Target-mag selected bbox must equally sized long edges, got ${targetMagSelectedBbox.size}"
+          _ <- Fox.runIf(request.body.interactionType == SAMInteractionType.BOUNDING_BOX)(
+            bool2Fox(request.body.selectionTopLeftX.isDefined &&
+              request.body.selectionTopLeftY.isDefined && request.body.selectionBottomRightX.isDefined && request.body.selectionBottomRightY.isDefined)) ?~> "Missing selectionTopLeft and selectionBottomRight parameters for bounding box interaction."
+          _ <- Fox.runIf(request.body.interactionType == SAMInteractionType.POINT)(bool2Fox(
+            request.body.pointX.isDefined && request.body.pointY.isDefined)) ?~> "Missing pointX and pointY parameters for point interaction."
           beforeDataLoading = Instant.now
           data <- datastoreClient.getLayerData(
             organizationName,
@@ -427,10 +443,13 @@ class DatasetController @Inject()(userService: UserService,
           mask <- wKRemoteSegmentAnythingClient.getMask(
             data,
             dataLayer.elementClass,
+            request.body.interactionType,
             request.body.selectionTopLeftX,
             request.body.selectionTopLeftY,
             request.body.selectionBottomRightX,
             request.body.selectionBottomRightY,
+            request.body.pointX,
+            request.body.pointY,
             targetMagSelectedBbox.size,
             intensityMin,
             intensityMax

app/models/dataset/WKRemoteSegmentAnythingClient.scala

@@ -4,6 +4,8 @@ import com.scalableminds.util.geometry.Vec3Int
 import com.scalableminds.util.tools.Fox
 import com.scalableminds.webknossos.datastore.rpc.RPC
 import com.scalableminds.webknossos.datastore.models.datasource.ElementClass
+import controllers.SAMInteractionType
+import controllers.SAMInteractionType.SAMInteractionType
 import utils.WkConf
 
 import java.nio.{ByteBuffer, ByteOrder}
@@ -14,23 +16,36 @@ class WKRemoteSegmentAnythingClient @Inject()(rpc: RPC, conf: WkConf) {
   def getMask(
       imageData: Array[Byte],
       elementClass: ElementClass.Value,
-      selectionTopLeftX: Int,
-      selectionTopLeftY: Int,
-      selectionBottomRightX: Int,
-      selectionBottomRightY: Int,
-      dataShape: Vec3Int, // two of the axes will be 1024, the other is the "depth". Axis order varies depending on viewport
+      samInteractionType: SAMInteractionType,
+      selectionTopLeftXOpt: Option[Int],
+      selectionTopLeftYOpt: Option[Int],
+      selectionBottomRightXOpt: Option[Int],
+      selectionBottomRightYOpt: Option[Int],
+      pointXOpt: Option[Int],
+      pointYOpt: Option[Int],
+      dataShape: Vec3Int, // two of the axes will be at most 1024, the other is the "depth". Axis order varies depending on viewport
       intensityMin: Option[Float],
       intensityMax: Option[Float]): Fox[Array[Byte]] = {
-    val metadataLengthInBytes = 1 + 1 + 4 + 4 + 4 + 4 + 4 + 4 + 4 + 4 + 4
+    val interactionInputLengthInBytes = 1 + (
+      if (samInteractionType == SAMInteractionType.BOUNDING_BOX) 4 + 4 + 4 + 4 else 4 + 4
+    )
+    val metadataLengthInBytes = 1 + 1 + 4 + 4 + interactionInputLengthInBytes + 4 + 4 + 4
     val buffer = ByteBuffer.allocate(metadataLengthInBytes + imageData.length).order(ByteOrder.LITTLE_ENDIAN)
     buffer.put(ElementClass.encodeAsByte(elementClass))
     buffer.put(if (intensityMin.isDefined && intensityMax.isDefined) 1.toByte else 0.toByte)
     buffer.putFloat(intensityMin.getOrElse(0.0f))
     buffer.putFloat(intensityMax.getOrElse(0.0f))
-    buffer.putInt(selectionTopLeftX)
-    buffer.putInt(selectionTopLeftY)
-    buffer.putInt(selectionBottomRightX)
-    buffer.putInt(selectionBottomRightY)
+    if (samInteractionType == SAMInteractionType.BOUNDING_BOX) {
+      buffer.put(0.toByte) // Set bounding box interaction
+      buffer.putInt(selectionTopLeftXOpt.getOrElse(0))
+      buffer.putInt(selectionTopLeftYOpt.getOrElse(0))
+      buffer.putInt(selectionBottomRightXOpt.getOrElse(0))
+      buffer.putInt(selectionBottomRightYOpt.getOrElse(0))
+    } else { // Else only point interaction is possible
+      buffer.put(1.toByte) // Set point interaction
+      buffer.putInt(pointXOpt.getOrElse(0))
+      buffer.putInt(pointYOpt.getOrElse(0))
+    }
     buffer.putInt(dataShape.x)
     buffer.putInt(dataShape.y)
     buffer.putInt(dataShape.z)

docs/tutorial_volume_annotation.md

@@ -25,11 +25,11 @@ Create a new segment by selecting a new segment ID. Use the trace tool to draw m
 ![type:video](https://static.webknossos.org/assets/docs/tutorial-volume-annotation/03_new_segments_lasso.mp4){: autoplay loop muted}
 
 ## Smart Annotation Tools
-Now, let’s explore the smart tools of WEBKNOSSOS. Choose the quick-select tool and turn on the AI option. Draw a rectangle around the cell you want to segment. Create and annotate new segments by pressing the keyboard shortcut “C”.
+Now, let’s explore the smart tools of WEBKNOSSOS. Choose the quick-select tool and turn on the AI option. Click on a cell you want to segment or draw a rectangle around it. Create and annotate new segments by pressing the keyboard shortcut “C”.
 
 ![type:video](https://static.webknossos.org/assets/docs/tutorial-volume-annotation/04_new_AI_quick_select.mp4){: autoplay loop muted}
 
-Quick tip: When dealing with long and complex shapes, try drawing multiple rectangles on different areas. The annotations will merge and your complex shape will be segmented in no time.
+Quick tip: The current zoom is taken into account when using the quick select feature. Therefore, smaller structures can be segmented better when being zoomed in further. And another tip: When dealing with long and complex shapes, try drawing multiple rectangles on different areas. The annotations will merge and your complex shape will be segmented in no time.
 
 ![type:video](https://static.webknossos.org/assets/docs/tutorial-volume-annotation/04_new_tip_long_cell.mp4){: autoplay loop muted}
 

frontend/javascripts/admin/admin_rest_api.ts

@@ -2338,8 +2338,19 @@ export async function getSamMask(
   layerName: string,
   mag: Vector3,
   surroundingBoxMag1: BoundingBox, // in mag 1
-  selectionTopLeft: Vector2, // in target mag
-  selectionBottomRight: Vector2, // in target mag
+  prompt:
+    | {
+        type: "BOUNDING_BOX"; // relative to topleft
+        selectionTopLeftX: number; // int, in target mag
+        selectionTopLeftY: number; // int, in target mag
+        selectionBottomRightX: number; // int, in target mag
+        selectionBottomRightY: number; // int, in target mag
+      }
+    | {
+        type: "POINT";
+        pointX: number; // int, relative to topleft
+        pointY: number; // int, relative to topleft
+      },
   additionalCoordinates: AdditionalCoordinate[],
   intensityRange?: Vector2 | null,
 ): Promise<Uint8Array> {
@@ -2349,17 +2360,17 @@ export async function getSamMask(
     params.append("intensityMax", `${intensityRange[1]}`);
   }
 
+  const { type: interactionType, ...promptWithoutType } = prompt;
+
   const buffer = await Request.sendJSONReceiveArraybuffer(
     `/api/datasets/${dataset.owningOrganization}/${dataset.name}/layers/${layerName}/segmentAnythingMask?${params}`,
     {
       data: {
         mag,
         surroundingBoundingBox: surroundingBoxMag1.asServerBoundingBox(),
         additionalCoordinates,
-        selectionTopLeftX: selectionTopLeft[0],
-        selectionTopLeftY: selectionTopLeft[1],
-        selectionBottomRightX: selectionBottomRight[0],
-        selectionBottomRightY: selectionBottomRight[1],
+        interactionType,
+        ...promptWithoutType,
       },
       showErrorToast: false,
     },

frontend/javascripts/app.ts

@@ -1,4 +1,3 @@
-import window from "libs/window";
 import { createNanoEvents } from "nanoevents";
 
 class OxalisApplication {
@@ -13,6 +12,5 @@ class OxalisApplication {
 }
 
 const app = new OxalisApplication();
-(window as any).app = app;
 
 export default app;

frontend/javascripts/oxalis/api/cross_origin_api.ts

@@ -146,7 +146,6 @@ function CrossOriginApi() {
   }, []);
   // biome-ignore lint/correctness/useExhaustiveDependencies: Rerun each time window.webknossos changes.
   useEffect(() => {
-    // @ts-expect-error ts-migrate(2339) FIXME: Property 'webknossos' does not exist on type 'Wind... Remove this comment to see the full error message
     if (window.webknossos && window.parent) {
       // @ts-expect-error ts-migrate(2339) FIXME: Property 'webknossos' does not exist on type 'Wind... Remove this comment to see the full error message
       window.webknossos.apiReady().then(() => {
@@ -158,7 +157,6 @@ function CrossOriginApi() {
         );
       });
     }
-    // @ts-expect-error ts-migrate(2339) FIXME: Property 'webknossos' does not exist on type 'Wind... Remove this comment to see the full error message
   }, [window.webknossos]);
   return null;
 }

frontend/javascripts/oxalis/api/wk_dev.ts

@@ -6,6 +6,26 @@ import type ApiLoader from "./api_loader";
 import { type ApiInterface } from "./api_latest";
 import showFpsMeter from "libs/fps_meter";
 
+// Can be accessed via window.webknossos.DEV.flags. Only use this
+// for debugging or one off scripts.
+export const WkDevFlags = {
+  sam: {
+    useLocalMask: true,
+  },
+  bucketDebugging: {
+    // For visualizing buckets which are passed to the GPU
+    visualizeBucketsOnGPU: false,
+    // For visualizing buckets which are prefetched
+    visualizePrefetchedBuckets: false,
+    // For enforcing fallback rendering. enforcedZoomDiff == 2, means
+    // that buckets of currentZoomStep + 2 are rendered.
+    enforcedZoomDiff: undefined,
+  },
+  meshing: {
+    marchingCubeSizeInTargetMag: [64, 64, 64] as Vector3,
+  },
+};
+
 export default class WkDev {
   /*
    * This class is only exposed to simplify debugging via the command line.
@@ -15,6 +35,8 @@ export default class WkDev {
   apiLoader: ApiLoader;
   _api!: ApiInterface;
 
+  flags = WkDevFlags;
+
   constructor(apiLoader: ApiLoader) {
     this.apiLoader = apiLoader;
     this.apiLoader.apiReady().then(async (api) => {

frontend/javascripts/oxalis/controller.tsx

@@ -171,7 +171,6 @@ class Controller extends React.PureComponent<PropsWithRouter, State> {
     initializeSceneController();
     this.initKeyboard();
     this.initTaskScript();
-    // @ts-expect-error ts-migrate(2339) FIXME: Property 'webknossos' does not exist on type '(Win... Remove this comment to see the full error message
     window.webknossos = new ApiLoader(Model);
     app.vent.emit("webknossos:ready");
     Store.dispatch(wkReadyAction());

frontend/javascripts/oxalis/controller/combinations/tool_controls.ts

@@ -13,6 +13,7 @@ import {
   handleClickSegment,
 } from "oxalis/controller/combinations/segmentation_handlers";
 import {
+  computeQuickSelectForPointAction,
   computeQuickSelectForRectAction,
   confirmQuickSelectAction,
   hideBrushAction,
@@ -53,6 +54,7 @@ import {
   setActiveUserBoundingBoxId,
 } from "oxalis/model/actions/ui_actions";
 import ArbitraryView from "oxalis/view/arbitrary_view";
+import features from "features";
 
 export type ActionDescriptor = {
   leftClick?: string;
@@ -726,6 +728,19 @@ export class QuickSelectTool {
 
         quickSelectGeometry.setCoordinates(startPos, currentPos);
       },
+      leftClick: (pos: Point2, _plane: OrthoView, _event: MouseEvent, _isTouch: boolean) => {
+        const state = Store.getState();
+        const clickedPos = V3.floor(calculateGlobalPos(state, pos));
+        isDragging = false;
+
+        const quickSelectConfig = state.userConfiguration.quickSelect;
+        const isAISelectAvailable = features().segmentAnythingEnabled;
+        const isQuickSelectHeuristic = quickSelectConfig.useHeuristic || !isAISelectAvailable;
+
+        if (!isQuickSelectHeuristic) {
+          Store.dispatch(computeQuickSelectForPointAction(clickedPos, quickSelectGeometry));
+        }
+      },
       rightClick: (pos: Point2, plane: OrthoView, event: MouseEvent, isTouch: boolean) => {
         SkeletonHandlers.handleOpenContextMenu(planeView, pos, plane, isTouch, event);
       },

frontend/javascripts/oxalis/model/actions/volumetracing_actions.ts

@@ -47,6 +47,7 @@ export type SetHasEditableMappingAction = ReturnType<typeof setHasEditableMappin
 export type SetMappingIsLockedAction = ReturnType<typeof setMappingIsLockedAction>;
 
 export type ComputeQuickSelectForRectAction = ReturnType<typeof computeQuickSelectForRectAction>;
+export type ComputeQuickSelectForPointAction = ReturnType<typeof computeQuickSelectForPointAction>;
 export type FineTuneQuickSelectAction = ReturnType<typeof fineTuneQuickSelectAction>;
 export type CancelQuickSelectAction = ReturnType<typeof cancelQuickSelectAction>;
 export type ConfirmQuickSelectAction = ReturnType<typeof confirmQuickSelectAction>;
@@ -94,6 +95,7 @@ export type VolumeTracingAction =
   | SetMappingIsLockedAction
   | InitializeEditableMappingAction
   | ComputeQuickSelectForRectAction
+  | ComputeQuickSelectForPointAction
   | FineTuneQuickSelectAction
   | CancelQuickSelectAction
   | ConfirmQuickSelectAction
@@ -397,6 +399,16 @@ export const computeQuickSelectForRectAction = (
     quickSelectGeometry,
   }) as const;
 
+export const computeQuickSelectForPointAction = (
+  position: Vector3,
+  quickSelectGeometry: QuickSelectGeometry,
+) =>
+  ({
+    type: "COMPUTE_QUICK_SELECT_FOR_POINT",
+    position,
+    quickSelectGeometry,
+  }) as const;
+
 export const fineTuneQuickSelectAction = (
   segmentMode: "dark" | "light",
   threshold: number,

frontend/javascripts/oxalis/model/bucket_data_handling/bucket.ts

@@ -30,18 +30,6 @@ export type BucketDataArray =
   | Uint32Array
   | Float32Array
   | BigUint64Array;
-export const bucketDebuggingFlags = {
-  // For visualizing buckets which are passed to the GPU
-  visualizeBucketsOnGPU: false,
-  // For visualizing buckets which are prefetched
-  visualizePrefetchedBuckets: false,
-  // For enforcing fallback rendering. enforcedZoomDiff == 2, means
-  // that buckets of currentZoomStep + 2 are rendered.
-  enforcedZoomDiff: undefined,
-};
-// Exposing this variable allows debugging on deployed systems
-// @ts-ignore
-window.bucketDebuggingFlags = bucketDebuggingFlags;
 
 const WARNING_THROTTLE_THRESHOLD = 10000;