001package org.nasdanika.ai;
002
003import java.awt.image.BufferedImage;
004import java.io.ByteArrayOutputStream;
005import java.io.IOException;
006import java.security.MessageDigest;
007import java.security.NoSuchAlgorithmException;
008import java.util.Map;
009
010import javax.imageio.ImageIO;
011
012/**
013 * Caches image embeddings in a map which can be loaded and saved between runs
014 * Uses image digest as caching key
015 */
016public class CachingImageEmbeddingGenerator<E> extends MapCachingEmbeddingGenerator<BufferedImage, E, String> implements ImageEmbeddingGenerator<E> {
017        
018        private String algorithm;
019
020        public CachingImageEmbeddingGenerator(ImageEmbeddingGenerator<E> target, Map<String, E> cache, String algorithm) {
021                super(target, cache);
022                this.algorithm = algorithm;
023        }
024
025        /**
026         * Uses SHA-512 algorithm
027         * @param cache
028         */
029        public CachingImageEmbeddingGenerator(ImageEmbeddingGenerator<E> target, Map<String, E> cache) {
030                this(target, cache, "SHA-512");
031        }
032        
033        protected String computeKey(BufferedImage input) {
034                try {
035                        ByteArrayOutputStream baos = new ByteArrayOutputStream();
036                        try (baos) {
037                                ImageIO.write(input, "png", baos);
038                        }
039                        byte[] imageBytes = baos.toByteArray();
040                        MessageDigest digest = MessageDigest.getInstance(algorithm);
041                        byte[] dBytes = digest.digest(imageBytes);
042                        
043                    StringBuilder sb = new StringBuilder();
044                    for (byte b : dBytes) {
045                        sb.append(String.format("%02x", b));
046                    }
047                    return sb.toString();
048                } catch (IOException | NoSuchAlgorithmException e) {
049                        throw new IllegalArgumentException("Error computing image digest: " + e, e);
050                }
051        }
052        
053}