# teeML using Marlin's Oyster

## Introduction

In the realm of secure machine learning, the need for privacy and efficiency is paramount. Traditional methods like zkML (zero-knowledge machine learning) have provided a way to ensure privacy but often at the cost of increased response times. Enter teeML, a novel approach leveraging Trusted Execution Environments (TEEs) to deploy machine learning models. This method not only maitains sufficient security but also significantly reduces response times & preserves privacy. In this blog post we cover how we leveraged teeML for Hyle's [vibe check demo](https://vibe.hyle.eu/) & how you can also do the same with [Marlin's Oyster](https://www.marlin.org/oyster).

## What is Oyster?

Oyster is a platform by Marlin that allows developers to deploy applications within secure enclaves. These enclaves provide a trusted execution environment, ensuring that the code and data are protected from external tampering. For more details, you can refer to the [Oyster documentation](https://docs.marlin.org/learn/what-is-oyster).

## How to verify integrity of a teeML?
In zkML systems the integrity is verified by verifying the zk-proofs like STARK in the case of Hyle's vibe check demo. These proofs can be verified on the settelment layer or anyother place as per user's preference. Generating these proofs is computationally intensive & takes considerable resources.

teeML also lets you verify the integrity of the computations that happened inside the enclave. Encalves sign the response with the enclave's private key. This signature is called "Attestation". The attestation can be verified similar to a zk-proof. To read more about how attestation verification works read Marlin documentation [here](https://docs.marlin.org/learn/oyster/core-concepts/remote-attestations).

## Deploy your teeML - Step-by-Step Guide

### Train and Export the Model
Train your machine learning model and export it. For Hyle's Vibe check we have exported the model as a pickle (.pkl) file. This file will be used within the enclave to make predictions. You can train you own model, however [here](https://github.com/Hyle-org/onnx_cairo?tab=readme-ov-file#train-the-xgboost-model) is the python script to train XGBoost classifier used in vibe check demo.

Code snippet to export you model to pickle (.pkl) file is given below:
> Please note that the code given below is not complete.

import joblib

"""Train using builtin categorical data support from XGBoost"""

X_train, X_test, y_train, y_test = train_test_split(
    X, y, random_state=1994, test_size=0.2
)

clf = xgb.XGBClassifier(
    **params,
    eval_metric="auc",
    enable_categorical=True,
    max_cat_to_onehot=1,  # We use optimal partitioning exclusively
)

clf.fit(X, y, eval_set=[(X_test, y_test), (X_train, y_train)])

model_file = "my_model.pkl"

# Save the model
joblib.dump(clf, model_file)

### Setup the Enclave
Check out this [repo](https://github.com/marlinprotocol/hyle-teeML-setup/tree/master) for the code to setup the enclave. The setup includes installing necessary packages and copying the model and server scripts. Below is the Flask HTTP server wrapped around the model prediction logic. This server will handle incoming requests, make predictions using the model, and serve predictions along with the attestations.

import pandas as pd
import joblib
from flask import Flask, request, jsonify
from cryptography.hazmat.primitives.asymmetric import ed25519
from base64 import b64encode
from hashlib import sha256
from datetime import datetime

def prediction(input, model):
input_df = pd.DataFrame([input])
result = model.predict(input_df)
return result

# Load the local classifier
print("Loading XGBClassifier...")
model = joblib.load("serialized_optimized.pkl")

# Load the secret key from the id.sec file
with open("/app/id.sec", "rb") as key_file:
secret_key_bytes = key_file.read()

# The secret key consists of 32 bytes for the private scalar and 32 bytes for the public key.
private_scalar = secret_key_bytes[:32]
public_key_bytes = secret_key_bytes[32:]

# Create a private key object
private_key = ed25519.Ed25519PrivateKey.from_private_bytes(private_scalar)

app = Flask(__name__)

@app.route('/prediction', methods=['POST'])
def predict():
data = request.get_json()
image = data['image']
prediction_result = prediction(image, model)
image_hash = sha256(str(image).encode()).hexdigest()
timestamp = datetime.utcnow().isoformat()

data_to_sign = f"{str(prediction_result)}{image_hash}{timestamp}"

data_to_sign_bytes = data_to_sign.encode()
signature = private_key.sign(data_to_sign_bytes)
signature_base64 = b64encode(signature).decode('utf-8')
return jsonify({
'prediction': prediction_result.tolist(),
'attestation': signature_base64,
'image_identifier': image_hash,
'timestamp': timestamp
})

> An attestation for each prediction is generated by signing it with the enclave key. This ensures that the prediction is trustworthy and has not been tampered with.

### Build and Run the Enclave
Once the enclave setup is ready you just need to run a couple of commands to generate the enclave image file. The steps for the same are highligted [here](https://github.com/marlinprotocol/hyle-teeML-setup/tree/master?tab=readme-ov-file#build-the-enclave-builder-image). After the enclave image file is generated you can deploy your Oyster enclave by following the steps mentioned [here](https://docs.marlin.org/user-guides/oyster/instances/quickstart/deploy).

## Hyle Vibe Check Frontend

The Hyle [Vibe Check frontend](https://github.com/Hyle-org/vibe-check/tree/main/vibe-check-frontend) is a Vue.js application that was originally designed to use zkML for facial feature classification. By switching to teeML, we were able to significantly reduce the response times. Following changes were made:

### Switch from zkML to teeML
The vibe check frontend used Cairo wasm runner & Prover to fetch predictions & generate zk-proof of the predicitons respectively. This can be seen in the image below. The code for Cairo wasm runner & prover was removed from the frontend, instead teeML API is now used to fetch predictions & encalve attestations. This means that now instead of generating zk-proofs we use the attestations from the enclave.


### Image Processing
The frontend changed image processing to the RGB image necessarye to feed-in to the ML model for predictions. Below code snippet highlights the new image processing logic added to the frontend:

const checkVibe = async (image: ImageBitmap, zoomingPromise: Promise<any>, x: number, y: number, width: number, height: number) => {
    vibeCheckStatus.value = null;
    status.value = "checking_vibe";

    console.log("Face detection coordinates:", { x, y, width, height });

    const small = document.createElement("canvas");
    const smallCtx = small.getContext("2d")!;
    small.width = 48;
    small.height = 48;

    // Draw the entire face detection area
    smallCtx.drawImage(image, x, y, width, height, 0, 0, 48, 48);
    
    const imageData = smallCtx.getImageData(0, 0, 48, 48);
    const processedData = new Float32Array(48 * 48);

    // Apply Gaussian blur to reduce noise
    const blurredData = new Float32Array(48 * 48);
    const kernel = [
        [1, 2, 1],
        [2, 4, 2],
        [1, 2, 1]
    ]; // Gaussian kernel for better blurring

    for (let y = 0; y < 48; y++) {
        for (let x = 0; x < 48; x++) {
            let sum = 0;
            let count = 0;
            for (let dy = -1; dy <= 1; dy++) {
                for (let dx = -1; dx <= 1; dx++) {
                    const nx = x + dx;
                    const ny = y + dy;
                    if (nx >= 0 && nx < 48 && ny >= 0 && ny < 48) {
                        const idx = (ny * 48 + nx) * 4;
                        const pixelValue = (imageData.data[idx] + imageData.data[idx + 1] + imageData.data[idx + 2]) / 3;
                        sum += pixelValue * kernel[dy + 1][dx + 1]; // Apply kernel weight
                        count += kernel[dy + 1][dx + 1];
                    }
                }
            }
            blurredData[y * 48 + x] = sum / count; // Average for blur
        }
    }

    // Adaptive thresholding
    const blockSize = 5; // Increased block size for better averaging
    const C = 5; // Increased constant to reduce false positives

    for (let y = 0; y < 48; y++) {
        for (let x = 0; x < 48; x++) {
            let sum = 0;
            let count = 0;
            for (let dy = -blockSize; dy <= blockSize; dy++) {
                for (let dx = -blockSize; dx <= blockSize; dx++) {
                    const nx = x + dx;
                    const ny = y + dy;
                    if (nx >= 0 && nx < 48 && ny >= 0 && ny < 48) {
                        sum += blurredData[ny * 48 + nx];
                        count++;
                    }
                }
            }
            const threshold = sum / count - C;
            const pixelValue = blurredData[y * 48 + x];
            processedData[y * 48 + x] = pixelValue < threshold ? 1 : 0;
        }
    }

    // Focus on the mouth and lip region
    const mouthRegion = new Float32Array(48 * 48);
    const mouthY = Math.floor(height * 0.55); // Adjusted for better mouth detection
    const mouthHeight = Math.floor(height * 0.3); // Increased height of the mouth region

    for (let y = mouthY; y < mouthY + mouthHeight; y++) {
        for (let x = 0; x < 48; x++) {
            const index = y * 48 + x;
            if (index < mouthRegion.length) {
                mouthRegion[index] = processedData[index]; // Copy the processed data for the mouth region
            }
        }
    }

    // Morphological opening to reduce black spots
    const openedData = new Float32Array(48 * 48);
    for (let y = 0; y < 48; y++) {
        for (let x = 0; x < 48; x++) {
            if (mouthRegion[y * 48 + x] === 1) {
                openedData[y * 48 + x] = 1; // Keep the pixel
            } else {
                // Check surrounding pixels
                let isSurrounded = false;
                for (let dy = -1; dy <= 1; dy++) {
                    for (let dx = -1; dx <= 1; dx++) {
                        const nx = x + dx;
                        const ny = y + dy;
                        if (nx >= 0 && nx < 48 && ny >= 0 && ny < 48 && mouthRegion[ny * 48 + nx] === 1) {
                            isSurrounded = true;
                        }
                    }
                }
                if (isSurrounded) {
                    openedData[y * 48 + x] = 0.5; // Lighten the pixel
                }
            }
        }
    }

    // Combine mouth and lip regions for final processing
    const lipRegion = new Float32Array(48 * 48);
    const lipY = Math.floor(height * 0.65); // Adjusted for better lip detection
    const lipHeight = Math.floor(height * 0.1); // Height of the lip region

    for (let y = lipY; y < lipY + lipHeight; y++) {
        for (let x = 0; x < 48; x++) {
            const index = y * 48 + x;
            if (index < lipRegion.length) {
                lipRegion[index] = processedData[index]; // Copy the processed data for the lip region
            }
        }
    }

    // Ensure lips are included in the processed data
    for (let i = 0; i < 48 * 48; i++) {
        if (lipRegion[i] === 1) {
            processedData[i] = 1; // Ensure lips are included in the processed data
        }
    }
   
    // Visualize the processed image
    // appendFloatArrayAsImage(processedData, "Processed Image"); // Debug

    try {
        const result = await callTeeApi(Array.from(processedData));
        console.log("TEE Result:", result);
        await zoomingPromise;

        if (result.prediction[0] === 1) {
            vibeCheckStatus.value = "success_vibe";
            status.value = "success_vibe";
        } else {
            vibeCheckStatus.value = "failed_vibe";
            status.value = "failed_vibe";
        }
    } catch (error) {
        console.error("Error calling TEE API:", error);
        status.value = "failed_vibe";
    }
}

## Performance Improvement

By using teeML, the response times for the Hyle Vibe Check application were significantly improved. The zkML approach required generating zk-proofs, which involved complex computations and resulted in long wait times for users. In contrast, teeML leverages the secure enclave to make predictions and generate attestations quickly, providing a much smoother user experience.

**zkML**: High security but long response times due to complex proof generation.
**teeML**: High security with significantly reduced response times by leveraging TEEs.

## Conclusion

teeML offers a powerful alternative to zkML by combining the security of trusted execution environments with the efficiency of direct model predictions. Deploying ML models within Marlin's Oyster enclave not only ensures data privacy but also enhances the user experience by reducing response times. This makes teeML an ideal choice for applications requiring both security and performance.

For more details on deploying applications within Oyster, refer to the [Oyster documentation](https://docs.marlin.org/learn/what-is-oyster). For other use-cases of TEEs check out this thread - https://x.com/MarlinProtocol/status/1696217078452150614