Some fixes in CLIP feature

vitrivr · Nov 18, 2023 · 0a6d3e0 · 0a6d3e0
1 parent 72d55b3
commit 0a6d3e0
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 6 deletions.
diff --git a/...in/kotlin/org/vitrivr/engine/base/features/external/implementations/clip/CLIPRetriever.kt b/...in/kotlin/org/vitrivr/engine/base/features/external/implementations/clip/CLIPRetriever.kt
@@ -3,12 +3,14 @@ package org.vitrivr.engine.base.features.external.implementations.clip
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.flow
+import org.vitrivr.engine.base.features.averagecolor.AverageColorRetriever
 import org.vitrivr.engine.core.context.QueryContext
 import org.vitrivr.engine.core.features.AbstractRetriever
 import org.vitrivr.engine.core.model.content.element.ContentElement
 import org.vitrivr.engine.core.model.content.element.ImageContent
 import org.vitrivr.engine.core.model.descriptor.vector.FloatVectorDescriptor
 import org.vitrivr.engine.core.model.metamodel.Schema
+import org.vitrivr.engine.core.model.query.proximity.Distance
 import org.vitrivr.engine.core.model.query.proximity.ProximityQuery
 import org.vitrivr.engine.core.model.retrievable.Retrieved
 
@@ -25,11 +27,43 @@ import org.vitrivr.engine.core.model.retrievable.Retrieved
  * @author Rahel Arnold
  * @version 1.0.0
  */
-class CLIPRetriever(field: Schema.Field<ContentElement<*>, FloatVectorDescriptor>, query: FloatVectorDescriptor, context: QueryContext) : AbstractRetriever<ContentElement<*>, FloatVectorDescriptor>(field, query, context) {
-    override fun toFlow(scope: CoroutineScope): Flow<Retrieved> = flow {
-        val query = ProximityQuery(this@CLIPRetriever.query) /* TODO: Not sure, if the default setting should be used here. */
-        this@CLIPRetriever.reader.getAll(query).forEach {
-            emit(it)
+class CLIPRetriever(
+    field: Schema.Field<ContentElement<*>, FloatVectorDescriptor>,
+    query: FloatVectorDescriptor,
+    context: QueryContext
+) : AbstractRetriever<ContentElement<*>, FloatVectorDescriptor>(field, query, context) {
+
+    companion object {
+        fun scoringFunction(retrieved: Retrieved.RetrievedWithDistance): Float = 1f - retrieved.distance
+
+    }
+
+    override fun toFlow(scope: CoroutineScope): Flow<Retrieved> {
+
+        val k = context.getProperty(field.fieldName, "limit")?.toIntOrNull() ?: 1000 //TODO get limit
+        val returnDescriptor =
+            context.getProperty(field.fieldName, "returnDescriptor")?.toBooleanStrictOrNull() ?: false
+
+
+        val reader = field.getReader()
+
+        val query = ProximityQuery(
+            descriptor = this@CLIPRetriever.query,
+            k = k,
+            distance = Distance.COSINE,
+            withDescriptor = returnDescriptor
+        )
+
+        return flow {
+            reader.getAll(query).forEach {
+                emit(
+                    if (it is Retrieved.RetrievedWithDistance) {
+                        Retrieved.PlusScore(it, AverageColorRetriever.scoringFunction(it))
+                    } else {
+                        it
+                    }
+                )
+            }
         }
     }
 }
diff --git a/vitrivr-engine-core/src/main/kotlin/org/vitrivr/engine/core/model/metamodel/Analyser.kt b/vitrivr-engine-core/src/main/kotlin/org/vitrivr/engine/core/model/metamodel/Analyser.kt
@@ -23,7 +23,7 @@ import kotlin.reflect.KClass
  */
 interface Analyser<C: ContentElement<*>, D: Descriptor> {
     /** The [KClass]es of the [ContentElement] accepted by this [Analyser].  */
-    val contentClasses: Set<KClass<C>>
+    val contentClasses: Set<KClass<out ContentElement<*>>>
 
     /** The [KClass] of the [Descriptor] generated by this [Analyser].  */
     val descriptorClass: KClass<D>