diff --git a/settings.gradle.kts b/settings.gradle.kts index 4a52fc02..5cfaf46a 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -46,6 +46,10 @@ include("skainet-data:skainet-data-simple") // ====== IO include("skainet-io:skainet-io-core") include("skainet-io:skainet-io-gguf") +include("skainet-io:skainet-io-image") + +// ====== models +include("skainet-models:skainet-model-yolo") // ====== Integrations include("skainet-integrations:skainet-simple-cpu") diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index ca0070b9..7e6814d0 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -7,6 +7,7 @@ import sk.ainet.lang.types.DType import sk.ainet.lang.ops.TensorOp import sk.ainet.lang.ops.InProgress import sk.ainet.lang.tensor.data.TensorDataFactory +import sk.ainet.lang.tensor.ops.UpsampleMode @InProgress("cpu", owner = "team:cpu", issue = "task-ops.md#defaultcpuops") public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory) : TensorOps { @@ -536,6 +537,37 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory return CpuTensor(outData, this, input.dtype) } + @TensorOp() + @InProgress("cpu", owner = "team:cpu", issue = "task-ops.md#op-upsample2d") + override fun upsample2d( + input: Tensor, + scale: Pair, + mode: UpsampleMode, + alignCorners: Boolean + ): Tensor { + require(input.rank == 4) { "upsample2d: input must be 4D (N, C, H, W)" } + val (scaleH, scaleW) = scale + require(scaleH > 0 && scaleW > 0) { "upsample2d: scale factors must be positive" } + require(mode == UpsampleMode.Nearest) { "upsample2d: only Nearest mode is implemented on CPU backend" } + + val n = input.shape[0] + val c = input.shape[1] + val inH = input.shape[2] + val inW = input.shape[3] + val outH = inH * scaleH + val outW = inW * scaleW + val outShape = Shape(n, c, outH, outW) + + val outData = dataFactory.init(outShape, input.dtype) { idx -> + val oh = idx[2] + val ow = idx[3] + val ih = oh / scaleH + val iw = ow / scaleW + input.data.get(idx[0], idx[1], ih, iw) + } + return CpuTensor(outData, this, input.dtype) + } + @TensorOp() @InProgress("cpu", owner = "team:cpu", issue = "task-ops.md#op-maxpool2d") override fun maxPool2d( diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsUpsampleTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsUpsampleTest.kt new file mode 100644 index 00000000..c045dfe5 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsUpsampleTest.kt @@ -0,0 +1,39 @@ +package sk.ainet.sk.ainet.exec.tensor.ops + +import kotlin.test.Test +import kotlin.test.assertEquals +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.data +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.dsl.tensor +import sk.ainet.lang.tensor.ops.UpsampleMode +import sk.ainet.lang.types.FP32 + +class DefaultCpuOpsUpsampleTest { + private val ctx = DirectCpuExecutionContext() + private val ops = ctx.ops + + @Test + fun nearest_mode_repeats_values() { + data(ctx) { _ -> + val input = tensor { + shape(1, 1, 2, 2) { + init { idx -> (1 + idx[2] * 2 + idx[3]).toFloat() } + } + } + + val upsampled = ops.upsample2d( + input = input, + scale = 2 to 2, + mode = UpsampleMode.Nearest, + alignCorners = false + ) + + assertEquals(Shape(1, 1, 4, 4), upsampled.shape) + assertEquals(1f, upsampled.data[0, 0, 0, 0]) + assertEquals(2f, upsampled.data[0, 0, 0, 2]) + assertEquals(3f, upsampled.data[0, 0, 2, 0]) + assertEquals(4f, upsampled.data[0, 0, 3, 3]) + } + } +} diff --git a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/lang/trace/TracingTensorOps.kt b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/lang/trace/TracingTensorOps.kt index a577f451..a17e4712 100644 --- a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/lang/trace/TracingTensorOps.kt +++ b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/lang/trace/TracingTensorOps.kt @@ -3,6 +3,7 @@ package sk.ainet.lang.trace import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.ops.TensorOps +import sk.ainet.lang.tensor.ops.UpsampleMode import sk.ainet.lang.types.DType /** @@ -77,12 +78,49 @@ public class TracingTensorOps( kernelSize: Pair, stride: Pair, padding: Pair - ): Tensor = base.maxPool2d(input, kernelSize, stride, padding) + ): Tensor { + val out = base.maxPool2d(input, kernelSize, stride, padding) + val attrs = mapOf("kernel" to kernelSize, "stride" to stride, "padding" to padding) + sink.onOpExecuted( + OpTrace( + opType = "maxPool2d", + inputs = listOf(session.refOf(input)), + outputs = listOf(session.refOf(out)), + attributes = attrs + ) + ) + return out + } + + override fun upsample2d( + input: Tensor, + scale: Pair, + mode: UpsampleMode, + alignCorners: Boolean + ): Tensor { + val out = base.upsample2d(input, scale, mode, alignCorners) + val inputs = listOf(session.refOf(input)) + val outputs = listOf(session.refOf(out)) + val attrs = OpAttributeFactory.shapesAndDTypes(listOf(input), listOf(out)) + mapOf( + "scale" to listOf(scale.first, scale.second), + "mode" to mode.name, + "alignCorners" to alignCorners + ) + sink.onOpExecuted(OpTrace(opType = "upsample2d", inputs = inputs, outputs = outputs, attributes = attrs)) + return out + } // ---- Shape ops ---- override fun reshape(tensor: Tensor, newShape: Shape): Tensor = base.reshape(tensor, newShape) override fun flatten(tensor: Tensor, startDim: Int, endDim: Int): Tensor = base.flatten(tensor, startDim, endDim) - override fun concat(tensors: List>, dim: Int): Tensor = base.concat(tensors, dim) + override fun concat(tensors: List>, dim: Int): Tensor { + val out = base.concat(tensors, dim) + val inputs = tensors.map { session.refOf(it) } + val outputs = listOf(session.refOf(out)) + val attrs = mapOf("dim" to dim, "count" to tensors.size) + sink.onOpExecuted(OpTrace(opType = "concat", inputs = inputs, outputs = outputs, attributes = attrs)) + return out + } override fun split(tensor: Tensor, splitSize: Int, dim: Int): List> = base.split(tensor, splitSize, dim) override fun squeeze(tensor: Tensor, dim: Int?): Tensor = base.squeeze(tensor, dim) override fun unsqueeze(tensor: Tensor, dim: Int): Tensor = base.unsqueeze(tensor, dim) @@ -97,6 +135,15 @@ public class TracingTensorOps( return out } + override fun silu(tensor: Tensor): Tensor { + val out = base.silu(tensor) + val inputs = listOf(session.refOf(tensor)) + val outputs = listOf(session.refOf(out)) + val attrs = OpAttributeFactory.unary(tensor, out) + sink.onOpExecuted(OpTrace(opType = "silu", inputs = inputs, outputs = outputs, attributes = attrs)) + return out + } + override fun softmax(tensor: Tensor, dim: Int): Tensor { val out = base.softmax(tensor, dim) val inputs = listOf(session.refOf(tensor)) @@ -114,7 +161,6 @@ public class TracingTensorOps( sink.onOpExecuted(OpTrace(opType = "sigmoid", inputs = inputs, outputs = outputs, attributes = attrs)) return out } - override fun silu(tensor: Tensor): Tensor = base.silu(tensor) override fun gelu(tensor: Tensor): Tensor = base.gelu(tensor) // ---- Reductions ---- diff --git a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt index 99cd1686..9acf53be 100644 --- a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt +++ b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt @@ -239,6 +239,38 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor return out } + override fun upsample2d( + input: Tensor, + scale: Pair, + mode: UpsampleMode, + alignCorners: Boolean + ): Tensor { + val out = base.upsample2d(input, scale, mode, alignCorners) + val params = mapOf( + "scale" to scale.toList(), + "mode" to mode.name, + "alignCorners" to alignCorners + ) + record(Upsample2dOperation(params), listOf(input), listOf(out)) + return out + } + + override fun concat(tensors: List>, dim: Int): Tensor { + val out = base.concat(tensors, dim) + record( + ConcatRecordingOperation(mapOf("dim" to dim, "count" to tensors.size)), + tensors, + listOf(out) + ) + return out + } + + override fun silu(tensor: Tensor): Tensor { + val out = base.silu(tensor) + record(SimpleActivationOperation("silu"), listOf(tensor), listOf(out)) + return out + } + // --- Shape ops --- override fun reshape(tensor: Tensor, newShape: Shape): Tensor { val out = base.reshape(tensor, newShape) @@ -285,9 +317,7 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor } // Delegations for unrecorded/less critical ops - override fun concat(tensors: List>, dim: Int): Tensor = base.concat(tensors, dim) override fun split(tensor: Tensor, splitSize: Int, dim: Int): List> = base.split(tensor, splitSize, dim) - override fun silu(tensor: Tensor): Tensor = base.silu(tensor) override fun gelu(tensor: Tensor): Tensor = base.gelu(tensor) override fun sum(tensor: Tensor, dim: Int?): Tensor = base.sum(tensor, dim) override fun mean(tensor: Tensor, dim: Int?): Tensor = base.mean(tensor, dim) @@ -295,4 +325,52 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor override fun sqrt(tensor: Tensor): Tensor = base.sqrt(tensor) override fun convert(tensor: Tensor, targetType: TTo): Tensor = base.convert(tensor, targetType) override fun tril(tensor: Tensor, k: Int): Tensor = base.tril(tensor, k) -} \ No newline at end of file +} + +private class ConcatRecordingOperation( + override val parameters: Map = emptyMap() +) : BaseOperation("concat", "shape", parameters) { + override fun execute(inputs: List>): List> = + throw UnsupportedOperationException("execute not supported in recording op") + + override fun validateInputs(inputs: List): ValidationResult = + if (inputs.size < 2) ValidationResult.Invalid(listOf("concat requires at least 2 inputs")) else ValidationResult.Valid + + override fun inferOutputs(inputs: List): List { + // validateInputs ensures at least 2 inputs, so first() is safe + val first = inputs.first() + return listOf( + TensorSpec( + name = "concat_output", + shape = first.shape, + dtype = first.dtype, + requiresGrad = inputs.any { it.requiresGrad } + ) + ) + } + + override fun clone(newParameters: Map): Operation = ConcatRecordingOperation(newParameters) +} + +private class SimpleActivationOperation( + activationName: String +) : BaseOperation(activationName, "activation", emptyMap()) { + override fun execute(inputs: List>): List> = + throw UnsupportedOperationException("execute not supported in recording op") + + override fun validateInputs(inputs: List): ValidationResult = + if (inputs.size != 1) ValidationResult.Invalid(listOf("Activation requires 1 input")) else ValidationResult.Valid + + override fun inferOutputs(inputs: List): List = + listOf( + TensorSpec( + name = "activation_output", + // validateInputs ensures exactly 1 input, so first() is safe + shape = inputs.first().shape, + dtype = inputs.first().dtype, + requiresGrad = inputs.any { it.requiresGrad } + ) + ) + + override fun clone(newParameters: Map): Operation = SimpleActivationOperation(name) +} diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/GraphTensorOps.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/GraphTensorOps.kt index fc16e31e..9103e6a7 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/GraphTensorOps.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/GraphTensorOps.kt @@ -19,7 +19,7 @@ public class GraphTensorOps( override fun add(a: Tensor, b: Tensor): Tensor { val result = baseOps.add(a, b) if (executionContext.isRecording) { - executionContext.currentTape?.recordOperation(AddOperation(), listOf(a, b), listOf(result)) + executionContext.currentTape?.recordOperation(AddOperation(), listOf(a, b), listOf(result)) } return result } @@ -48,8 +48,18 @@ public class GraphTensorOps( padding: Pair ): Tensor = baseOps.maxPool2d(input, kernelSize, stride, padding) - override fun reshape(tensor: Tensor, newShape: Shape): Tensor = baseOps.reshape(tensor, newShape) - override fun flatten(tensor: Tensor, startDim: Int, endDim: Int): Tensor = baseOps.flatten(tensor, startDim, endDim) + override fun upsample2d( + input: Tensor, + scale: Pair, + mode: UpsampleMode, + alignCorners: Boolean + ): Tensor = baseOps.upsample2d(input, scale, mode, alignCorners) + + override fun reshape(tensor: Tensor, newShape: Shape): Tensor = + baseOps.reshape(tensor, newShape) + + override fun flatten(tensor: Tensor, startDim: Int, endDim: Int): Tensor = + baseOps.flatten(tensor, startDim, endDim) override fun relu(tensor: Tensor): Tensor = baseOps.relu(tensor) override fun softmax(tensor: Tensor, dim: Int): Tensor = baseOps.softmax(tensor, dim) @@ -57,8 +67,12 @@ public class GraphTensorOps( override fun silu(tensor: Tensor): Tensor = baseOps.silu(tensor) override fun gelu(tensor: Tensor): Tensor = baseOps.gelu(tensor) - override fun concat(tensors: List>, dim: Int): Tensor = baseOps.concat(tensors, dim) - override fun split(tensor: Tensor, splitSize: Int, dim: Int): List> = baseOps.split(tensor, splitSize, dim) + override fun concat(tensors: List>, dim: Int): Tensor = + baseOps.concat(tensors, dim) + + override fun split(tensor: Tensor, splitSize: Int, dim: Int): List> = + baseOps.split(tensor, splitSize, dim) + override fun squeeze(tensor: Tensor, dim: Int?): Tensor = baseOps.squeeze(tensor, dim) override fun unsqueeze(tensor: Tensor, dim: Int): Tensor = baseOps.unsqueeze(tensor, dim) @@ -66,6 +80,8 @@ public class GraphTensorOps( override fun mean(tensor: Tensor, dim: Int?): Tensor = baseOps.mean(tensor, dim) override fun variance(tensor: Tensor, dim: Int?): Tensor = baseOps.variance(tensor, dim) override fun sqrt(tensor: Tensor): Tensor = baseOps.sqrt(tensor) - override fun convert(tensor: Tensor, targetType: TTo): Tensor = baseOps.convert(tensor, targetType) + override fun convert(tensor: Tensor, targetType: TTo): Tensor = + baseOps.convert(tensor, targetType) + override fun tril(tensor: Tensor, k: Int): Tensor = baseOps.tril(tensor, k) } \ No newline at end of file diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/MinimalAddTensorOps.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/MinimalAddTensorOps.kt index 0b671fce..e1f3d0fe 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/MinimalAddTensorOps.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/MinimalAddTensorOps.kt @@ -4,6 +4,7 @@ import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.data.DenseTensorDataFactory import sk.ainet.lang.tensor.ops.TensorOps +import sk.ainet.lang.tensor.ops.UpsampleMode import sk.ainet.lang.tensor.ops.VoidTensorOps import sk.ainet.lang.types.DType import sk.ainet.lang.types.FP32 @@ -37,25 +38,59 @@ public class MinimalAddTensorOps : TensorOps { override fun divide(a: Tensor, b: Tensor): Tensor = delegate.divide(a, b) override fun matmul(a: Tensor, b: Tensor): Tensor = delegate.matmul(a, b) override fun transpose(tensor: Tensor): Tensor = delegate.transpose(tensor) - override fun conv2d(input: Tensor, weight: Tensor, bias: Tensor?, stride: Pair, padding: Pair, dilation: Pair, groups: Int): Tensor = + override fun conv2d( + input: Tensor, + weight: Tensor, + bias: Tensor?, + stride: Pair, + padding: Pair, + dilation: Pair, + groups: Int + ): Tensor = delegate.conv2d(input, weight, bias, stride, padding, dilation, groups) - override fun maxPool2d(input: Tensor, kernelSize: Pair, stride: Pair, padding: Pair): Tensor = + + override fun maxPool2d( + input: Tensor, + kernelSize: Pair, + stride: Pair, + padding: Pair + ): Tensor = delegate.maxPool2d(input, kernelSize, stride, padding) - override fun reshape(tensor: Tensor, newShape: Shape): Tensor = delegate.reshape(tensor, newShape) - override fun flatten(tensor: Tensor, startDim: Int, endDim: Int): Tensor = delegate.flatten(tensor, startDim, endDim) + + override fun upsample2d( + input: Tensor, + scale: Pair, + mode: UpsampleMode, + alignCorners: Boolean + ): Tensor = delegate.upsample2d(input, scale, mode, alignCorners) + + override fun reshape(tensor: Tensor, newShape: Shape): Tensor = + delegate.reshape(tensor, newShape) + + override fun flatten(tensor: Tensor, startDim: Int, endDim: Int): Tensor = + delegate.flatten(tensor, startDim, endDim) + override fun relu(tensor: Tensor): Tensor = delegate.relu(tensor) override fun softmax(tensor: Tensor, dim: Int): Tensor = delegate.softmax(tensor, dim) override fun sigmoid(tensor: Tensor): Tensor = delegate.sigmoid(tensor) override fun silu(tensor: Tensor): Tensor = delegate.silu(tensor) override fun gelu(tensor: Tensor): Tensor = delegate.gelu(tensor) - override fun concat(tensors: List>, dim: Int): Tensor = delegate.concat(tensors, dim) - override fun split(tensor: Tensor, splitSize: Int, dim: Int): List> = delegate.split(tensor, splitSize, dim) + override fun concat(tensors: List>, dim: Int): Tensor = + delegate.concat(tensors, dim) + + override fun split(tensor: Tensor, splitSize: Int, dim: Int): List> = + delegate.split(tensor, splitSize, dim) + override fun squeeze(tensor: Tensor, dim: Int?): Tensor = delegate.squeeze(tensor, dim) - override fun unsqueeze(tensor: Tensor, dim: Int): Tensor = delegate.unsqueeze(tensor, dim) + override fun unsqueeze(tensor: Tensor, dim: Int): Tensor = + delegate.unsqueeze(tensor, dim) + override fun sum(tensor: Tensor, dim: Int?): Tensor = delegate.sum(tensor, dim) override fun mean(tensor: Tensor, dim: Int?): Tensor = delegate.mean(tensor, dim) override fun variance(tensor: Tensor, dim: Int?): Tensor = delegate.variance(tensor, dim) override fun sqrt(tensor: Tensor): Tensor = delegate.sqrt(tensor) - override fun convert(tensor: Tensor, targetType: TTo): Tensor = delegate.convert(tensor, targetType) + override fun convert(tensor: Tensor, targetType: TTo): Tensor = + delegate.convert(tensor, targetType) + override fun tril(tensor: Tensor, k: Int): Tensor = delegate.tril(tensor, k) } diff --git a/skainet-io/skainet-io-core/build.gradle.kts b/skainet-io/skainet-io-core/build.gradle.kts index 6d7fdbf6..ae1c9239 100644 --- a/skainet-io/skainet-io-core/build.gradle.kts +++ b/skainet-io/skainet-io-core/build.gradle.kts @@ -47,8 +47,6 @@ kotlin { browser() } - linuxX64 () - sourceSets { val commonMain by getting { dependencies { diff --git a/skainet-io/skainet-io-image/build.gradle.kts b/skainet-io/skainet-io-image/build.gradle.kts new file mode 100644 index 00000000..66e1a6b5 --- /dev/null +++ b/skainet-io/skainet-io-image/build.gradle.kts @@ -0,0 +1,76 @@ +@file:OptIn(ExperimentalWasmDsl::class) + +import org.jetbrains.kotlin.gradle.ExperimentalKotlinGradlePluginApi +import org.jetbrains.kotlin.gradle.ExperimentalWasmDsl +import org.jetbrains.kotlin.gradle.dsl.JvmTarget + +plugins { + alias(libs.plugins.kotlinMultiplatform) + alias(libs.plugins.androidLibrary) + alias(libs.plugins.vanniktech.mavenPublish) +} + +kotlin { + targets.configureEach { + compilations.configureEach { + compileTaskProvider.get().compilerOptions { + freeCompilerArgs.add("-Xexpect-actual-classes") + } + } + } + + jvm() + + androidTarget { + publishLibraryVariants("release") + @OptIn(ExperimentalKotlinGradlePluginApi::class) + compilerOptions { + jvmTarget.set(JvmTarget.JVM_11) + } + } + + iosArm64() + iosSimulatorArm64() + macosArm64 () + linuxX64 () + linuxArm64 () + + js { + browser() + } + + @OptIn(ExperimentalWasmDsl::class) + wasmJs { + browser() + } + + + + sourceSets { + val commonMain by getting { + dependencies { + implementation(project(":skainet-lang:skainet-lang-core")) + implementation(libs.kotlinx.io.core) + implementation(libs.kotlinx.coroutines) + } + } + + val commonTest by getting { + dependencies { + implementation(libs.kotlin.test) + } + } + } +} + +android { + namespace = "sk.ainet.io.image" + compileSdk = libs.versions.android.compileSdk.get().toInt() + defaultConfig { + minSdk = libs.versions.android.minSdk.get().toInt() + } + compileOptions { + sourceCompatibility = JavaVersion.VERSION_11 + targetCompatibility = JavaVersion.VERSION_11 + } +} diff --git a/skainet-io/skainet-io-image/gradle.properties b/skainet-io/skainet-io-image/gradle.properties new file mode 100644 index 00000000..1ddf7182 --- /dev/null +++ b/skainet-io/skainet-io-image/gradle.properties @@ -0,0 +1,2 @@ +POM_ARTIFACT_ID=skainet-io-image +POM_NAME=skainet neural network scripting API \ No newline at end of file diff --git a/skainet-io/skainet-io-image/src/androidMain/kotlin/sk/ainet/io/image/ImageInterop.android.kt b/skainet-io/skainet-io-image/src/androidMain/kotlin/sk/ainet/io/image/ImageInterop.android.kt new file mode 100644 index 00000000..0f4d1d0b --- /dev/null +++ b/skainet-io/skainet-io-image/src/androidMain/kotlin/sk/ainet/io/image/ImageInterop.android.kt @@ -0,0 +1,129 @@ +package sk.ainet.io.image + +import android.graphics.Bitmap +import sk.ainet.context.ExecutionContext +import sk.ainet.context.data +import sk.ainet.lang.tensor.dsl.* +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP16 +import android.graphics.Bitmap.Config + +actual typealias PlatformBitmapImage = Bitmap + +/** + * Converts this Bitmap to a packed RGB (R,G,B) ByteArray. + * Output order: row-major (top→bottom, left→right), 3 bytes per pixel. + * Works for any input config; internally uses ARGB_8888 for correctness. + */ +private fun PlatformBitmapImage.toRgbByteArray(): ByteArray { + // Ensure ARGB_8888 so channel extraction is predictable + val src = if (config != Bitmap.Config.ARGB_8888) { + copy(Bitmap.Config.ARGB_8888, /* isMutable = */false) + } else this + + val w = src.width + val h = src.height + + val pixels = IntArray(w * h) + src.getPixels(pixels, 0, w, 0, 0, w, h) + + val rgb = ByteArray(w * h * 3) + var o = 0 + // Each pixel is Android ARGB: 0xAARRGGBB + for (p in pixels) { + rgb[o++] = ((p shr 16) and 0xFF).toByte() // R + rgb[o++] = ((p shr 8) and 0xFF).toByte() // G + rgb[o++] = (p and 0xFF).toByte() // B + } + // If we created a copy, we can let GC reclaim it; 'src' will be collected + return rgb +} + +actual fun platformImageToArgb( + image: PlatformBitmapImage, + ctx: ExecutionContext +): Tensor { + val src = if (image.config != Bitmap.Config.ARGB_8888) { + image.copy(Bitmap.Config.ARGB_8888, false) + } else image + + val width = src.width + val height = src.height + + // 1) Pull pixels as ARGB_8888 + val argb = IntArray(width * height) + src.getPixels(argb, 0, width, /*x=*/0, /*y=*/0, width, height) + + // 2) Convert to FloatArray in RGB order, CHW layout to match (1,3,H,W) + val rgbChw = FloatArray(width * height * 3) + for (y in 0 until height) { + for (x in 0 until width) { + val px = argb[y * width + x] + val r = (px shr 16) and 0xFF + val g = (px shr 8) and 0xFF + val b = px and 0xFF + val hwIndex = y * width + x + rgbChw[hwIndex] = r.toFloat() + rgbChw[1 * height * width + hwIndex] = g.toFloat() + rgbChw[2 * height * width + hwIndex] = b.toFloat() + } + } + + // 3) Create a tensor with shape (1, C, H, W) + return data(ctx) { + tensor { + shape(1, 3, height, width) { + fromArray(rgbChw) + } + } + } +} + +actual fun platformImageToRgbByteArray(image: PlatformBitmapImage): ByteArray { + val src = if (image.config != Bitmap.Config.ARGB_8888) { + image.copy(Bitmap.Config.ARGB_8888, /*mutable=*/false) + } else image + return src.toRgbByteArray() +} + +actual fun platformImageSize(image: PlatformBitmapImage): Pair = + image.width to image.height + +actual fun argbToPlatformImage( + image: Tensor, + ctx: ExecutionContext +): PlatformBitmapImage { + // Tensor is NCHW according to this module's convention used in platformImageToArgb + val shape = image.data.shape + val channels = shape[1] + val height = shape[2] + val width = shape[3] + + val pixels = IntArray(width * height) + var i = 0 + for (y in 0 until height) { + for (x in 0 until width) { + val argb = when (channels) { + 1 -> { + val v = image.data[0, 0, y, x].toInt().coerceIn(0, 255) + (0xFF shl 24) or (v shl 16) or (v shl 8) or v + } + 3 -> { + val r = image.data[0, 0, y, x].toInt().coerceIn(0, 255) + val g = image.data[0, 1, y, x].toInt().coerceIn(0, 255) + val b = image.data[0, 2, y, x].toInt().coerceIn(0, 255) + (0xFF shl 24) or (r shl 16) or (g shl 8) or b + } + else -> { + val v = image.data[0, 0, y, x].toInt().coerceIn(0, 255) + (0xFF shl 24) or (v shl 16) or (v shl 8) or v + } + } + pixels[i++] = argb + } + } + + val bmp = Bitmap.createBitmap(width, height, Config.ARGB_8888) + bmp.setPixels(pixels, 0, width, 0, 0, width, height) + return bmp +} diff --git a/skainet-io/skainet-io-image/src/commonMain/kotlin/sk/ainet/io/image/ImageInterop.kt b/skainet-io/skainet-io-image/src/commonMain/kotlin/sk/ainet/io/image/ImageInterop.kt new file mode 100644 index 00000000..631956b7 --- /dev/null +++ b/skainet-io/skainet-io-image/src/commonMain/kotlin/sk/ainet/io/image/ImageInterop.kt @@ -0,0 +1,19 @@ +package sk.ainet.io.image + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP16 + +expect class PlatformBitmapImage + +// Convert Platform image to an RGB tensor with shape (1, 3, H, W) +expect fun platformImageToArgb(image: PlatformBitmapImage, ctx: ExecutionContext): Tensor + +// Convert RGB tensor (1, 3/1, H, W) back to platform image +expect fun argbToPlatformImage(image: Tensor, ctx: ExecutionContext): PlatformBitmapImage + +// RGB bytes in row-major order, 3 bytes per pixel +expect fun platformImageToRgbByteArray(image: PlatformBitmapImage): ByteArray + +// Returns width to Pair.first and height to Pair.second +expect fun platformImageSize(image: PlatformBitmapImage): Pair \ No newline at end of file diff --git a/skainet-io/skainet-io-image/src/iosMain/kotlin/sk/ainet/io/image/ImageInterop.ios.kt b/skainet-io/skainet-io-image/src/iosMain/kotlin/sk/ainet/io/image/ImageInterop.ios.kt new file mode 100644 index 00000000..6cf60f48 --- /dev/null +++ b/skainet-io/skainet-io-image/src/iosMain/kotlin/sk/ainet/io/image/ImageInterop.ios.kt @@ -0,0 +1,23 @@ +package sk.ainet.io.image + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP16 + +actual class PlatformBitmapImage + +actual fun platformImageToArgb(image: PlatformBitmapImage, ctx: ExecutionContext): Tensor { + throw NotImplementedError("platformImageToArgb is not implemented for iOS yet") +} + +actual fun argbToPlatformImage(image: Tensor, ctx: ExecutionContext): PlatformBitmapImage { + throw NotImplementedError("argbToPlatformImage is not implemented for iOS yet") +} + +actual fun platformImageToRgbByteArray(image: PlatformBitmapImage): ByteArray { + throw NotImplementedError("platformImageToRgbByteArray is not implemented for iOS yet") +} + +actual fun platformImageSize(image: PlatformBitmapImage): Pair { + throw NotImplementedError("platformImageSize is not implemented for iOS yet") +} diff --git a/skainet-io/skainet-io-image/src/jsMain/kotlin/sk/ainet/io/image/ImageInterop.js.kt b/skainet-io/skainet-io-image/src/jsMain/kotlin/sk/ainet/io/image/ImageInterop.js.kt new file mode 100644 index 00000000..0b75afea --- /dev/null +++ b/skainet-io/skainet-io-image/src/jsMain/kotlin/sk/ainet/io/image/ImageInterop.js.kt @@ -0,0 +1,23 @@ +package sk.ainet.io.image + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP16 + +actual class PlatformBitmapImage + +actual fun platformImageToArgb(image: PlatformBitmapImage, ctx: ExecutionContext): Tensor { + throw NotImplementedError("platformImageToArgb is not implemented for JS yet") +} + +actual fun argbToPlatformImage(image: Tensor, ctx: ExecutionContext): PlatformBitmapImage { + throw NotImplementedError("argbToPlatformImage is not implemented for JS yet") +} + +actual fun platformImageToRgbByteArray(image: PlatformBitmapImage): ByteArray { + throw NotImplementedError("platformImageToRgbByteArray is not implemented for JS yet") +} + +actual fun platformImageSize(image: PlatformBitmapImage): Pair { + throw NotImplementedError("platformImageSize is not implemented for JS yet") +} diff --git a/skainet-io/skainet-io-image/src/jvmMain/kotlin/sk/ainet/io/image/ImageInterop.jvm.kt b/skainet-io/skainet-io-image/src/jvmMain/kotlin/sk/ainet/io/image/ImageInterop.jvm.kt new file mode 100644 index 00000000..78b54c3b --- /dev/null +++ b/skainet-io/skainet-io-image/src/jvmMain/kotlin/sk/ainet/io/image/ImageInterop.jvm.kt @@ -0,0 +1,100 @@ +package sk.ainet.io.image + +import sk.ainet.context.ExecutionContext +import sk.ainet.context.data +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.dsl.tensor +import sk.ainet.lang.types.FP16 +import java.awt.Color +import java.awt.image.BufferedImage + +actual typealias PlatformBitmapImage = BufferedImage + +actual fun platformImageToArgb( + image: PlatformBitmapImage, + ctx: ExecutionContext +): Tensor { + val width = image.width + val height = image.height + + val argb = IntArray(width * height) + image.getRGB(0, 0, width, height, argb, 0, width) + + val rgbChw = FloatArray(width * height * 3) + for (y in 0 until height) { + for (x in 0 until width) { + val px = argb[y * width + x] + val r = (px shr 16) and 0xFF + val g = (px shr 8) and 0xFF + val b = px and 0xFF + val hw = y * width + x + rgbChw[hw] = r.toFloat() + rgbChw[1 * height * width + hw] = g.toFloat() + rgbChw[2 * height * width + hw] = b.toFloat() + } + } + + return data(ctx) { + tensor { + shape(1, 3, height, width) { + fromArray(rgbChw) + } + } + } +} + +actual fun argbToPlatformImage( + image: Tensor, + ctx: ExecutionContext +): PlatformBitmapImage { + val shape = image.data.shape + val channels = shape[1] + val height = shape[2] + val width = shape[3] + + val pixels = IntArray(width * height) + var i = 0 + for (y in 0 until height) { + for (x in 0 until width) { + val argb = when (channels) { + 1 -> { + val v = image.data[0, 0, y, x].toInt().coerceIn(0, 255) + (0xFF shl 24) or (v shl 16) or (v shl 8) or v + } + 3 -> { + val r = image.data[0, 0, y, x].toInt().coerceIn(0, 255) + val g = image.data[0, 1, y, x].toInt().coerceIn(0, 255) + val b = image.data[0, 2, y, x].toInt().coerceIn(0, 255) + (0xFF shl 24) or (r shl 16) or (g shl 8) or b + } + else -> { + val v = image.data[0, 0, y, x].toInt().coerceIn(0, 255) + (0xFF shl 24) or (v shl 16) or (v shl 8) or v + } + } + pixels[i++] = argb + } + } + + val out = BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB) + out.setRGB(0, 0, width, height, pixels, 0, width) + return out +} + +actual fun platformImageToRgbByteArray(image: PlatformBitmapImage): ByteArray { + val w = image.width + val h = image.height + val argb = IntArray(w * h) + image.getRGB(0, 0, w, h, argb, 0, w) + val out = ByteArray(w * h * 3) + var o = 0 + for (p in argb) { + out[o++] = ((p shr 16) and 0xFF).toByte() + out[o++] = ((p shr 8) and 0xFF).toByte() + out[o++] = (p and 0xFF).toByte() + } + return out +} + +actual fun platformImageSize(image: PlatformBitmapImage): Pair = + image.width to image.height diff --git a/skainet-io/skainet-io-image/src/linuxArm64Main/kotlin/sk/ainet/io/image/ImageInterop.linuxArm64.kt b/skainet-io/skainet-io-image/src/linuxArm64Main/kotlin/sk/ainet/io/image/ImageInterop.linuxArm64.kt new file mode 100644 index 00000000..98e50103 --- /dev/null +++ b/skainet-io/skainet-io-image/src/linuxArm64Main/kotlin/sk/ainet/io/image/ImageInterop.linuxArm64.kt @@ -0,0 +1,23 @@ +package sk.ainet.io.image + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP16 + +actual class PlatformBitmapImage + +actual fun platformImageToArgb(image: PlatformBitmapImage, ctx: ExecutionContext): Tensor { + throw NotImplementedError("platformImageToArgb is not implemented for Linux arm64 yet") +} + +actual fun argbToPlatformImage(image: Tensor, ctx: ExecutionContext): PlatformBitmapImage { + throw NotImplementedError("argbToPlatformImage is not implemented for Linux arm64 yet") +} + +actual fun platformImageToRgbByteArray(image: PlatformBitmapImage): ByteArray { + throw NotImplementedError("platformImageToRgbByteArray is not implemented for Linux arm64 yet") +} + +actual fun platformImageSize(image: PlatformBitmapImage): Pair { + throw NotImplementedError("platformImageSize is not implemented for Linux arm64 yet") +} diff --git a/skainet-io/skainet-io-image/src/linuxX64Main/kotlin/sk/ainet/io/image/ImageInterop.linuxX64.kt b/skainet-io/skainet-io-image/src/linuxX64Main/kotlin/sk/ainet/io/image/ImageInterop.linuxX64.kt new file mode 100644 index 00000000..aa0c2dfd --- /dev/null +++ b/skainet-io/skainet-io-image/src/linuxX64Main/kotlin/sk/ainet/io/image/ImageInterop.linuxX64.kt @@ -0,0 +1,23 @@ +package sk.ainet.io.image + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP16 + +actual class PlatformBitmapImage + +actual fun platformImageToArgb(image: PlatformBitmapImage, ctx: ExecutionContext): Tensor { + throw NotImplementedError("platformImageToArgb is not implemented for Linux x64 yet") +} + +actual fun argbToPlatformImage(image: Tensor, ctx: ExecutionContext): PlatformBitmapImage { + throw NotImplementedError("argbToPlatformImage is not implemented for Linux x64 yet") +} + +actual fun platformImageToRgbByteArray(image: PlatformBitmapImage): ByteArray { + throw NotImplementedError("platformImageToRgbByteArray is not implemented for Linux x64 yet") +} + +actual fun platformImageSize(image: PlatformBitmapImage): Pair { + throw NotImplementedError("platformImageSize is not implemented for Linux x64 yet") +} diff --git a/skainet-io/skainet-io-image/src/macosArm64Main/kotlin/sk/ainet/io/image/ImageInterop.macos.kt b/skainet-io/skainet-io-image/src/macosArm64Main/kotlin/sk/ainet/io/image/ImageInterop.macos.kt new file mode 100644 index 00000000..a495f179 --- /dev/null +++ b/skainet-io/skainet-io-image/src/macosArm64Main/kotlin/sk/ainet/io/image/ImageInterop.macos.kt @@ -0,0 +1,184 @@ +@file:OptIn(kotlinx.cinterop.ExperimentalForeignApi::class) + +package sk.ainet.io.image + +import kotlinx.cinterop.* +import platform.AppKit.NSImage +import platform.CoreFoundation.CFDataCreate +import platform.CoreFoundation.kCFAllocatorDefault +import platform.CoreGraphics.* +import platform.CoreGraphics.CGImageAlphaInfo +import platform.CoreGraphics.CGColorRenderingIntent +import platform.CoreGraphics.CGSize +import platform.CoreGraphics.CGSizeMake +import sk.ainet.context.ExecutionContext +import sk.ainet.context.data +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.dsl.tensor +import sk.ainet.lang.types.FP16 + +actual typealias PlatformBitmapImage = NSImage + +actual fun platformImageToArgb(image: PlatformBitmapImage, ctx: ExecutionContext): Tensor = + data(ctx) { + // Draw into RGBA8 buffer and read back as RGB + val (w, h) = platformImageSize(image) + val rgba = drawImageIntoRgbaBuffer(image, w, h) + + // Convert interleaved RGB bytes to CHW floats + val rgbChw = FloatArray(w * h * 3) + var p = 0 + for (y in 0 until h) { + for (x in 0 until w) { + val hw = y * w + x + val r = rgba[p].toUByte().toInt(); p++ + val g = rgba[p].toUByte().toInt(); p++ + val b = rgba[p].toUByte().toInt(); p++ + p++ // skip alpha + rgbChw[hw] = r.toFloat() + rgbChw[1 * h * w + hw] = g.toFloat() + rgbChw[2 * h * w + hw] = b.toFloat() + } + } + + tensor { + shape(1, 3, h, w) { + fromArray(rgbChw) + } + } + } + +actual fun argbToPlatformImage(image: Tensor, ctx: ExecutionContext): PlatformBitmapImage { + val shape = image.data.shape + val channels = shape[1] + val height = shape[2] + val width = shape[3] + + // Prepare RGBA bytes + val bytes = ByteArray(width * height * 4) + var o = 0 + for (y in 0 until height) { + for (x in 0 until width) { + when (channels) { + 1 -> { + val v = image.data[0, 0, y, x].toInt().coerceIn(0, 255).toByte() + bytes[o++] = v // R + bytes[o++] = v // G + bytes[o++] = v // B + bytes[o++] = (-1).toByte() // A = 255 + } + 3 -> { + val r = image.data[0, 0, y, x].toInt().coerceIn(0, 255).toByte() + val g = image.data[0, 1, y, x].toInt().coerceIn(0, 255).toByte() + val b = image.data[0, 2, y, x].toInt().coerceIn(0, 255).toByte() + bytes[o++] = r + bytes[o++] = g + bytes[o++] = b + bytes[o++] = (-1).toByte() // A = 255 + } + else -> { + val v = image.data[0, 0, y, x].toInt().coerceIn(0, 255).toByte() + bytes[o++] = v + bytes[o++] = v + bytes[o++] = v + bytes[o++] = (-1).toByte() + } + } + } + } + + // Create CGImage from RGBA bytes and wrap into NSImage + val cgImage = createCgImageFromRgba(bytes, width, height) + val size: CValue = CGSizeMake(width.toDouble(), height.toDouble()) + return NSImage(cGImage = cgImage, size = size) +} + +actual fun platformImageToRgbByteArray(image: PlatformBitmapImage): ByteArray { + val (w, h) = platformImageSize(image) + val rgba = drawImageIntoRgbaBuffer(image, w, h) + val out = ByteArray(w * h * 3) + var i = 0 + var p = 0 + while (p < rgba.size) { + out[i++] = rgba[p] // R + out[i++] = rgba[p + 1] // G + out[i++] = rgba[p + 2] // B + p += 4 // skip A + } + return out +} + +actual fun platformImageSize(image: PlatformBitmapImage): Pair { + val cg = image.CGImageForProposedRect(null, null, null) + return if (cg != null) { + CGImageGetWidth(cg).toInt() to CGImageGetHeight(cg).toInt() + } else { + // Fallback to point size if no CGImage representation (values in points) + val sz = image.size + val w = sz.useContents { width } + val h = sz.useContents { height } + w.toInt() to h.toInt() + } +} + +// Helpers +private fun drawImageIntoRgbaBuffer(image: NSImage, width: Int, height: Int): ByteArray = memScoped { + val cg = image.CGImageForProposedRect(null, null, null) + ?: error("NSImage has no CGImage representation") + + val colorSpace = CGColorSpaceCreateDeviceRGB() + ?: error("Failed to create RGB color space") + val bytesPerRow = width * 4 + val bitmapInfo: UInt = CGImageAlphaInfo.kCGImageAlphaPremultipliedLast.value + + val buffer = ByteArray(width * height * 4) + buffer.usePinned { pinned -> + val ctx = CGBitmapContextCreate( + pinned.addressOf(0), + width.convert(), + height.convert(), + 8.convert(), + bytesPerRow.convert(), + colorSpace, + bitmapInfo + ) ?: error("Failed to create bitmap context") + + // Flip the context vertically to match typical top-left origin drawing + val rect = CGRectMake(0.0, 0.0, width.toDouble(), height.toDouble()) + // Draw without additional transforms; NSImage/CGImage coordinates are bottom-left origin. + CGContextDrawImage(ctx, rect, cg) + } + buffer +} + +private fun createCgImageFromRgba(bytes: ByteArray, width: Int, height: Int): CGImageRef = memScoped { + val colorSpace = CGColorSpaceCreateDeviceRGB() + ?: error("Failed to create RGB color space") + val bytesPerRow = width * 4 + + val provider = bytes.usePinned { pinned -> + val cfData = CFDataCreate( + kCFAllocatorDefault, + pinned.addressOf(0).reinterpret(), + bytes.size.convert() + ) ?: error("Failed to create CFData") + CGDataProviderCreateWithCFData(cfData) + ?: error("Failed to create CGDataProvider") + } + + val bitmapInfo: UInt = CGImageAlphaInfo.kCGImageAlphaPremultipliedLast.value + + CGImageCreate( + width.convert(), + height.convert(), + 8.convert(), + 32.convert(), + bytesPerRow.convert(), + colorSpace, + bitmapInfo, + provider, + null, + true, + CGColorRenderingIntent.kCGRenderingIntentDefault + ) ?: error("Failed to create CGImage") +} diff --git a/skainet-io/skainet-io-image/src/wasmJsMain/kotlin/sk/ainet/io/image/ImageInterop.wasmJs.kt b/skainet-io/skainet-io-image/src/wasmJsMain/kotlin/sk/ainet/io/image/ImageInterop.wasmJs.kt new file mode 100644 index 00000000..48ab8eb0 --- /dev/null +++ b/skainet-io/skainet-io-image/src/wasmJsMain/kotlin/sk/ainet/io/image/ImageInterop.wasmJs.kt @@ -0,0 +1,23 @@ +package sk.ainet.io.image + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP16 + +actual class PlatformBitmapImage + +actual fun platformImageToArgb(image: PlatformBitmapImage, ctx: ExecutionContext): Tensor { + throw NotImplementedError("platformImageToArgb is not implemented for Wasm JS yet") +} + +actual fun argbToPlatformImage(image: Tensor, ctx: ExecutionContext): PlatformBitmapImage { + throw NotImplementedError("argbToPlatformImage is not implemented for Wasm JS yet") +} + +actual fun platformImageToRgbByteArray(image: PlatformBitmapImage): ByteArray { + throw NotImplementedError("platformImageToRgbByteArray is not implemented for Wasm JS yet") +} + +actual fun platformImageSize(image: PlatformBitmapImage): Pair { + throw NotImplementedError("platformImageSize is not implemented for Wasm JS yet") +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Upsample2d.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Upsample2d.kt new file mode 100644 index 00000000..b6c939ad --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Upsample2d.kt @@ -0,0 +1,40 @@ +package sk.ainet.lang.nn + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.ops.UpsampleMode +import sk.ainet.lang.types.DType + +/** + * 2D upsampling layer for increasing spatial resolution. + * + * Supports nearest-neighbor mode (typical for YOLO necks). Bilinear can be wired + * once backends add support. + */ +public class Upsample2d( + public val scale: Pair = 2 to 2, + public val mode: UpsampleMode = UpsampleMode.Nearest, + public val alignCorners: Boolean = false, + override val name: String = "Upsample2d", +) : Module() { + + init { + require(scale.first > 0 && scale.second > 0) { "Upsample2d scale factors must be positive" } + require(mode == UpsampleMode.Nearest || !alignCorners) { + "alignCorners applies only to bilinear mode" + } + } + + override val modules: List> + get() = emptyList() + + override fun forward(input: Tensor, ctx: ExecutionContext): Tensor = + sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) { + input.ops.upsample2d( + input = input, + scale = scale, + mode = mode, + alignCorners = alignCorners + ) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt index 96fdfeee..58da8700 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt @@ -7,12 +7,14 @@ import sk.ainet.lang.nn.Input import sk.ainet.lang.nn.Linear import sk.ainet.lang.nn.MaxPool2d import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.Upsample2d import sk.ainet.lang.nn.normalization.BatchNormalization import sk.ainet.lang.nn.normalization.GroupNormalization import sk.ainet.lang.nn.normalization.LayerNormalization import sk.ainet.lang.nn.topology.MLP import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.ops.UpsampleMode import sk.ainet.lang.types.DType import sk.ainet.context.ExecutionContext import sk.ainet.lang.nn.DefaultNeuralNetworkExecutionContext @@ -294,6 +296,29 @@ public interface NeuralNetworkDsl : NetworkDslItem { content: MAXPOOL2D.() -> Unit ) + /** + * Creates a 2D upsampling layer for increasing spatial resolution. + * + * @param scale Upsampling factors for height and width + * @param mode Interpolation mode (nearest default) + * @param alignCorners Alignment flag for bilinear mode (ignored for nearest) + * @param id Optional identifier for the layer + */ + public fun upsample2d( + scale: Pair = 2 to 2, + mode: UpsampleMode = UpsampleMode.Nearest, + alignCorners: Boolean = false, + id: String = "" + ) + + /** + * Creates a 2D upsampling layer with parameters configured in the DSL block. + */ + public fun upsample2d( + id: String = "", + content: UPSAMPLE2D.() -> Unit + ) + /** * Groups layers into a sequential block for better organization. * @@ -372,6 +397,16 @@ public interface MAXPOOL2D : NetworkDslItem { public fun padding(size: Int) } +@NetworkDsl +public interface UPSAMPLE2D : NetworkDslItem { + public var scale: Pair + public var mode: UpsampleMode + public var alignCorners: Boolean + + // Helper setter to allow concise Int-based configuration in DSL blocks + public fun scale(factor: Int) +} + /** * Scope for weights initialization with implicit shape context. @@ -674,6 +709,33 @@ public class MaxPool2dImpl( } } +public class Upsample2dImpl( + override val executionContext: ExecutionContext, + initialScale: Pair, + initialMode: UpsampleMode, + initialAlignCorners: Boolean, + private val id: String, +) : UPSAMPLE2D { + + override var scale: Pair = initialScale + override var mode: UpsampleMode = initialMode + override var alignCorners: Boolean = initialAlignCorners + + override fun scale(factor: Int) { + this.scale = factor to factor + } + + public fun create(): Upsample2d { + require(scale.first > 0 && scale.second > 0) { "Upsample2d scale must be > 0. Set it in the DSL block." } + return Upsample2d( + scale = scale, + mode = mode, + alignCorners = alignCorners, + name = getDefaultName(id, "Upsample2d", 0) + ) + } +} + // Stage implementation public class StageImpl( @@ -924,6 +986,32 @@ public class StageImpl( modules += impl.create() } + override fun upsample2d( + scale: Pair, + mode: UpsampleMode, + alignCorners: Boolean, + id: String + ) { + modules += Upsample2d( + scale = scale, + mode = mode, + alignCorners = alignCorners, + name = getDefaultName(id, "Upsample2d", modules.size) + ) + } + + override fun upsample2d(id: String, content: UPSAMPLE2D.() -> Unit) { + val impl = Upsample2dImpl( + executionContext = executionContext, + initialScale = 2 to 2, + initialMode = UpsampleMode.Nearest, + initialAlignCorners = false, + id = getDefaultName(id, "Upsample2d", modules.size) + ) + impl.content() + modules += impl.create() + } + override fun softmax(dim: Int, id: String) { modules += Softmax(dim, getDefaultName(id, "Softmax", modules.size)) // Softmax does not change feature dimension @@ -1177,6 +1265,35 @@ public class NeuralNetworkDslImpl( modules.add(impl.create()) } + override fun upsample2d( + scale: Pair, + mode: UpsampleMode, + alignCorners: Boolean, + id: String + ) { + modules += Upsample2d( + scale = scale, + mode = mode, + alignCorners = alignCorners, + name = getDefaultName(id, "Upsample2d", modules.size) + ) + } + + override fun upsample2d( + id: String, + content: UPSAMPLE2D.() -> Unit + ) { + val impl = Upsample2dImpl( + executionContext = executionContext, + initialScale = 2 to 2, + initialMode = UpsampleMode.Nearest, + initialAlignCorners = false, + id = getDefaultName(id, "Upsample2d", modules.size) + ) + impl.content() + modules.add(impl.create()) + } + override fun softmax(dim: Int, id: String) { modules += Softmax(dim, getDefaultName(id, "Softmax", modules.size)) // Softmax does not change feature dimension diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorExtensions.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorExtensions.kt index e90464f7..ee058092 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorExtensions.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorExtensions.kt @@ -1,6 +1,7 @@ package sk.ainet.lang.tensor import sk.ainet.lang.types.DType +import sk.ainet.lang.tensor.ops.UpsampleMode // Tensor extension functions that delegate to the ops component public fun Tensor.t(): Tensor = ops.transpose(this) @@ -26,5 +27,10 @@ public fun Tensor.mean(dim: Int? = null): Tensor = op public fun Tensor.variance(dim: Int? = null): Tensor = ops.variance(this, dim) public fun Tensor.sqrt(): Tensor = ops.sqrt(this) public fun Tensor.tril(k: Int = 0): Tensor = ops.tril(this, k) +public fun Tensor.upsample2d( + scale: Pair = 2 to 2, + mode: UpsampleMode = UpsampleMode.Nearest, + alignCorners: Boolean = false +): Tensor = ops.upsample2d(this, scale, mode, alignCorners) -// Global matmul function for the Linear layer usage pattern (removed due to duplicate with extension function) \ No newline at end of file +// Global matmul function for the Linear layer usage pattern (removed due to duplicate with extension function) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt index d18b77c1..6a611e68 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt @@ -315,6 +315,49 @@ public class MaxPool2dOperation( override fun clone(newParameters: Map): Operation = MaxPool2dOperation(newParameters) } +public class Upsample2dOperation( + parameters: Map = emptyMap() +) : BaseOperation("upsample2d", "nn", parameters) { + + override fun execute(inputs: List>): List> { + require(inputs.size == 1) { "Upsample2d operation requires exactly 1 input" } + throw UnsupportedOperationException("Direct execution not supported in graph mode") + } + + override fun validateInputs(inputs: List): ValidationResult { + if (inputs.size != 1) { + return ValidationResult.Invalid(listOf("Upsample2d operation requires exactly 1 input, got ${inputs.size}")) + } + return ValidationResult.Valid + } + + override fun inferOutputs(inputs: List): List { + require(inputs.size == 1) { "Upsample2d operation requires exactly 1 input" } + val inputShape = inputs[0].shape + val scale = (parameters["scale"] as? List<*>)?.mapNotNull { (it as? Number)?.toInt() } + val outputShape = if (inputShape != null && inputShape.size == 4 && scale != null && scale.size == 2) { + listOf( + inputShape[0], + inputShape[1], + inputShape[2] * scale[0], + inputShape[3] * scale[1] + ) + } else { + inputShape + } + return listOf( + TensorSpec( + name = "upsample2d_output", + shape = outputShape, + dtype = inputs[0].dtype, + requiresGrad = inputs[0].requiresGrad + ) + ) + } + + override fun clone(newParameters: Map): Operation = Upsample2dOperation(newParameters) +} + /** * Shape operations */ @@ -569,4 +612,4 @@ public class UnsqueezeOperation( } override fun clone(newParameters: Map): Operation = UnsqueezeOperation(newParameters) -} \ No newline at end of file +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt index ea041f19..6d61c62c 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt @@ -34,6 +34,12 @@ public interface TensorOps { stride: Pair = kernelSize, padding: Pair = 0 to 0 ): Tensor + public fun upsample2d( + input: Tensor, + scale: Pair, + mode: UpsampleMode = UpsampleMode.Nearest, + alignCorners: Boolean = false + ): Tensor // Shape operations public fun reshape(tensor: Tensor, newShape: Shape): Tensor @@ -66,4 +72,4 @@ public interface TensorOps { tensor: Tensor, targetType: TTo ): Tensor -} \ No newline at end of file +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/UpsampleMode.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/UpsampleMode.kt new file mode 100644 index 00000000..5819bbf9 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/UpsampleMode.kt @@ -0,0 +1,12 @@ +package sk.ainet.lang.tensor.ops + +/** + * Interpolation modes for 2D upsampling. + * + * YOLO-style models typically use nearest-neighbor. Bilinear is provided for parity + * but may ignore `alignCorners` when not applicable. + */ +public enum class UpsampleMode { + Nearest, + Bilinear +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt index c59eb343..cbeb3952 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt @@ -5,6 +5,7 @@ import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.VoidOpsTensor import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.tensor.ops.UpsampleMode import sk.ainet.lang.types.DType public class VoidTensorOps : TensorOps { @@ -142,6 +143,17 @@ public class VoidTensorOps : TensorOps { return VoidOpsTensor(resultData, input.dtype) } + override fun upsample2d( + input: Tensor, + scale: Pair, + mode: UpsampleMode, + alignCorners: Boolean + ): Tensor { + val resultShape = calculateUpsample2dShape(input.shape, scale) + val resultData = dataFactory.zeros(resultShape, input.dtype) + return VoidOpsTensor(resultData, input.dtype) + } + override fun reshape(tensor: Tensor, newShape: Shape): Tensor { val resultShape = calculateReshapeTargetShape(tensor.shape, newShape) val resultData = dataFactory.zeros(resultShape, tensor.dtype) @@ -637,6 +649,21 @@ public class VoidTensorOps : TensorOps { return result } + /** + * Calculates the result shape for 2D upsampling. + */ + private fun calculateUpsample2dShape(shape: Shape, scale: Pair): Shape { + require(shape.rank == 4) { + "Upsample2d expects 4D input (N, C, H, W), got ${shape.dimensions.contentToString()}" + } + val (scaleH, scaleW) = scale + require(scaleH > 0 && scaleW > 0) { "Upsample2d scale factors must be positive" } + val dims = shape.dimensions.copyOf() + dims[2] = dims[2] * scaleH + dims[3] = dims[3] * scaleW + return Shape(dims) + } + /** * Calculates the result shape for squeeze operation */ diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/Upsample2dVoidOpsShapeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/Upsample2dVoidOpsShapeTest.kt new file mode 100644 index 00000000..7b2f0056 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/Upsample2dVoidOpsShapeTest.kt @@ -0,0 +1,32 @@ +package sk.ainet.lang.tensor.ops + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.types.FP32 + +class Upsample2dVoidOpsShapeTest { + private val dataFactory = DenseTensorDataFactory() + private val ops = VoidTensorOps() + + private fun tensor(shape: Shape): VoidOpsTensor = + VoidOpsTensor(dataFactory.zeros(shape, FP32::class), FP32::class) + + @Test + fun upsample2d_scales_spatial_dims() { + val x = tensor(Shape(1, 3, 10, 20)) + val y = ops.upsample2d(x, scale = 2 to 3, mode = UpsampleMode.Nearest, alignCorners = false) + assertEquals(listOf(1, 3, 20, 60), y.shape.dimensions.toList()) + } + + @Test + fun upsample2d_invalid_rank_throws() { + val x = tensor(Shape(1, 3, 10)) + assertFailsWith { + ops.upsample2d(x, scale = 2 to 2, mode = UpsampleMode.Nearest, alignCorners = false) + } + } +} diff --git a/skainet-lang/skainet-lang-models/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/Yolo8.kt b/skainet-lang/skainet-lang-models/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/Yolo8.kt deleted file mode 100644 index 48276b25..00000000 --- a/skainet-lang/skainet-lang-models/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/Yolo8.kt +++ /dev/null @@ -1,54 +0,0 @@ -package sk.ainet.lang.model.dnn.yolo - -import sk.ainet.context.ExecutionContext -import sk.ainet.lang.model.Model -import sk.ainet.lang.model.ModelCard -import sk.ainet.lang.nn.Module -import sk.ainet.lang.nn.definition -import sk.ainet.lang.nn.network -import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.types.FP32 - -public class Yolo8 : Model, Tensor> { - - private fun buildModel(executionContext: ExecutionContext): Module = definition { - network(executionContext) { - // Placeholder input definition; actual YOLOv8 graph to be wired when available - input(1) - } - } - - // Backward-compatible helper for old call sites - public fun model(executionContext: ExecutionContext): Module = create(executionContext) - - override fun create(executionContext: ExecutionContext): Module = buildModel(executionContext) - - override suspend fun calculate( - module: Module, - inputValue: Tensor, - executionContext: ExecutionContext, - reportProgress: suspend (current: Int, total: Int, message: String?) -> Unit - ): Tensor { - reportProgress(0, 1, "starting yolo8 forward") - val out = module.forward(inputValue, executionContext) - reportProgress(1, 1, "done") - return out - } - - override fun modelCard(): ModelCard { - return ModelCard( - license = "apache-2.0", - libraryName = "skainet", - pipelineTag = "object-detection", - language = listOf("en"), - modalities = listOf("vision"), - baseModel = "yolov8", - contextLength = 0, - datasets = emptyList(), - metrics = emptyList(), - modelIndex = emptyList(), - intendedUse = "YOLOv8-style object detection placeholder implemented with new Model API.", - limitations = "Architecture is a placeholder; no pretrained weights or real detection head wired yet." - ) - } -} \ No newline at end of file diff --git a/skainet-models/skainet-model-yolo/build.gradle.kts b/skainet-models/skainet-model-yolo/build.gradle.kts new file mode 100644 index 00000000..8946c1ab --- /dev/null +++ b/skainet-models/skainet-model-yolo/build.gradle.kts @@ -0,0 +1,66 @@ +import org.jetbrains.kotlin.gradle.ExperimentalKotlinGradlePluginApi +import org.jetbrains.kotlin.gradle.ExperimentalWasmDsl +import org.jetbrains.kotlin.gradle.dsl.JvmTarget + +plugins { + alias(libs.plugins.kotlinMultiplatform) + alias(libs.plugins.androidLibrary) + alias(libs.plugins.vanniktech.mavenPublish) + alias(libs.plugins.kover) + alias(libs.plugins.binary.compatibility.validator) +} + +kotlin { + explicitApi() + + androidTarget { + @OptIn(ExperimentalKotlinGradlePluginApi::class) + compilerOptions { + jvmTarget.set(JvmTarget.JVM_11) + } + } + + iosArm64() + iosSimulatorArm64() + macosArm64 () + linuxX64 () + linuxArm64 () + + jvm() + + js { + browser() + } + + @OptIn(ExperimentalWasmDsl::class) + wasmJs { + browser() + } + + sourceSets { + commonMain.dependencies { + implementation(project(":skainet-lang:skainet-lang-core")) + implementation(project(":skainet-lang:skainet-lang-models")) + implementation(libs.kotlinx.coroutines) + } + + commonTest.dependencies { + implementation(libs.kotlin.test) + implementation(project(":skainet-backends:skainet-backend-cpu")) + + } + } +} + +android { + namespace = "sk.ainet.lang.model.yolo" + compileSdk = libs.versions.android.compileSdk.get().toInt() + + defaultConfig { + minSdk = libs.versions.android.minSdk.get().toInt() + } + compileOptions { + sourceCompatibility = JavaVersion.VERSION_11 + targetCompatibility = JavaVersion.VERSION_11 + } +} \ No newline at end of file diff --git a/skainet-models/skainet-model-yolo/gradle.properties b/skainet-models/skainet-model-yolo/gradle.properties new file mode 100644 index 00000000..cd80139a --- /dev/null +++ b/skainet-models/skainet-model-yolo/gradle.properties @@ -0,0 +1,2 @@ +POM_ARTIFACT_ID=skainet-model-yolo +POM_NAME=skainet neural network scripting API \ No newline at end of file diff --git a/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/Yolo8.kt b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/Yolo8.kt new file mode 100644 index 00000000..5b115d31 --- /dev/null +++ b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/Yolo8.kt @@ -0,0 +1,98 @@ +package sk.ainet.lang.model.dnn.yolo + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.model.Model +import sk.ainet.lang.model.ModelCard +import sk.ainet.lang.nn.Module +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 + +public class Yolo8( + private val config: YoloConfig = YoloConfig() +) : Model, Tensor> { + + private fun buildModel(executionContext: ExecutionContext): Module = + Yolo8Module(executionContext, config) + + // Backward-compatible helper for old call sites + public fun model(executionContext: ExecutionContext): Module = create(executionContext) + + override fun create(executionContext: ExecutionContext): Module = buildModel(executionContext) + + override suspend fun calculate( + module: Module, + inputValue: Tensor, + executionContext: ExecutionContext, + reportProgress: suspend (current: Int, total: Int, message: String?) -> Unit + ): Tensor { + reportProgress(0, 1, "starting yolo8 forward") + val out = module.forward(inputValue, executionContext) + reportProgress(1, 1, "done") + return out + } + + /** + * Runs the YOLO graph and returns all head outputs (small/medium/large) for downstream post-processing. + */ + public fun calculateHeads( + module: Module, + inputValue: Tensor, + executionContext: ExecutionContext + ): HeadOutputs { + require(module is Yolo8Module) { "calculateHeads expects a Yolo8Module instance" } + return module.forwardHeads(inputValue, executionContext) + } + + /** + * Full inference: forward pass + decode + NMS. Expects the input tensor to be + * preprocessed to the configured model size (see [YoloPreprocess]). + */ + public suspend fun infer( + module: Module, + input: YoloInput, + executionContext: ExecutionContext, + reportProgress: suspend (current: Int, total: Int, message: String?) -> Unit = { _, _, _ -> } + ): List { + require(module is Yolo8Module) { "infer expects a Yolo8Module instance" } + reportProgress(0, 2, "yolo forward") + val heads = module.forwardHeads(input.tensor, executionContext) + reportProgress(1, 2, "decode") + val detections = YoloDecoder.decode(heads, config, input) + reportProgress(2, 2, "done") + return detections + } + + override fun modelCard(): ModelCard { + return ModelCard( + license = "apache-2.0", + libraryName = "skainet", + pipelineTag = "object-detection", + language = listOf("en"), + modalities = listOf("vision"), + baseModel = "yolov8", + contextLength = 0, + datasets = emptyList(), + metrics = emptyList(), + modelIndex = emptyList(), + intendedUse = "YOLOv8-style object detection skeleton using SKaiNET DSL blocks.", + limitations = "Weights and image preprocessing are not included; outputs depend on provided weights and input normalization." + ) + } +} + +private class Yolo8Module( + private val executionContext: ExecutionContext, + private val config: YoloConfig, override val name: String = "Yolo8" +) : Module() { + + private val graph = Yolo8Graph(executionContext, config) + + override val modules: List> + get() = graph.modules() + + override fun forward(input: Tensor, ctx: ExecutionContext): Tensor = + graph.forward(input, ctx).large + + fun forwardHeads(input: Tensor, ctx: ExecutionContext): HeadOutputs = + graph.forward(input, ctx) +} diff --git a/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/Yolo8Graph.kt b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/Yolo8Graph.kt new file mode 100644 index 00000000..a158bd31 --- /dev/null +++ b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/Yolo8Graph.kt @@ -0,0 +1,396 @@ +package sk.ainet.lang.model.dnn.yolo + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.definition +import sk.ainet.lang.nn.network +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.ops.UpsampleMode +import sk.ainet.lang.tensor.silu +import sk.ainet.lang.types.FP32 + +/** + * Minimal YOLOv8-style graph assembled from reusable blocks. + * + * This is intentionally compact: it wires the main backbone, a lightweight neck + * with upsampling/concats, and three detection heads. The output currently returns + * the largest-scale head; smaller heads are kept for future concatenation when + * post-processing is added. + */ +internal class Yolo8Graph(private val executionContext: ExecutionContext, private val config: YoloConfig) { + + private val stem = ConvBnSiLU( + inChannels = 3, + outChannels = 32, + kernel = 3, + stride = 2, + padding = 1, + name = "stem", + executionContext = executionContext + ) + + private val stage1 = ConvBnSiLU( + inChannels = 32, + outChannels = 64, + kernel = 3, + stride = 2, + padding = 1, + name = "stage1_down", + executionContext = executionContext + ) + private val c2f1 = C2f( + inChannels = 64, + outChannels = 64, + bottlenecks = 1, + name = "c2f1", + executionContext = executionContext + ) + + private val stage2 = ConvBnSiLU( + inChannels = 64, + outChannels = 128, + kernel = 3, + stride = 2, + padding = 1, + name = "stage2_down", + executionContext = executionContext + ) + private val c2f2 = C2f( + inChannels = 128, + outChannels = 128, + bottlenecks = 1, + name = "c2f2", + executionContext = executionContext + ) + + private val stage3 = ConvBnSiLU( + inChannels = 128, + outChannels = 256, + kernel = 3, + stride = 2, + padding = 1, + name = "stage3_down", + executionContext = executionContext + ) + private val c2f3 = C2f( + inChannels = 256, + outChannels = 256, + bottlenecks = 1, + name = "c2f3", + executionContext = executionContext + ) + + private val stage4 = ConvBnSiLU( + inChannels = 256, + outChannels = 512, + kernel = 3, + stride = 2, + padding = 1, + name = "stage4_down", + executionContext = executionContext + ) + private val c2f4 = C2f( + inChannels = 512, + outChannels = 512, + bottlenecks = 1, + name = "c2f4", + executionContext = executionContext + ) + + private val sppf = Sppf( + channels = 512, + name = "sppf", + executionContext = executionContext + ) + + // Neck + private val up1 = UpsampleWrapper(scale = 2 to 2, name = "up1", executionContext = executionContext) + private val c2fNeck1 = C2f( + inChannels = 512 + 256, + outChannels = 256, + bottlenecks = 1, + name = "c2f_neck1", + executionContext = executionContext + ) + + private val up2 = UpsampleWrapper(scale = 2 to 2, name = "up2", executionContext = executionContext) + private val c2fNeck2 = C2f( + inChannels = 256 + 128, + outChannels = 128, + bottlenecks = 1, + name = "c2f_neck2", + executionContext = executionContext + ) + + // Heads (three scales) + private val headSmall = DetectHead( + inChannels = 128, + outChannels = config.numClasses + 5, + name = "detect_small", + executionContext = executionContext + ) + private val headMedium = DetectHead( + inChannels = 256, + outChannels = config.numClasses + 5, + name = "detect_medium", + executionContext = executionContext + ) + private val headLarge = DetectHead( + inChannels = 512, + outChannels = config.numClasses + 5, + name = "detect_large", + executionContext = executionContext + ) + + fun modules(): List> = listOf( + stem.module, + stage1.module, c2f1, + stage2.module, c2f2, + stage3.module, c2f3, + stage4.module, c2f4, + sppf, + up1, + c2fNeck1, + up2, + c2fNeck2, + headSmall, + headMedium, + headLarge + ) + + fun forward(input: Tensor, ctx: ExecutionContext): HeadOutputs { + val x0 = stem.forward(input, ctx) + val x1 = c2f1.forward(stage1.forward(x0, ctx), ctx) // 1/4 + val x2 = c2f2.forward(stage2.forward(x1, ctx), ctx) // 1/8 + val x3 = c2f3.forward(stage3.forward(x2, ctx), ctx) // 1/16 + val x4 = c2f4.forward(stage4.forward(x3, ctx), ctx) // 1/32 + val p5 = sppf.forward(x4, ctx) + + // Neck path + val p4 = c2fNeck1.forward( + concatAlongChannels( + up1.forward(p5, ctx), + x3, + ctx + ), + ctx + ) + val p3 = c2fNeck2.forward( + concatAlongChannels( + up2.forward(p4, ctx), + x2, + ctx + ), + ctx + ) + + // Heads: currently return the large-scale output; smaller scales stay available for later post-processing wiring. + val small = headSmall.forward(p3, ctx) + val medium = headMedium.forward(p4, ctx) + val large = headLarge.forward(p5, ctx) + return HeadOutputs(small = small, medium = medium, large = large) + } + + private fun concatAlongChannels( + a: Tensor, + b: Tensor, + ctx: ExecutionContext + ): Tensor = ctx.ops.concat(listOf(a, b), dim = 1) +} + +internal class ConvBnSiLU( + private val inChannels: Int, + private val outChannels: Int, + private val kernel: Int, + private val stride: Int, + private val padding: Int, + private val name: String, + executionContext: ExecutionContext +) { + val module: Module = definition { + network(executionContext) { + sequential { + stage(name) { + conv2d( + outChannels = outChannels, + kernelSize = kernel to kernel, + stride = stride to stride, + padding = padding to padding + ) { + this.inChannels = this@ConvBnSiLU.inChannels + } + batchNorm(numFeatures = outChannels) + activation { tensor -> tensor.silu() } + } + } + } + } + + fun forward(input: Tensor, ctx: ExecutionContext): Tensor = + module.forward(input, ctx) +} + +internal class C2f( + private val inChannels: Int, + private val outChannels: Int, + private val bottlenecks: Int, + override val name: String, + private val executionContext: ExecutionContext +) : Module() { + + private val hidden = outChannels / 2 + private val reduce = ConvBnSiLU( + inChannels = inChannels, + outChannels = hidden * 2, + kernel = 1, + stride = 1, + padding = 0, + name = "${name}_reduce", + executionContext = executionContext + ) + private val bottleneckBlocks = List(bottlenecks) { idx -> + ConvBnSiLU( + inChannels = hidden * 2, + outChannels = hidden * 2, + kernel = 3, + stride = 1, + padding = 1, + name = "${name}_b$idx", + executionContext = executionContext + ) + } + private val expand = ConvBnSiLU( + inChannels = hidden * (2 + bottlenecks), + outChannels = outChannels, + kernel = 1, + stride = 1, + padding = 0, + name = "${name}_expand", + executionContext = executionContext + ) + + override val modules: List> = buildList { + add(reduce.module) + bottleneckBlocks.forEach { add(it.module) } + add(expand.module) + } + + override fun forward(input: Tensor, ctx: ExecutionContext): Tensor { + val y0 = reduce.forward(input, ctx) + val parts = mutableListOf(y0) + var last = y0 + bottleneckBlocks.forEach { block -> + last = block.forward(last, ctx) + parts += last + } + val concat = ctx.ops.concat(parts, dim = 1) + return expand.forward(concat, ctx) + } +} + +public data class HeadOutputs( + val small: Tensor, + val medium: Tensor, + val large: Tensor +) + +internal class Sppf( + private val channels: Int, + override val name: String, + private val executionContext: ExecutionContext +) : Module() { + + private val reduce = ConvBnSiLU( + inChannels = channels, + outChannels = channels, + kernel = 1, + stride = 1, + padding = 0, + name = "${name}_reduce", + executionContext = executionContext + ) + private val expand = ConvBnSiLU( + inChannels = channels * 4, + outChannels = channels, + kernel = 1, + stride = 1, + padding = 0, + name = "${name}_expand", + executionContext = executionContext + ) + + override val modules: List> = listOf(reduce.module, expand.module) + + override fun forward(input: Tensor, ctx: ExecutionContext): Tensor { + val x = reduce.forward(input, ctx) + val p1 = ctx.ops.maxPool2d(x, kernelSize = 5 to 5, stride = 1 to 1, padding = 2 to 2) + val p2 = ctx.ops.maxPool2d(p1, kernelSize = 5 to 5, stride = 1 to 1, padding = 2 to 2) + val p3 = ctx.ops.maxPool2d(p2, kernelSize = 5 to 5, stride = 1 to 1, padding = 2 to 2) + val concat = ctx.ops.concat(listOf(x, p1, p2, p3), dim = 1) + return expand.forward(concat, ctx) + } +} + +internal class UpsampleWrapper( + private val scale: Pair, + override val name: String, + private val executionContext: ExecutionContext +) : Module() { + private val module: Module = definition { + network(executionContext) { + sequential { + stage(name) { + upsample2d(scale = scale, mode = UpsampleMode.Nearest) + } + } + } + } + + override val modules: List> = listOf(module) + + override fun forward(input: Tensor, ctx: ExecutionContext): Tensor = + module.forward(input, ctx) +} + +internal class DetectHead( + private val inChannels: Int, + private val outChannels: Int, + override val name: String, + private val executionContext: ExecutionContext +) : Module() { + + private val stem = ConvBnSiLU( + inChannels = inChannels, + outChannels = inChannels, + kernel = 3, + stride = 1, + padding = 1, + name = "${name}_stem", + executionContext = executionContext + ) + + private val proj: Module = definition { + network(executionContext) { + sequential { + stage("${name}_proj") { + conv2d( + outChannels = outChannels, + kernelSize = 1 to 1, + stride = 1 to 1, + padding = 0 to 0 + ) { + this.inChannels = inChannels + this.bias = true + } + } + } + } + } + + override val modules: List> = listOf(stem.module, proj) + + override fun forward(input: Tensor, ctx: ExecutionContext): Tensor { + val h = stem.forward(input, ctx) + return proj.forward(h, ctx) + } +} diff --git a/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/YoloConfig.kt b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/YoloConfig.kt new file mode 100644 index 00000000..1e1e9fde --- /dev/null +++ b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/YoloConfig.kt @@ -0,0 +1,13 @@ +package sk.ainet.lang.model.dnn.yolo + +/** + * YOLO configuration used for decoding and thresholds. + */ +public data class YoloConfig( + val numClasses: Int = 80, + val inputSize: Int = 640, + val confThreshold: Float = 0.25f, + val iouThreshold: Float = 0.45f, + val maxDetections: Int = 300, + val classNames: List = emptyList() +) diff --git a/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/YoloPost.kt b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/YoloPost.kt new file mode 100644 index 00000000..8b8fce4a --- /dev/null +++ b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/YoloPost.kt @@ -0,0 +1,178 @@ +package sk.ainet.lang.model.dnn.yolo + +import kotlin.math.exp +import kotlin.math.max +import kotlin.math.min +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 + +public data class Detection( + val box: Box, + val score: Float, + val classId: Int, + val label: String? = null +) + +public data class Box( + val x1: Float, + val y1: Float, + val x2: Float, + val y2: Float +) { + val width: Float get() = max(0f, x2 - x1) + val height: Float get() = max(0f, y2 - y1) +} + +public data class YoloInput( + val tensor: Tensor, + val originalWidth: Int, + val originalHeight: Int, + val letterboxScale: Float = 1f, + val padW: Int = 0, + val padH: Int = 0 +) + +internal object YoloDecoder { + fun decode( + heads: HeadOutputs, + config: YoloConfig, + inputMeta: YoloInput + ): List { + val small = decodeHead(heads.small, stride = 8, config, inputMeta) + val medium = decodeHead(heads.medium, stride = 16, config, inputMeta) + val large = decodeHead(heads.large, stride = 32, config, inputMeta) + return nms(small + medium + large, config) + } + + private fun decodeHead( + head: Tensor, + stride: Int, + config: YoloConfig, + inputMeta: YoloInput + ): List { + val (batch, channels, h, w) = head.shape.dimensions + require(batch == 1) { "Only batch size 1 is supported for decode, got $batch" } + val classes = config.numClasses + require(channels == classes + 5) { + "Head channels ($channels) must equal classes+5 (${classes + 5})" + } + val out = mutableListOf() + var y = 0 + while (y < h) { + var x = 0 + while (x < w) { + val base = floatArrayOf( + head.data.get(0, 0, y, x), + head.data.get(0, 1, y, x), + head.data.get(0, 2, y, x), + head.data.get(0, 3, y, x), + head.data.get(0, 4, y, x) + ) + val obj = sigmoid(base[4]) + if (obj >= config.confThreshold) { + val clsScore = argMaxClass(head, y, x, classes) + val score = obj * clsScore.score + if (score >= config.confThreshold) { + val box = decodeBox(base, x, y, stride, inputMeta) + val label = config.classNames.getOrNull(clsScore.classId) + out += Detection(box, score, clsScore.classId, label) + } + } + x++ + } + y++ + } + return out + } + + private data class ClassScore(val classId: Int, val score: Float) + + private fun argMaxClass( + head: Tensor, + y: Int, + x: Int, + classes: Int + ): ClassScore { + var bestScore = Float.NEGATIVE_INFINITY + var bestId = -1 + var c = 0 + while (c < classes) { + val v = head.data.get(0, 5 + c, y, x) + if (v > bestScore) { + bestScore = v + bestId = c + } + c++ + } + return ClassScore(bestId, sigmoid(bestScore)) + } + + private fun decodeBox( + raw: FloatArray, + gridX: Int, + gridY: Int, + stride: Int, + meta: YoloInput + ): Box { + val x = (sigmoid(raw[0]) * 2f - 0.5f + gridX) * stride + val y = (sigmoid(raw[1]) * 2f - 0.5f + gridY) * stride + val w = (sigmoid(raw[2]) * 2f).let { it * it } * stride + val h = (sigmoid(raw[3]) * 2f).let { it * it } * stride + val halfW = w / 2f + val halfH = h / 2f + val lx = x - halfW + val ty = y - halfH + val rx = x + halfW + val by = y + halfH + + // Remove letterbox padding and rescale to original dimensions + val scale = meta.letterboxScale + val padW = meta.padW + val padH = meta.padH + val mappedX1 = (lx - padW) / scale + val mappedY1 = (ty - padH) / scale + val mappedX2 = (rx - padW) / scale + val mappedY2 = (by - padH) / scale + + val x1 = mappedX1.coerceIn(0f, meta.originalWidth.toFloat()) + val y1 = mappedY1.coerceIn(0f, meta.originalHeight.toFloat()) + val x2 = mappedX2.coerceIn(0f, meta.originalWidth.toFloat()) + val y2 = mappedY2.coerceIn(0f, meta.originalHeight.toFloat()) + return Box(x1, y1, x2, y2) + } + + private fun sigmoid(x: Float): Float = (1f / (1f + exp(-x))) + + private fun nms(detections: List, config: YoloConfig): List { + if (detections.isEmpty()) return emptyList() + val sorted = detections.sortedByDescending { it.score }.toMutableList() + val keep = mutableListOf() + while (sorted.isNotEmpty() && keep.size < config.maxDetections) { + val best = sorted.removeAt(0) + keep += best + val iterator = sorted.iterator() + while (iterator.hasNext()) { + val other = iterator.next() + if (best.classId == other.classId) { + val iou = iou(best.box, other.box) + if (iou > config.iouThreshold) { + iterator.remove() + } + } + } + } + return keep + } + + private fun iou(a: Box, b: Box): Float { + val interX1 = max(a.x1, b.x1) + val interY1 = max(a.y1, b.y1) + val interX2 = min(a.x2, b.x2) + val interY2 = min(a.y2, b.y2) + val interW = max(0f, interX2 - interX1) + val interH = max(0f, interY2 - interY1) + val interArea = interW * interH + val union = a.width * a.height + b.width * b.height - interArea + return if (union <= 0f) 0f else interArea / union + } +} diff --git a/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/YoloPreprocess.kt b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/YoloPreprocess.kt new file mode 100644 index 00000000..d0a68e1f --- /dev/null +++ b/skainet-models/skainet-model-yolo/src/commonMain/kotlin/sk/ainet/lang/model/dnn/yolo/YoloPreprocess.kt @@ -0,0 +1,52 @@ +package sk.ainet.lang.model.dnn.yolo + +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 + +/** + * Preprocessing helpers for YOLOv8 inputs. + * + * For now we assume the caller supplies an already resized and normalized NCHW tensor + * matching [YoloConfig.inputSize]. This keeps dependencies minimal and avoids + * platform-specific image handling. Letterbox metadata is carried to map boxes + * back to the original image size. + */ +public object YoloPreprocess { + + /** + * Wraps a tensor that is already in the expected model shape [1, 3, inputSize, inputSize]. + * + * @param tensor normalized tensor (values typically in [0,1]) + * @param originalWidth width of the source image + * @param originalHeight height of the source image + * @param inputSize target model resolution (defaults to config.inputSize) + * @param padW horizontal padding applied during letterbox (pixels) + * @param padH vertical padding applied during letterbox (pixels) + * @param scale scaling factor used during letterbox (original -> model space) + */ + public fun fromReadyTensor( + tensor: Tensor, + originalWidth: Int, + originalHeight: Int, + inputSize: Int, + padW: Int = 0, + padH: Int = 0, + scale: Float = 1f + ): YoloInput { + require(tensor.shape.rank == 4) { "Expected NCHW tensor rank 4, got ${tensor.shape.rank}" } + require(tensor.shape[0] == 1 && tensor.shape[1] == 3) { + "Expected batch=1 and channels=3, got shape ${tensor.shape.dimensions.contentToString()}" + } + require(tensor.shape[2] == inputSize && tensor.shape[3] == inputSize) { + "Tensor must already be resized to $inputSize, got ${tensor.shape[2]}x${tensor.shape[3]}" + } + return YoloInput( + tensor = tensor, + originalWidth = originalWidth, + originalHeight = originalHeight, + letterboxScale = scale, + padW = padW, + padH = padH + ) + } +}