import FS from 'fs' ; import { Transcoder } from '../../cache/transcoder' ; import type { Database } from '../../types' ; import TensorFlow from './tensorflow' ; export const MODEL_WIDTH = 224 ; export const MODEL_HEIGHT = 224 ; export const SUBSET_CRITERIA = { tags : { exists : true } , type : { equalsAny : [ 'still' ] } , thumbnail : true , } ; export async function loadFile ( absolutePath ) { return new Promise < Buffer > ( ( resolve , reject ) => { FS . readFile ( absolutePath , ( err , data ) => { if ( err ) { reject ( err ) ; } else { resolve ( data ) ; } } ) ; } ) ; } const IMAGE_NET_RGB = TensorFlow . tensor1d ( [ 123.68 , 116.779 , 103.939 ] ) ; export async function loadImageFileVgg16 ( absolutePath ) { const raw = await loadFile ( absolutePath ) ; const res = TensorFlow . tidy ( ( ) => { const common = TensorFlow . decodeImage ( raw ) . resizeNearestNeighbor ( [ MODEL_WIDTH , MODEL_HEIGHT ] ) . toFloat ( ) ; return common . sub ( IMAGE_NET_RGB ) . reverse ( 2 ) . expandDims ( ) ; } ) ; return res as TensorFlow . Tensor3D | TensorFlow . Tensor4D ; } const MOBILENET_SCALAR = TensorFlow . scalar ( 127.5 ) ; export async function loadImageFileMobileNet ( absolutePath , width = MODEL_WIDTH , height = MODEL_HEIGHT ) { const raw = await loadFile ( absolutePath ) ; const res = TensorFlow . tidy ( ( ) => { const common = TensorFlow . decodeImage ( raw ) . resizeNearestNeighbor ( [ width , height ] ) . toFloat ( ) ; return common . sub ( MOBILENET_SCALAR ) . div ( MOBILENET_SCALAR ) . expandDims ( ) ; } ) ; return res as TensorFlow . Tensor3D | TensorFlow . Tensor4D ; } export async function loadImageFileCommon ( absolutePath , width = MODEL_WIDTH , height = MODEL_HEIGHT ) { const raw = await loadFile ( absolutePath ) ; const res = TensorFlow . tidy ( ( ) => { const common = TensorFlow . decodeImage ( raw ) . resizeNearestNeighbor ( [ width , height ] ) . toFloat ( ) ; return ( common . sub ( MOBILENET_SCALAR ) . div ( 255 ) . expandDims ( ) ) ; } ) ; return res as TensorFlow . Tensor3D | TensorFlow . Tensor4D ; } export type DatasetMapperFunction = ( hash : string ) => Promise < TensorFlow . TensorContainer > ; export function createDatasetMapper ( database , tags ) { const transcoder = new Transcoder ( database ) ; return async ( hash ) => { const media = await database . getMedia ( hash ) ; if ( ! media ) { throw new Error ( `Couldn't find media to generate preview: ${ hash }` ) ; } const labelTensor = TensorFlow . tensor1d ( tags . map ( ( tag ) => ( media . tags . includes ( tag ) ? 1 : 0 ) ) , 'int32' ) ; const reshapedLabelTensor = labelTensor . reshape ( [ 1 , tags . length ] ) ; TensorFlow . dispose ( labelTensor ) ; return { xs : await loadImageFileMobileNet ( transcoder . getThumbnailPath ( media ) ) , ys : reshapedLabelTensor , } ; } ; }