Инференс ML моделей в Java: TensorFlow, DL4J и DJL

@Configuration
public class ModelConfig {
    @Bean
    public SavedModelBundle tensorflowModel() {
        // Загружаем модель при инициализации бина
        return SavedModelBundle.load("/models/classifier", "serve");
    }
    
    @PreDestroy
    public void cleanup() {
        // Явно освобождаем ресурсы при остановке
        tensorflowModel().close();
    }
}

@Component
public class PredictionService {
    private final MeterRegistry registry;
    private final Timer inferenceTimer;
    
    public PredictionService(MeterRegistry registry) {
        this.registry = registry;
        this.inferenceTimer = Timer.builder("ml.inference")
            .tag("model", "classifier")
            .register(registry);
    }
    
    public Result predict(Input data) {
        return inferenceTimer.record(() -> {
            // Реальный вызов модели здесь
            return model.predict(data);
        });
    }
}

// TensorFlow Lite квантизация post-training
Interpreter.Options options = new Interpreter.Options();
options.setNumThreads(4);
// Модель уже квантизирована при конвертации
Interpreter tflite = new Interpreter(loadModelFile(), options);
 
// DL4J требует явного указания precision
MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
    .dataType(DataType.INT8) // Используем int8 precision
    .layer(new DenseLayer.Builder().nIn(784).nOut(256).build())
    // остальные слои...
    .build();

MlflowClient client = new MlflowClient("http://mlflow-server:5000");
ModelVersion model = client.getLatestVersions("classifier", "Production").get(0);
String modelUri = model.getSource();
// Загружаем модель по URI

// Правильный подход с try-with-resources
public Prediction predict(Input data) {
    try (SavedModelBundle model = SavedModelBundle.load(modelPath, "serve")) {
        // Работаем с моделью
        return runInference(model, data);
    } // Автоматически освобождает нативные ресурсы
}
 
// Лучше - загрузить один раз при старте
private final SavedModelBundle model;
 
@PostConstruct
public void initModel() {
    this.model = SavedModelBundle.load(modelPath, "serve");
}
 
@PreDestroy  
public void cleanup() {
    if (model != null) {
        model.close();
    }
}

Criteria<BufferedImage, Classifications> criteria = Criteria.builder()
    .setTypes(BufferedImage.class, Classifications.class)
    .optModelPath(Paths.get("/models/resnet"))
    .optEngine("PyTorch")
    .optProgress(new ProgressBar())
    .build();
 
ZooModel<BufferedImage, Classifications> model = criteria.loadModel();

@PostConstruct
public void warmUpModel() {
// Загружаем модель
this.model = loadModel();
 
// Генерируем dummy input нужной размерности
float[][] dummyData = createDummyInput(inputShape);
 
// Прогреваем 100 итераций
for (int i = 0; i < 100; i++) {
    model.predict(dummyData);
}
 
logger.info("Model warm-up completed");
}

public class ModelPool {
private final BlockingQueue<Predictor> pool;
 
public ModelPool(int size, ZooModel model) {
    this.pool = new ArrayBlockingQueue<>(size);
    for (int i = 0; i < size; i++) {
        pool.offer(model.newPredictor());
    }
}
 
public Result predict(Input data) throws InterruptedException {
    Predictor predictor = pool.take(); // Блокируется если пула пуст
    try {
        return predictor.predict(data);
    } finally {
        pool.offer(predictor); // Возвращаем в пул
    }
}
}

@PreDestroy
public void shutdown() {
if (model != null) {
    logger.info("Closing model resources");
    model.close(); // Освобождаем native memory
}
}

public float[][][] preprocessImage(BufferedImage img) {
    // Resize до нужного размера
    BufferedImage resized = new BufferedImage(224, 224, BufferedImage.TYPE_INT_RGB);
    Graphics2D g = resized.createGraphics();
    g.setRenderingHint(RenderingHints.KEY_INTERPOLATION, 
                       RenderingHints.VALUE_INTERPOLATION_BILINEAR);
    g.drawImage(img, 0, 0, 224, 224, null);
    g.dispose();
    
    // Конвертация в нормализованный массив
    float[][][] data = new float[3][224][224];
    for (int y = 0; y < 224; y++) {
        for (int x = 0; x < 224; x++) {
            int rgb = resized.getRGB(x, y);
            // Извлекаем каналы и нормализуем (-1 до 1 для ImageNet)
            data[0][y][x] = ((rgb >> 16) & 0xFF) / 127.5f - 1.0f; // R
            data[1][y][x] = ((rgb >> 8) & 0xFF) / 127.5f - 1.0f;  // G
            data[2][y][x] = (rgb & 0xFF) / 127.5f - 1.0f;         // B
        }
    }
    return data;
}

public class ResilientModelLoader {
    private static final int MAX_RETRIES = 3;
    private static final long RETRY_DELAY_MS = 1000;
    
    public SavedModelBundle loadWithRetry(String modelPath) {
        Exception lastException = null;
        
        for (int attempt = 1; attempt <= MAX_RETRIES; attempt++) {
            try {
                // Проверяем существование файла
                if (!Files.exists(Paths.get(modelPath))) {
                    throw new ModelNotFoundException("Model not found: " + modelPath);
                }
                
                // Загружаем модель
                SavedModelBundle model = SavedModelBundle.load(modelPath, "serve");
                logger.info("Model loaded successfully on attempt {}", attempt);
                return model;
                
            } catch (IllegalArgumentException e) {
                // Невалидный путь или тег - retry бесполезен
                logger.error("Invalid model configuration", e);
                throw new ModelLoadException("Model configuration error", e);
                
            } catch (Exception e) {
                lastException = e;
                logger.warn("Failed to load model, attempt {}/{}", attempt, MAX_RETRIES, e);
                
                if (attempt < MAX_RETRIES) {
                    try {
                        Thread.sleep(RETRY_DELAY_MS * attempt); // Exponential backoff
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        break;
                    }
                }
            }
        }
        
        // Все попытки исчерпаны
        throw new ModelLoadException("Failed to load model after retries", lastException);
    }
}

public class EfficientPredictor {
    // Переиспользуемый буфер вместо аллокации на каждый запрос
    private final ThreadLocal<float[][]> inputBuffer = 
        ThreadLocal.withInitial(() -> new float[1][784]);
    
    public Result predict(InputData data) {
        float[][] buffer = inputBuffer.get();
        // Копируем данные в существующий буфер
        System.arraycopy(data.getFeatures(), 0, buffer[0], 0, 784);
        
        // Выполняем inference без лишних аллокаций
        return model.predict(buffer);
    }
}

public class BatchingPredictor {
    private final BlockingQueue<PredictionRequest> queue = new LinkedBlockingQueue<>(1000);
    private final int batchSize = 32;
    private final long maxWaitMs = 10;
    
    public CompletableFuture<Result> predictAsync(InputData data) {
        CompletableFuture<Result> future = new CompletableFuture<>();
        queue.offer(new PredictionRequest(data, future));
        return future;
    }
    
    // Отдельный поток собирает батчи и выполняет inference
    private void processBatches() {
        List<PredictionRequest> batch = new ArrayList<>(batchSize);
        while (running) {
            batch.clear();
            queue.drainTo(batch, batchSize);
            
            if (!batch.isEmpty()) {
                // Собрали батч, выполняем предсказание
                List<Result> results = model.predictBatch(
                    batch.stream()
                        .map(PredictionRequest::getData)
                        .collect(Collectors.toList())
                );
                
                // Возвращаем результаты через futures
                for (int i = 0; i < batch.size(); i++) {
                    batch.get(i).getFuture().complete(results.get(i));
                }
            }
        }
    }
}

// Pre-allocated batch buffer для эффективности
private final float[][][] batchBuffer = new float[MAX_BATCH_SIZE][3][224][224];
private final AtomicInteger currentBatchSize = new AtomicInteger(0);
 
public void addToBatch(float[][][] input) {
    int index = currentBatchSize.getAndIncrement();
    // Копируем в существующий буфер вместо аллокации нового
    System.arraycopy(input[0], 0, batchBuffer[index], 0, input[0].length);
}

@RestController
public class AsyncPredictionController {
    private final ExecutorService inferenceExecutor = 
        Executors.newFixedThreadPool(4); // Ограничиваем параллелизм
    
    @PostMapping("/predict")
    public CompletableFuture<PredictionResponse> predict(@RequestBody InputData data) {
        return CompletableFuture.supplyAsync(() -> {
            // Тяжёлая работа в отдельном пуле
            float[] features = preprocessor.transform(data);
            Result result = model.predict(features);
            return new PredictionResponse(result);
        }, inferenceExecutor);
    }
}

@Component
public class MonitoredPredictor {
    private final Timer inferenceTimer;
    private final Counter errorCounter;
    
    public MonitoredPredictor(MeterRegistry registry) {
        this.inferenceTimer = Timer.builder("inference.duration")
            .description("Время выполнения предсказания модели")
            .tag("model", "classifier")
            .publishPercentiles(0.5, 0.95, 0.99)
            .register(registry);
            
        this.errorCounter = Counter.builder("inference.errors")
            .description("Количество ошибок при inference")
            .tag("model", "classifier")
            .register(registry);
    }
    
    public Result predict(Input data) {
        return inferenceTimer.record(() -> {
            try {
                return model.predict(data);
            } catch (Exception e) {
                errorCounter.increment();
                throw e;
            }
        });
    }
}

@PostConstruct
public void warmUp() {
    float[][] dummyInput = new float[1][784];
    Arrays.fill(dummyInput[0], 0.5f);
    
    // Первые 50 вызовов самые важные для JIT
    for (int i = 0; i < 100; i++) {
        model.predict(dummyInput);
    }
    
    logger.info("Model warmed up and ready");
}

// Модель загружается из classpath, а не с диска
@Bean
public SavedModelBundle model() {
    // Распаковка из ресурсов jar если нужно
    Path modelPath = extractModelFromResources();
    return SavedModelBundle.load(modelPath.toString(), "serve");
}

@GetMapping("/health")
public ResponseEntity<String> health() {
    try {
        // Быстрая проверка что модель отвечает
        float[][] testInput = new float[1][784];
        Result result = model.predict(testInput);
        return ResponseEntity.ok("Healthy");
    } catch (Exception e) {
        return ResponseEntity.status(503).body("Model unavailable");
    }
}

private final LoadingCache<String, Result> cache = Caffeine.newBuilder()
    .maximumSize(10_000) // Лимит по количеству записей
    .expireAfterWrite(Duration.ofMinutes(15)) // TTL важен для свежести
    .recordStats() // Метрики hit rate
    .build(key -> model.predict(parseInput(key)));
 
public Result predictWithCache(InputData data) {
    String cacheKey = generateKey(data); // Hash входных данных
    return cache.get(cacheKey);
}

public Result predictWithFallback(InputData data) {
try {
  // Пробуем основную модель
  return primaryModel.predict(data);
} catch (Exception e) {
  logger.warn("Primary model failed, trying fallback", e);
  
  try {
      // Упрощённая модель как запасной вариант
      return fallbackModel.predict(data);
  } catch (Exception e2) {
      logger.warn("Fallback model failed, using cache", e2);
      
      // Ищем похожий результат в кеше
      return cache.getSimilar(data)
          .orElseGet(() -> getDefaultPrediction(data));
  }
}
}

public interface ImageClassifier {
    ClassificationResult classify(BufferedImage image);
    String getEngineName();
    void warmUp();
}
 
@Component
public class TensorFlowClassifier implements ImageClassifier {
    private SavedModelBundle model;
    // Реализация через TensorFlow Java
}
 
@Component  
public class DL4JClassifier implements ImageClassifier {
    private MultiLayerNetwork model;
    // Реализация через Deeplearning4j
}
 
@Component
public class DJLClassifier implements ImageClassifier {
    private Predictor<BufferedImage, Classifications> predictor;
    // Реализация через Deep Java Library
}

@RestController
@RequestMapping("/api/v1")
public class ClassificationController {
    
    private final Map<String, ImageClassifier> classifiers;
    private final MeterRegistry meterRegistry;
    
    public ClassificationController(List<ImageClassifier> classifierList, 
                                   MeterRegistry meterRegistry) {
        // Spring инжектит все имплементации автоматически
        this.classifiers = classifierList.stream()
            .collect(Collectors.toMap(
                ImageClassifier::getEngineName,
                classifier -> classifier
            ));
        this.meterRegistry = meterRegistry;
    }
    
    @PostMapping(value = "/classify", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
    public ResponseEntity<ClassificationResponse> classify(
            @RequestParam("image") MultipartFile file,
            @RequestParam(value = "engine", defaultValue = "tensorflow") String engine) {
        
        Timer.Sample sample = Timer.start(meterRegistry);
        
        try {
            BufferedImage image = ImageIO.read(file.getInputStream());
            if (image == null) {
                return ResponseEntity.badRequest()
                    .body(new ClassificationResponse("Invalid image format"));
            }
            
            ImageClassifier classifier = classifiers.get(engine.toLowerCase());
            if (classifier == null) {
                return ResponseEntity.badRequest()
                    .body(new ClassificationResponse("Unknown engine: " + engine));
            }
            
            ClassificationResult result = classifier.classify(image);
            
            sample.stop(Timer.builder("classification.duration")
                .tag("engine", engine)
                .register(meterRegistry));
            
            return ResponseEntity.ok(ClassificationResponse.from(result, engine));
            
        } catch (IOException e) {
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
                .body(new ClassificationResponse("Failed to process image: " + e.getMessage()));
        }
    }
}

@ControllerAdvice
public class GlobalExceptionHandler {
    
    private static final Logger logger = LoggerFactory.getLogger(GlobalExceptionHandler.class);
    
    @ExceptionHandler(ModelLoadException.class)
    public ResponseEntity<ErrorResponse> handleModelLoadException(ModelLoadException e) {
        logger.error("Failed to load model", e);
        return ResponseEntity.status(HttpStatus.SERVICE_UNAVAILABLE)
            .body(new ErrorResponse("ML model temporarily unavailable", e.getMessage()));
    }
    
    @ExceptionHandler(InferenceException.class)
    public ResponseEntity<ErrorResponse> handleInferenceException(InferenceException e) {
        logger.error("Inference failed", e);
        return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
            .body(new ErrorResponse("Prediction failed", e.getMessage()));
    }
}

@Component
public class TensorFlowClassifier implements ImageClassifier {
    
    private static final Logger logger = LoggerFactory.getLogger(TensorFlowClassifier.class);
    private static final int IMAGE_SIZE = 224;
    private static final int NUM_CLASSES = 1000;
    
    private SavedModelBundle model;
    private final MeterRegistry meterRegistry;
    private final String modelPath;
    
    public TensorFlowClassifier(MeterRegistry meterRegistry,
                               @Value("${models.tensorflow.path}") String modelPath) {
        this.meterRegistry = meterRegistry;
        this.modelPath = modelPath;
        loadModel();
        warmUp();
    }
    
    private void loadModel() {
        try {
            this.model = SavedModelBundle.load(modelPath, "serve");
            logger.info("TensorFlow model loaded from {}", modelPath);
        } catch (Exception e) {
            throw new ModelLoadException("Failed to load TensorFlow model", e);
        }
    }
    
    @Override
    public ClassificationResult classify(BufferedImage image) {
        Timer.Sample sample = Timer.start(meterRegistry);
        
        try {
            // Конвертация изображения в тензор нужной размерности
            float[][][][] inputData = preprocessImage(image);
            
            try (TFloat32 inputTensor = TFloat32.tensorOf(
                    StdArrays.ndCopyOf(inputData))) {
                
                // Выполнение inference через сессию TensorFlow
                Map<String, Tensor> feedDict = new HashMap<>();
                feedDict.put("serving_default_input:0", inputTensor);
                
                List<Tensor<?>> outputs = model.session()
                    .runner()
                    .feed("serving_default_input:0", inputTensor)
                    .fetch("StatefulPartitionedCall:0")
                    .run();
                
                // Извлечение результатов из выходного тензора
                try (TFloat32 outputTensor = (TFloat32) outputs.get(0)) {
                    float[][] predictions = new float[1][NUM_CLASSES];
                    outputTensor.read(StdArrays.array2dCopyOf(predictions));
                    
                    return processOutput(predictions[0]);
                }
            }
        } catch (Exception e) {
            meterRegistry.counter("classification.errors", "engine", "tensorflow")
                .increment();
            throw new InferenceException("TensorFlow inference failed", e);
        } finally {
            sample.stop(meterRegistry.timer("classification.duration", 
                "engine", "tensorflow"));
        }
    }
    
    private float[][][][] preprocessImage(BufferedImage image) {
        // Resize с сохранением пропорций
        BufferedImage resized = new BufferedImage(IMAGE_SIZE, IMAGE_SIZE, 
            BufferedImage.TYPE_INT_RGB);
        Graphics2D g = resized.createGraphics();
        g.setRenderingHint(RenderingHints.KEY_INTERPOLATION,
            RenderingHints.VALUE_INTERPOLATION_BILINEAR);
        g.drawImage(image, 0, 0, IMAGE_SIZE, IMAGE_SIZE, null);
        g.dispose();
        
        // Конвертация в формат NHWC (batch, height, width, channels)
        float[][][][] data = new float[1][IMAGE_SIZE][IMAGE_SIZE][3];
        for (int y = 0; y < IMAGE_SIZE; y++) {
            for (int x = 0; x < IMAGE_SIZE; x++) {
                int rgb = resized.getRGB(x, y);
                // ImageNet нормализация
                data[0][y][x][0] = ((rgb >> 16) & 0xFF) / 255.0f; // R
                data[0][y][x][1] = ((rgb >> 8) & 0xFF) / 255.0f;  // G  
                data[0][y][x][2] = (rgb & 0xFF) / 255.0f;         // B
            }
        }
        return data;
    }
    
    private ClassificationResult processOutput(float[] predictions) {
        // Находим топ-5 классов с наивысшими вероятностями
        List<ClassLabel> topPredictions = new ArrayList<>();
        for (int i = 0; i < predictions.length; i++) {
            topPredictions.add(new ClassLabel(i, predictions[i]));
        }
        
        topPredictions.sort((a, b) -> Float.compare(b.probability, a.probability));
        return new ClassificationResult(topPredictions.subList(0, 
            Math.min(5, topPredictions.size())));
    }
    
    @Override
    public void warmUp() {
        logger.info("Warming up TensorFlow model...");
        BufferedImage dummy = new BufferedImage(IMAGE_SIZE, IMAGE_SIZE, 
            BufferedImage.TYPE_INT_RGB);
        for (int i = 0; i < 50; i++) {
            classify(dummy);
        }
        logger.info("TensorFlow model warmed up");
    }
    
    @Override
    public String getEngineName() {
        return "tensorflow";
    }
    
    @PreDestroy
    public void cleanup() {
        if (model != null) {
            model.close();
            logger.info("TensorFlow model resources released");
        }
    }
}

@Component
public class DL4JClassifier implements ImageClassifier {
    
    private static final Logger logger = LoggerFactory.getLogger(DL4JClassifier.class);
    private static final int IMAGE_SIZE = 224;
    
    private MultiLayerNetwork model;
    private final MeterRegistry meterRegistry;
    
    public DL4JClassifier(MeterRegistry meterRegistry,
                         @Value("${models.dl4j.path}") String modelPath) {
        this.meterRegistry = meterRegistry;
        loadModel(modelPath);
        warmUp();
    }
    
    private void loadModel(String path) {
        try {
            File modelFile = new File(path);
            this.model = ModelSerializer.restoreMultiLayerNetwork(modelFile);
            logger.info("DL4J model loaded from {}", path);
        } catch (IOException e) {
            throw new ModelLoadException("Failed to load DL4J model", e);
        }
    }
    
    @Override
    public ClassificationResult classify(BufferedImage image) {
        Timer.Sample sample = Timer.start(meterRegistry);
        
        try (INDArray input = preprocessImage(image)) {
            // DL4J inference - thread-safe после init
            INDArray output = model.output(input);
            
            ClassificationResult result = extractPredictions(output);
            
            // Явное освобождение output array
            output.close();
            
            return result;
            
        } catch (Exception e) {
            meterRegistry.counter("classification.errors", "engine", "dl4j")
                .increment();
            throw new InferenceException("DL4J inference failed", e);
        } finally {
            sample.stop(meterRegistry.timer("classification.duration",
                "engine", "dl4j"));
        }
    }
    
    private INDArray preprocessImage(BufferedImage image) {
        // Resize через Java Graphics2D
        BufferedImage resized = ImageUtils.resize(image, IMAGE_SIZE, IMAGE_SIZE);
        
        // Создаём INDArray нужной размерности (batch, channels, height, width)
        INDArray arr = Nd4j.create(1, 3, IMAGE_SIZE, IMAGE_SIZE);
        
        for (int y = 0; y < IMAGE_SIZE; y++) {
            for (int x = 0; x < IMAGE_SIZE; x++) {
                int rgb = resized.getRGB(x, y);
                // Нормализация в диапазон [-1, 1]
                arr.putScalar(new int[]{0, 0, y, x}, 
                    (((rgb >> 16) & 0xFF) / 127.5f) - 1.0f);
                arr.putScalar(new int[]{0, 1, y, x},
                    (((rgb >> 8) & 0xFF) / 127.5f) - 1.0f);
                arr.putScalar(new int[]{0, 2, y, x},
                    ((rgb & 0xFF) / 127.5f) - 1.0f);
            }
        }
        
        return arr;
    }
    
    private ClassificationResult extractPredictions(INDArray output) {
        float[] predictions = output.toFloatVector();
        
        List<ClassLabel> labels = IntStream.range(0, predictions.length)
            .mapToObj(i -> new ClassLabel(i, predictions[i]))
            .sorted((a, b) -> Float.compare(b.probability, a.probability))
            .limit(5)
            .collect(Collectors.toList());
            
        return new ClassificationResult(labels);
    }
    
    @Override
    public void warmUp() {
        logger.info("Warming up DL4J model...");
        BufferedImage dummy = new BufferedImage(IMAGE_SIZE, IMAGE_SIZE,
            BufferedImage.TYPE_INT_RGB);
        for (int i = 0; i < 50; i++) {
            classify(dummy);
        }
        logger.info("DL4J model warmed up");
    }
    
    @Override
    public String getEngineName() {
        return "dl4j";
    }
}

@Component
public class DJLClassifier implements ImageClassifier {
    
    private static final Logger logger = LoggerFactory.getLogger(DJLClassifier.class);
    
    private ZooModel<Image, Classifications> model;
    private final ThreadLocal<Predictor<Image, Classifications>> predictorPool;
    private final MeterRegistry meterRegistry;
    
    public DJLClassifier(MeterRegistry meterRegistry,
                        @Value("${models.djl.path}") String modelPath,
                        @Value("${models.djl.engine:PyTorch}") String engine) {
        this.meterRegistry = meterRegistry;
        loadModel(modelPath, engine);
        
        // Thread-local predictor для безопасности потоков
        this.predictorPool = ThreadLocal.withInitial(() -> model.newPredictor());
        
        warmUp();
    }
    
    private void loadModel(String path, String engineName) {
        try {
            Criteria<Image, Classifications> criteria = Criteria.builder()
                .setTypes(Image.class, Classifications.class)
                .optModelPath(Paths.get(path))
                .optEngine(engineName)
                .optProgress(new ProgressBar())
                .build();
                
            this.model = criteria.loadModel();
            logger.info("DJL model loaded with {} engine from {}", engineName, path);
            
        } catch (Exception e) {
            throw new ModelLoadException("Failed to load DJL model", e);
        }
    }
    
    @Override
    public ClassificationResult classify(BufferedImage image) {
        Timer.Sample sample = Timer.start(meterRegistry);
        
        try {
            // Конвертация BufferedImage в DJL Image
            Image djlImage = ImageFactory.getInstance().fromImage(image);
            
            // Predictor потокобезопасен через ThreadLocal
            Predictor<Image, Classifications> predictor = predictorPool.get();
            Classifications classifications = predictor.predict(djlImage);
            
            return convertToResult(classifications);
            
        } catch (Exception e) {
            meterRegistry.counter("classification.errors", "engine", "djl")
                .increment();
            throw new InferenceException("DJL inference failed", e);
        } finally {
            sample.stop(meterRegistry.timer("classification.duration",
                "engine", "djl"));
        }
    }
    
    private ClassificationResult convertToResult(Classifications classifications) {
        List<ClassLabel> labels = classifications.topK(5).stream()
            .map(c -> new ClassLabel(
                c.getClassId(),
                (float) c.getProbability(),
                c.getClassName()
            ))
            .collect(Collectors.toList());
            
        return new ClassificationResult(labels);
    }
    
    @Override
    public void warmUp() {
        logger.info("Warming up DJL model...");
        BufferedImage dummy = new BufferedImage(224, 224, 
            BufferedImage.TYPE_INT_RGB);
        for (int i = 0; i < 50; i++) {
            classify(dummy);
        }
        logger.info("DJL model warmed up");
    }
    
    @Override
    public String getEngineName() {
        return "djl";
    }
}

public record ClassLabel(int classId, float probability, String className) {
    public ClassLabel(int classId, float probability) {
        this(classId, probability, ImageNetLabels.getLabel(classId));
    }
}
 
public record ClassificationResult(List<ClassLabel> predictions,
                                  long inferenceTimeMs) {
    public ClassificationResult(List<ClassLabel> predictions) {
        this(predictions, 0L);
    }
}
 
public record ClassificationResponse(String engine,
                                    List<ClassLabel> predictions,
                                    long inferenceTimeMs,
                                    String error) {
    public static ClassificationResponse from(ClassificationResult result, 
                                             String engine) {
        return new ClassificationResponse(engine, result.predictions(),
            result.inferenceTimeMs(), null);
    }
    
    public ClassificationResponse(String error) {
        this(null, null, 0L, error);
    }
}

public class ImageNetLabels {
    private static final Map<Integer, String> LABELS = loadLabels();
    
    private static Map<Integer, String> loadLabels() {
        try (InputStream is = ImageNetLabels.class
                .getResourceAsStream("/imagenet_classes.txt");
             BufferedReader reader = new BufferedReader(
                 new InputStreamReader(is))) {
            
            Map<Integer, String> labels = new HashMap<>();
            String line;
            int index = 0;
            while ((line = reader.readLine()) != null) {
                labels.put(index++, line.trim());
            }
            return labels;
        } catch (IOException e) {
            throw new RuntimeException("Failed to load ImageNet labels", e);
        }
    }
    
    public static String getLabel(int classId) {
        return LABELS.getOrDefault(classId, "unknown");
    }
}

# Стадия сборки
FROM maven:3.9-eclipse-temurin-17 AS builder
WORKDIR /build
COPY pom.xml .
COPY src src
RUN mvn clean package -DskipTests
 
# Рантайм стадия  
FROM eclipse-temurin:17-jre-alpine
WORKDIR /app
 
# Копируем только jar и модели
COPY --from=builder /build/target/*.jar app.jar
COPY models /app/models
 
# Настройка JVM для контейнера
ENV JAVA_OPTS="-XX:+UseZGC -XX:MaxRAMPercentage=75.0 -XX:+HeapDumpOnOutOfMemoryError"
 
EXPOSE 8080
ENTRYPOINT ["sh", "-c", "java $JAVA_OPTS -jar app.jar"]