Mobile Deployment¶

import UIKit
import LibTorch_Lite

class ImageClassifier {
    private var module: TorchModule
    private let inputSize: Int = 224

    // ImageNet normalization
    private let mean: [Float] = [0.485, 0.456, 0.406]
    private let std: [Float] = [0.229, 0.224, 0.225]

    init(modelPath: String) throws {
        guard let module = TorchModule(fileAtPath: modelPath) else {
            throw ClassifierError.modelLoadFailed
        }
        self.module = module
    }

    func predict(image: UIImage) throws -> Prediction {
        // Preprocess
        guard let tensor = preprocess(image: image) else {
            throw ClassifierError.preprocessingFailed
        }

        // Run inference
        guard let output = module.predict(image: tensor) else {
            throw ClassifierError.inferenceFailed
        }

        // Parse output
        let scores = output.floatArray
        guard let (classId, confidence) = topPrediction(scores: scores) else {
            throw ClassifierError.invalidOutput
        }

        return Prediction(classId: classId, confidence: confidence)
    }

    func predictTop5(image: UIImage) throws -> [Prediction] {
        guard let tensor = preprocess(image: image) else {
            throw ClassifierError.preprocessingFailed
        }

        guard let output = module.predict(image: tensor) else {
            throw ClassifierError.inferenceFailed
        }

        let scores = output.floatArray
        return topKPredictions(scores: scores, k: 5)
    }

    private func preprocess(image: UIImage) -> UnsafeMutableRawPointer? {
        // Resize to 224x224
        guard let resized = image.resized(to: CGSize(
            width: inputSize,
            height: inputSize
        )) else {
            return nil
        }

        // Convert to pixel buffer
        guard let pixelBuffer = resized.pixelBuffer() else {
            return nil
        }

        // Normalize and convert to tensor format
        return normalize(pixelBuffer: pixelBuffer)
    }

    private func normalize(pixelBuffer: CVPixelBuffer) -> UnsafeMutableRawPointer? {
        CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
        defer {
            CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly)
        }

        guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else {
            return nil
        }

        let width = CVPixelBufferGetWidth(pixelBuffer)
        let height = CVPixelBufferGetHeight(pixelBuffer)
        let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)

        // Allocate tensor memory [1, 3, 224, 224]
        let tensorSize = 1 * 3 * height * width
        let tensorData = UnsafeMutablePointer<Float>.allocate(
            capacity: tensorSize
        )

        // Convert RGBA to RGB and normalize
        for y in 0..<height {
            for x in 0..<width {
                let pixelOffset = y * bytesPerRow + x * 4
                let pixel = baseAddress.advanced(by: pixelOffset)
                    .assumingMemoryBound(to: UInt8.self)

                // Extract RGB values
                let r = Float(pixel[0]) / 255.0
                let g = Float(pixel[1]) / 255.0
                let b = Float(pixel[2]) / 255.0

                // Normalize and store in CHW format
                let idx = y * width + x
                tensorData[idx] = (r - mean[0]) / std[0]
                tensorData[height * width + idx] = (g - mean[1]) / std[1]
                tensorData[2 * height * width + idx] = (b - mean[2]) / std[2]
            }
        }

        return UnsafeMutableRawPointer(tensorData)
    }

    private func topPrediction(scores: [Float]) -> (Int, Float)? {
        guard let maxIndex = scores.enumerated().max(
            by: { $0.element < $1.element }
        ) else {
            return nil
        }

        return (maxIndex.offset, maxIndex.element)
    }

    private func topKPredictions(scores: [Float], k: Int) -> [Prediction] {
        let indexed = scores.enumerated().map { ($0.offset, $0.element) }
        let sorted = indexed.sorted { $0.1 > $1.1 }

        return sorted.prefix(k).map { idx, score in
            Prediction(classId: idx, confidence: score)
        }
    }
}

struct Prediction {
    let classId: Int
    let confidence: Float
}

enum ClassifierError: Error {
    case modelLoadFailed
    case preprocessingFailed
    case inferenceFailed
    case invalidOutput
}

// UIImage extensions
extension UIImage {
    func resized(to size: CGSize) -> UIImage? {
        UIGraphicsBeginImageContextWithOptions(size, false, 1.0)
        defer { UIGraphicsEndImageContext() }
        draw(in: CGRect(origin: .zero, size: size))
        return UIGraphicsGetImageFromCurrentImageContext()
    }

    func pixelBuffer() -> CVPixelBuffer? {
        let width = Int(size.width)
        let height = Int(size.height)

        let attrs = [
            kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
            kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue
        ] as CFDictionary

        var pixelBuffer: CVPixelBuffer?
        let status = CVPixelBufferCreate(
            kCFAllocatorDefault,
            width,
            height,
            kCVPixelFormatType_32ARGB,
            attrs,
            &pixelBuffer
        )

        guard status == kCVReturnSuccess, let buffer = pixelBuffer else {
            return nil
        }

        CVPixelBufferLockBaseAddress(buffer, [])
        defer { CVPixelBufferUnlockBaseAddress(buffer, []) }

        let context = CGContext(
            data: CVPixelBufferGetBaseAddress(buffer),
            width: width,
            height: height,
            bitsPerComponent: 8,
            bytesPerRow: CVPixelBufferGetBytesPerRow(buffer),
            space: CGColorSpaceCreateDeviceRGB(),
            bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue
        )

        guard let ctx = context, let cgImage = self.cgImage else {
            return nil
        }

        ctx.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))

        return buffer
    }
}

import UIKit

class ViewController: UIViewController, UIImagePickerControllerDelegate,
                      UINavigationControllerDelegate {

    @IBOutlet weak var imageView: UIImageView!
    @IBOutlet weak var resultLabel: UILabel!
    @IBOutlet weak var confidenceLabel: UILabel!

    private var classifier: ImageClassifier?
    private let classNames = loadClassNames() // Load from file

    override func viewDidLoad() {
        super.viewDidLoad()

        // Load model
        guard let modelPath = Bundle.main.path(
            forResource: "mobile_model",
            ofType: "pt"
        ) else {
            showError("Model file not found")
            return
        }

        do {
            classifier = try ImageClassifier(modelPath: modelPath)
        } catch {
            showError("Failed to load model: \(error)")
        }
    }

    @IBAction func selectImageTapped(_ sender: UIButton) {
        let picker = UIImagePickerController()
        picker.delegate = self
        picker.sourceType = .photoLibrary
        present(picker, animated: true)
    }

    @IBAction func takePictureTapped(_ sender: UIButton) {
        guard UIImagePickerController.isSourceTypeAvailable(.camera) else {
            showError("Camera not available")
            return
        }

        let picker = UIImagePickerController()
        picker.delegate = self
        picker.sourceType = .camera
        present(picker, animated: true)
    }

    func imagePickerController(
        _ picker: UIImagePickerController,
        didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey: Any]
    ) {
        picker.dismiss(animated: true)

        guard let image = info[.originalImage] as? UIImage else {
            return
        }

        imageView.image = image
        classifyImage(image)
    }

    private func classifyImage(_ image: UIImage) {
        guard let classifier = classifier else {
            showError("Classifier not loaded")
            return
        }

        DispatchQueue.global(qos: .userInitiated).async {
            do {
                let prediction = try classifier.predict(image: image)

                DispatchQueue.main.async {
                    self.showResult(prediction: prediction)
                }
            } catch {
                DispatchQueue.main.async {
                    self.showError("Classification failed: \(error)")
                }
            }
        }
    }

    private func showResult(prediction: Prediction) {
        let className = classNames[prediction.classId]
        let confidence = prediction.confidence * 100

        resultLabel.text = "Class: \(className)"
        confidenceLabel.text = String(format: "Confidence: %.1f%%", confidence)
    }

    private func showError(_ message: String) {
        let alert = UIAlertController(
            title: "Error",
            message: message,
            preferredStyle: .alert
        )
        alert.addAction(UIAlertAction(title: "OK", style: .default))
        present(alert, animated: true)
    }

    private static func loadClassNames() -> [String] {
        // Load from bundled file or return defaults
        return (0..<1000).map { "Class \($0)" }
    }
}