creek

  1import 'dart:io';
  2import 'dart:convert';
  3import 'package:creekui/utils/image_utils.dart';
  4import 'package:flutter/foundation.dart';
  5import 'package:image/image.dart' as img;
  6import 'package:onnxruntime/onnxruntime.dart';
  7
  8class TextureAnalyzerService {
  9  OrtSession? _session;
 10  Map<String, List<double>>? _centroids;
 11
 12  // DINOv2 Small (S14) Config
 13  static const int INPUT_SIZE = 224;
 14  static const int EMBED_DIM = 384;
 15
 16  // Logic params matching your Python script
 17  static const double RELATIVE_THRESH = 0.85;
 18  static const double ABSOLUTE_FLOOR = 0.05; // Updated to match Python
 19
 20  Future<void> initialize({
 21    required String modelPath,
 22    required String jsonPath,
 23  }) async {
 24    if (_session != null) return;
 25
 26    try {
 27      debugPrint("Initializing Texture Service (DINOv2 S14)...");
 28
 29      OrtEnv.instance.init();
 30
 31      // 1. Load Model
 32      final sessionOptions = OrtSessionOptions();
 33      _session = OrtSession.fromFile(File(modelPath), sessionOptions);
 34      sessionOptions.release();
 35
 36      // 2. Load Centroids
 37      final jsonStr = await File(jsonPath).readAsString();
 38      final Map<String, dynamic> jsonMap = json.decode(jsonStr);
 39      _centroids = {};
 40      jsonMap.forEach((k, v) => _centroids![k] = List<double>.from(v));
 41    } catch (e) {
 42      debugPrint("Texture Service Init Error: $e");
 43    }
 44  }
 45
 46  Future<Map<String, dynamic>> analyze(
 47    String path, {
 48    String? modelPath,
 49    String? jsonPath,
 50  }) async {
 51    if (modelPath == null || jsonPath == null) {
 52      return {'success': false, 'scores': {}, 'error': 'Paths missing'};
 53    }
 54    await initialize(modelPath: modelPath, jsonPath: jsonPath);
 55    if (_session == null || _centroids == null) {
 56      return {'success': false, 'scores': {}, 'error': 'Init failed'};
 57    }
 58
 59    OrtValueTensor? inputOrt;
 60    OrtRunOptions? runOptions;
 61    List<OrtValue?>? outputs;
 62
 63    try {
 64      // 1. Preprocess (Standard ImageNet Normalization)
 65      final bytes = await File(path).readAsBytes();
 66      final image = img.decodeImage(bytes);
 67      if (image == null) {
 68        return {'success': false, 'scores': {}, 'error': 'Decode failed'};
 69      }
 70
 71      final resized = img.copyResize(
 72        image,
 73        width: INPUT_SIZE,
 74        height: INPUT_SIZE,
 75      );
 76      final List<double> inputFloats = List.filled(
 77        1 * 3 * INPUT_SIZE * INPUT_SIZE,
 78        0.0,
 79      );
 80
 81      // ImageNet Stats
 82      const mean = [0.485, 0.456, 0.406];
 83      const std = [0.229, 0.224, 0.225];
 84
 85      int pixelIndex = 0;
 86      for (int y = 0; y < INPUT_SIZE; y++) {
 87        for (int x = 0; x < INPUT_SIZE; x++) {
 88          final pixel = resized.getPixel(x, y);
 89          // R
 90          inputFloats[pixelIndex] = ((pixel.r / 255.0) - mean[0]) / std[0];
 91          // G
 92          inputFloats[pixelIndex + (INPUT_SIZE * INPUT_SIZE)] =
 93              ((pixel.g / 255.0) - mean[1]) / std[1];
 94          // B
 95          inputFloats[pixelIndex + (2 * INPUT_SIZE * INPUT_SIZE)] =
 96              ((pixel.b / 255.0) - mean[2]) / std[2];
 97          pixelIndex++;
 98        }
 99      }
100
101      final float32List = Float32List.fromList(inputFloats);
102
103      // 2. Inference
104      inputOrt = OrtValueTensor.createTensorWithDataList(float32List, [
105        1,
106        3,
107        INPUT_SIZE,
108        INPUT_SIZE,
109      ]);
110      runOptions = OrtRunOptions();
111
112      // Run Inference
113      // Note: 'input' is standard for tf2onnx/torch.onnx exports
114      outputs = _session!.run(runOptions, {"input": inputOrt});
115
116      if (outputs.isEmpty) throw Exception("No output");
117
118      // 3. Extract Feature Vector
119      // The Python export script baked "Mean Pooling" into the model.
120      // So we get a single vector [1, 384] directly. No need to loop patches!
121      final rawOutput = outputs[0]?.value as List;
122      final List<double> embedding = [];
123      void flatten(dynamic data) {
124        if (data is num) {
125          embedding.add(data.toDouble());
126        } else if (data is List) {
127          for (var item in data) {
128            flatten(item);
129          }
130        }
131      }
132
133      flatten(rawOutput);
134
135      // 4. Normalize Embedding (L2 Norm)
136      // DINOv2 cosine similarity requires normalized vectors.
137      final normalizedEmbedding = l2Normalize(embedding);
138
139      // 5. Scoring (Dot Product = Cosine Similarity)
140      Map<String, double> rawScores = {};
141      _centroids!.forEach((name, centroid) {
142        double dot = dotProduct(normalizedEmbedding, centroid);
143        rawScores[name] = dot;
144      });
145
146      // 6. Filter & Sort
147      var sorted =
148          rawScores.entries.toList()
149            ..sort((a, b) => b.value.compareTo(a.value));
150      Map<String, dynamic> finalResults = {};
151
152      if (sorted.isNotEmpty) {
153        double topScore = sorted.first.value;
154        for (var entry in sorted) {
155          if (entry.value < ABSOLUTE_FLOOR) continue;
156          if (entry.value < (topScore * RELATIVE_THRESH)) break;
157
158          finalResults[entry.key] = entry.value;
159        }
160      }
161
162      return {'success': true, 'scores': finalResults, 'error': null};
163    } catch (e) {
164      debugPrint("Texture Analysis Failed: $e");
165      return {'success': false, 'scores': {}, 'error': e.toString()};
166    } finally {
167      inputOrt?.release();
168      runOptions?.release();
169      outputs?.forEach((element) => element?.release());
170    }
171  }
172
173  void dispose() {
174    _session?.release();
175  }
176}