Auto-Select via SAM

CHANGELOG.unreleased.md

@@ -13,6 +13,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 ### Added
 - Added segment groups so that segments can be organized in a hierarchy (similar to skeletons). [#6966](https://github.com/scalableminds/webknossos/pull/6966)
 - In addition to drag and drop, the selected tree(s) in the Skeleton tab can also be moved into another group by right-clicking the target group and selecting "Move selected tree(s) here". [#7005](https://github.com/scalableminds/webknossos/pull/7005)
+- Added a machine-learning based quick select mode. Activate it via the "AI" button in the toolbar after selecting the quick-select tool. [#7051](https://github.com/scalableminds/webknossos/pull/7051)
 - Added support for remote datasets encoded with [brotli](https://datatracker.ietf.org/doc/html/rfc7932). [#7041](https://github.com/scalableminds/webknossos/pull/7041)
 - Teams can be edited more straight-forwardly in a popup in the team edit page. [#7043](https://github.com/scalableminds/webknossos/pull/7043)
 - Annotations with Editable Mappings (a.k.a Supervoxel Proofreading) can now be merged. [#7026](https://github.com/scalableminds/webknossos/pull/7026)

app/controllers/DataSetController.scala

@@ -20,7 +20,7 @@ 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 utils.ObjectId
+import utils.{ObjectId, WkConf}
 
 import javax.inject.Inject
 import scala.collection.mutable.ListBuffer
@@ -42,17 +42,27 @@ object DatasetUpdateParameters extends TristateOptionJsonHelper {
     Json.configured(tristateOptionParsing).format[DatasetUpdateParameters]
 }
 
+case class SegmentAnythingEmbeddingParameters(
+    mag: Vec3Int,
+    boundingBox: BoundingBox
+)
+
+object SegmentAnythingEmbeddingParameters {
+  implicit val jsonFormat: Format[SegmentAnythingEmbeddingParameters] = Json.format[SegmentAnythingEmbeddingParameters]
+}
+
 @Api
 class DataSetController @Inject()(userService: UserService,
                                   userDAO: UserDAO,
                                   dataSetService: DataSetService,
-                                  dataSetDataLayerDAO: DataSetDataLayerDAO,
                                   dataStoreDAO: DataStoreDAO,
                                   dataSetLastUsedTimesDAO: DataSetLastUsedTimesDAO,
                                   organizationDAO: OrganizationDAO,
                                   teamDAO: TeamDAO,
+                                  wKRemoteSegmentAnythingClient: WKRemoteSegmentAnythingClient,
                                   teamService: TeamService,
                                   dataSetDAO: DataSetDAO,
+                                  conf: WkConf,
                                   analyticsService: AnalyticsService,
                                   mailchimpClient: MailchimpClient,
                                   exploreRemoteLayerService: ExploreRemoteLayerService,
@@ -520,4 +530,32 @@ Expects:
       case _             => Messages("dataSet.notFoundConsiderLogin", dataSetName)
     }
 
+  @ApiOperation(hidden = true, value = "")
+  def segmentAnythingEmbedding(organizationName: String,
+                               dataSetName: String,
+                               dataLayerName: String): Action[SegmentAnythingEmbeddingParameters] =
+    sil.SecuredAction.async(validateJson[SegmentAnythingEmbeddingParameters]) { implicit request =>
+      log() {
+        for {
+          _ <- bool2Fox(conf.Features.segmentAnythingEnabled) ?~> "segmentAnything.notEnabled"
+          _ <- bool2Fox(conf.SegmentAnything.uri.nonEmpty) ?~> "segmentAnything.noUri"
+          dataset <- dataSetDAO.findOneByNameAndOrganizationName(dataSetName, organizationName) ?~> notFoundMessage(
+            dataSetName) ~> NOT_FOUND
+          dataSource <- dataSetService.dataSourceFor(dataset) ?~> "dataSource.notFound" ~> NOT_FOUND
+          usableDataSource <- dataSource.toUsable ?~> "dataSet.notImported"
+          dataLayer <- usableDataSource.dataLayers.find(_.name == dataLayerName) ?~> "dataSet.noLayers"
+          datastoreClient <- dataSetService.clientFor(dataset)(GlobalAccessContext)
+          targetMagBbox: BoundingBox = request.body.boundingBox / request.body.mag
+          _ <- bool2Fox(targetMagBbox.dimensions.sorted == Vec3Int(1, 1024, 1024)) ?~> s"Target-mag bbox must be sized 1024×1024×1 (or transposed), got ${targetMagBbox.dimensions}"
+          data <- datastoreClient.getLayerData(organizationName,
+                                               dataset,
+                                               dataLayer.name,
+                                               request.body.boundingBox,
+                                               request.body.mag) ?~> "segmentAnything.getData.failed"
+          _ = logger.debug(s"Sending ${data.length} bytes to SAM server, element class is ${dataLayer.elementClass}...")
+          embedding <- wKRemoteSegmentAnythingClient.getEmbedding(data, dataLayer.elementClass) ?~> "segmentAnything.getEmbedding.failed"
+        } yield Ok(embedding)
+      }
+    }
+
 }

app/models/binary/DataSetService.scala

@@ -16,7 +16,7 @@ import com.scalableminds.webknossos.datastore.models.datasource.{
 import com.scalableminds.webknossos.datastore.rpc.RPC
 import com.scalableminds.webknossos.datastore.storage.TemporaryStore
 import com.typesafe.scalalogging.LazyLogging
-import models.folder.{FolderDAO, FolderService}
+import models.folder.FolderDAO
 
 import javax.inject.Inject
 import models.job.WorkerDAO
@@ -41,7 +41,6 @@ class DataSetService @Inject()(organizationDAO: OrganizationDAO,
                                dataStoreService: DataStoreService,
                                teamService: TeamService,
                                userService: UserService,
-                               folderService: FolderService,
                                val thumbnailCache: TemporaryStore[String, Array[Byte]],
                                rpc: RPC,
                                conf: WkConf)(implicit ec: ExecutionContext)
@@ -66,7 +65,7 @@ class DataSetService @Inject()(organizationDAO: OrganizationDAO,
     createDataSet(dataStore, organizationName, unreportedDatasource)
   }
 
-  def createDataSet(
+  private def createDataSet(
       dataStore: DataStore,
       owningOrganization: String,
       dataSource: InboxDataSource,
@@ -256,7 +255,7 @@ class DataSetService @Inject()(organizationDAO: OrganizationDAO,
       dataStore <- dataStoreFor(dataSet)
     } yield new WKRemoteDataStoreClient(dataStore, rpc)
 
-  def lastUsedTimeFor(_dataSet: ObjectId, userOpt: Option[User]): Fox[Instant] =
+  private def lastUsedTimeFor(_dataSet: ObjectId, userOpt: Option[User]): Fox[Instant] =
     userOpt match {
       case Some(user) =>
         (for {

app/models/binary/WKRemoteDataStoreClient.scala

@@ -1,6 +1,6 @@
 package models.binary
 
-import com.scalableminds.util.geometry.Vec3Int
+import com.scalableminds.util.geometry.{BoundingBox, Vec3Int}
 import com.scalableminds.util.tools.Fox
 import com.scalableminds.webknossos.datastore.rpc.RPC
 import com.scalableminds.webknossos.datastore.services.DirectoryStorageReport
@@ -29,6 +29,26 @@ class WKRemoteDataStoreClient(dataStore: DataStore, rpc: RPC) extends LazyLoggin
       .getWithBytesResponse
   }
 
+  def getLayerData(organizationName: String,
+                   dataset: DataSet,
+                   layerName: String,
+                   mag1BoundingBox: BoundingBox,
+                   mag: Vec3Int): Fox[Array[Byte]] = {
+    val targetMagBoundingBox = mag1BoundingBox / mag
+    logger.debug(s"Fetching raw data. Mag $mag, mag1 bbox: $mag1BoundingBox, target-mag bbox: $targetMagBoundingBox")
+    rpc(
+      s"${dataStore.url}/data/datasets/${urlEncode(organizationName)}/${dataset.urlEncodedName}/layers/$layerName/data")
+      .addQueryString("token" -> RpcTokenHolder.webKnossosToken)
+      .addQueryString("mag" -> mag.toMagLiteral())
+      .addQueryString("x" -> mag1BoundingBox.topLeft.x.toString)
+      .addQueryString("y" -> mag1BoundingBox.topLeft.y.toString)
+      .addQueryString("z" -> mag1BoundingBox.topLeft.z.toString)
+      .addQueryString("width" -> targetMagBoundingBox.width.toString)
+      .addQueryString("height" -> targetMagBoundingBox.height.toString)
+      .addQueryString("depth" -> targetMagBoundingBox.depth.toString)
+      .getWithBytesResponse
+  }
+
   def findPositionWithData(organizationName: String, dataSet: DataSet, dataLayerName: String): Fox[JsObject] =
     rpc(
       s"${dataStore.url}/data/datasets/${urlEncode(organizationName)}/${dataSet.urlEncodedName}/layers/$dataLayerName/findData")

app/models/binary/WKRemoteSegmentAnythingClient.scala

@@ -0,0 +1,15 @@
+package models.binary
+
+import com.scalableminds.util.tools.Fox
+import com.scalableminds.webknossos.datastore.rpc.RPC
+import com.scalableminds.webknossos.datastore.models.datasource.ElementClass
+import utils.WkConf
+
+import javax.inject.Inject
+
+class WKRemoteSegmentAnythingClient @Inject()(rpc: RPC, conf: WkConf) {
+  def getEmbedding(imageData: Array[Byte], elementClass: ElementClass.Value): Fox[Array[Byte]] =
+    rpc(s"${conf.SegmentAnything.uri}/predictions/sam_vit_l")
+      .addQueryString("elementClass" -> elementClass.toString)
+      .postBytesWithBytesResponse(imageData)
+}

app/utils/WkConf.scala

@@ -104,6 +104,7 @@ class WkConf @Inject()(configuration: Configuration) extends ConfigReader with L
     val exportTiffMaxVolumeMVx: Long = get[Long]("features.exportTiffMaxVolumeMVx")
     val exportTiffMaxEdgeLengthVx: Long = get[Long]("features.exportTiffMaxEdgeLengthVx")
     val openIdConnectEnabled: Boolean = get[Boolean]("features.openIdConnectEnabled")
+    val segmentAnythingEnabled: Boolean = get[Boolean]("features.segmentAnythingEnabled")
   }
 
   object Datastore {
@@ -230,6 +231,10 @@ class WkConf @Inject()(configuration: Configuration) extends ConfigReader with L
     val children = List(Loki)
   }
 
+  object SegmentAnything {
+    val uri: String = get[String]("segmentAnything.uri")
+  }
+
   val children =
     List(
       Http,
@@ -246,7 +251,8 @@ class WkConf @Inject()(configuration: Configuration) extends ConfigReader with L
       GoogleAnalytics,
       BackendAnalytics,
       Slick,
-      Voxelytics
+      Voxelytics,
+      SegmentAnything
     )
 
 }

conf/application.conf

@@ -121,6 +121,7 @@ features {
   # The Only valid item value is currently "ConnectomeView":
   optInTabs = []
   openIdConnectEnabled = false
+  segmentAnythingEnabled = false
 }
 
 # Serve annotations. Only active if the corresponding play module is enabled
@@ -289,10 +290,15 @@ voxelytics {
   }
 }
 
+segmentAnything {
+  uri = "http://localhost:8080"
+}
+
 # Avoid creation of a pid file
 pidfile.path = "/dev/null"
 
 
-# # uncomment these lines for faster restart during local backend development (but beware the then-missing features):
+# # Uncomment these lines for faster restart during local backend development (but beware the then-missing features):
+# # Uncommenting these lines also means that a DB has to be set up manually for a dev deployment.
 # slick.checkSchemaOnStartup = false
 # play.modules.disabled += "play.modules.swagger.SwaggerModule"