React state is always one step behind while making predictions on uploaded image using Tensorflowjs

Question

I'm developing a React app to classify Pokemon based on Images using Tensorflow.js.

What I want - Upload an image of a Pokemon, generate predictions for the same Pokemon.
What is actually happening - When I upload an image to make a prediction, the output is always for the previous image. So, 1st prediction is always garbage (random Pokemon). The predictions I get when I upload the 2nd Pokemon is always for the 1st Pokemon. 3rd Pokemon when uploaded gives prediction for the 2nd Pokemon and so on.

See the bottom of the question for a detailed version of where it works and where it doesn't.

Here are the relevant pieces of code -

I first check if the model is present in indexeddb, if yes, I load it into state model. If not, I fetch it from a server and store it in the state. This is what the first useEffect does upon the first render of the page.
I use another useEffect that runs whenever the findState.uploadedImage changes. This state is present in Redux-toolkit.

Here's a short demo of the problem => https://youtu.be/MX70zbupNWQ

Here's the app URL => https://poke-zoo.herokuapp.com/

Here's the Github repo => https://github.com/theairbend3r/poke-zoo/tree/master/frontend/src/features/find

Here's the file SearchOutput.js. This fetches the model and makes predictions.

const SearchOutput = () => {
  const findState = useSelector(selectorFind)
  const dispatch = useDispatch()
  const imageRef = useRef(null)

  const [model, setModel] = useState(null)
  const [predictions, setPredictions] = useState([])

  const MODEL_HTTP_URL = "api/pokeml/classify"
  const MODEL_INDEXEDDB_URL = "indexeddb://poke-model"

  useEffect(() => {
    async function fetchModel() {
      try {
        const localClassifierModel = await tf.loadLayersModel(
          MODEL_INDEXEDDB_URL
        )

        setModel(localClassifierModel)
        console.log("Model loaded from IndexedDB")
      } catch (e) {
        const classifierModel = await tf.loadLayersModel(MODEL_HTTP_URL)
        setModel(classifierModel)

        await classifierModel.save(MODEL_INDEXEDDB_URL)

        console.error(e)
      }
    }
    fetchModel()
  }, [])

  const getTopKPred = (pred, k) => {
    const predIdx = []
    const predNames = []

    const topkPred = [...pred].sort((a, b) => b - a).slice(0, k)

    topkPred.map(i => predIdx.push(pred.indexOf(i)))
    predIdx.map(i => predNames.push(idx2class[i]))

    return predNames
  }

  useEffect(() => {
    async function makePredictions() {
      if (imageRef && model) {
        try {
          const imgTensor = tf.browser
            .fromPixels(imageRef.current)
            .resizeNearestNeighbor([160, 160])
            .toFloat()
            .sub(127.5)
            .div(127.5)
            .expandDims()

          const y_pred = await model.predict(imgTensor).dataSync()
          const topkPredNames = getTopKPred(y_pred, 5)

          console.log(topkPredNames)
          return topkPredNames
        } catch (e) {
          console.log("Unable to run predictions.")
        }
      }
    }
    makePredictions()
  }, [findState.uploadedImage])

  return (
     
        {findState.uploadedImage && (
          
        )}


        
           {findState.matchesFound.length === 6 &&
             findState.matchesFound.map(poke => (
               
             ))}
       
     
)
}

Here's the file findSlice.js that stores the input image into the redux state.

import { createSlice } from "@reduxjs/toolkit"
import axios from "axios"

const initialState = {
  uploadedImage: "",
  model: null,
  matchesFound: [],
}

export const findSlice = createSlice({
  name: "find",
  initialState: initialState,
  reducers: {
    storeInputImage: (state, action) => {
      state.uploadedImage = action.payload.uploadedImage
    },
    setModel: (state, action) => {
      state.model = action.payload.model
    },
  },
})

export const selectorFind = state => state.find
export const { storeInputImage, setModel } = findSlice.actions
export default findSlice.reducer

Details about the problem.

### Desktop

#### Table

|         |   Ubuntu    |     Windows      |      MacOS       |
| ------- | :---------: | :--------------: | :--------------: |
| Firefox | not working |   not working    |   not working    |
| Chrome  | not working | somewhat working | somewhat working |
| Safari  |     NA      |        NA        | somewhat working |

#### Comments

|         |                                    Ubuntu                                    |                                   Windows                                    |                                    MacOS                                     |
| :-----: | :--------------------------------------------------------------------------: | :--------------------------------------------------------------------------: | :--------------------------------------------------------------------------: |
| Firefox | Predictions are always one step behin for both Captured and Uploaded images. | Predictions are always one step behin for both Captured and Uploaded images. | Predictions are always one step behin for both Captured and Uploaded images. |
| Chrome  | Works only on Captured Images. Uploaded images give same predictions always. | Works only on Captured Images. Uploaded images give same predictions always. | Works only on Captured Images. Uploaded images give same predictions always. |
| Safari  |                                      NA                                      |                                      NA                                      | Works only on Captured Images. Uploaded images give same predictions always. |

### Mobile

#### Table

|         |     Android      |       iOS        |
| ------- | :--------------: | :--------------: |
| Firefox | somewhat working |   not working    |
| Chrome  | somewhat working |   not working    |
| Safari  |        NA        | somewhat working |

#### Comments

|         | Android                                                                     | iOS                                                                          |
| :-----: | --------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
| Firefox | No predictions load for a captured image. Works on uploaded images only.    | Camera does not load. Uploaded images give same predictions always.          |
| Chrome  | Works on capture images only. Uploaded images give same predictions always. | Camera does not load. Uploaded images give same predictions always.          |
| Safari  | NA                                                                          | Works only on Captured Images. Uploaded images give same predictions always. |

Edit: Based on the suggestions below. This has not solved the problem tho. Putting it here for reference.

const SearchOutput = () => {
  const findState = useSelector(selectorFind)
  const [imageRef, setImageRef] = useState(null)

  const onChangeRef = useCallback(node => {
    setImageRef(node)
  }, [])

  const [model, setModel] = useState(null)
  const [predictions, setPredictions] = useState([])

  const MODEL_HTTP_URL = "api/pokeml/classify"
  const MODEL_INDEXEDDB_URL = "indexeddb://poke-model"


  useEffect(() => {
    async function fetchModel() {
      try {
        const localClassifierModel = await tf.loadLayersModel(
          MODEL_INDEXEDDB_URL
        )

        setModel(localClassifierModel)
        console.log("Model loaded from IndexedDB")
      } catch (e) {
        try {
          const classifierModel = await tf.loadLayersModel(MODEL_HTTP_URL)
          setModel(classifierModel)

          await classifierModel.save(MODEL_INDEXEDDB_URL)
          console.log("Model saved to IndexedDB")
        } catch (e) {
          console.log("Unable to load model at all: ", e)
        }
      }
    }
    fetchModel()
  }, [])

  useEffect(() => {
    async function makePredictions() {
      if (imageRef && model) {
        console.log(
          "Uploaded Image from inside the useEffect",
          findState.uploadedImage
        )
        console.log("ImageRef from inside the useEffect", imageRef.current)
        try {
          const imgTensor = tf.browser
            .fromPixels(imageRef.current)
            .resizeNearestNeighbor([160, 160])
            .toFloat()
            .sub(127)
            .div(127)
            .expandDims()

          const y_pred = await model.predict(imgTensor).data()
          console.log(y_pred)
          console.log(pokemonState)

          const topkPredNames = getTopKPredPokeObj(y_pred, 6, pokemonState)

          dispatch(storePredictions({ predictions: topkPredNames }))

          console.log(topkPredNames)

          return topkPredNames
        } catch (e) {
          console.log("Unable to run predictions.", e)
        }
      }
    }
    makePredictions()
  }, [findState.uploadedImage])

  return (
    
          {findState.uploadedImage && (
            
          )}
    
  )
}

export default SearchOutput

user_6138713 · Answer

I've read about this issue before and I haven't encountered this myself but this should help you.

The useRef hook can be a trap for your custom hook, if you combine it with a useEffect that skips rendering. Your first instinct will be to add ref.current to the second argument of useEffect, so it will update once the ref changes. But the ref isn’t updated till after your component has rendered — meaning, any useEffect that skips rendering, won’t see any changes to the ref before the next render pass.

As you can see on the behavior that you are encountering, the imageRef.current inside your useEffect is not returning the updated object which it is logging the previous value.

A logical way to this problem is to do something like:

  useEffect(() => {
    async function makePredictions() {
      //...
    }
    makePredictions()
  }, [findState.uploadedImage, imageRef.current])

But the problem with this is that changes in imageRef.current doesn't trigger a render in React. So based on React documentation on how to measure a DOM node, you should use useCallback instead of useRef.

Something like this should work:

  const [imageRef, setImageRef] = useState(null);
  const onChangeRef = useCallback(node => {
    // ref value changed to node
    setImageRef(node); // e.g. change ref state to trigger re-render
    if (node === null) { 
      // node is null, if DOM node of ref had been unmounted before
    } else {
      // ref value exists
    }
  }, []);
  
  useEffect(() => {
    async function makePredictions() {
      if (imageRef && model) {
        try {
          const imgTensor = tf.browser
            .fromPixels(imageRef)
            .resizeNearestNeighbor([160, 160])
            .toFloat()
            .sub(127.5)
            .div(127.5)
            .expandDims()

          const y_pred = await model.predict(imgTensor).dataSync()
          const topkPredNames = getTopKPred(y_pred, 5)

          console.log(topkPredNames)
          return topkPredNames
        } catch (e) {
          console.log("Unable to run predictions.")
        }
      }
    }
    makePredictions()
  }, [findState.uploadedImage, imageRef])

Instead of , you should use .

References:

Ref objects inside useEffect Hooks
How can I measure a DOM node?

user_13798886 · Answer

On first sight it looks like the cause could be the missing dependency of "model" in the useEffect that holds the makePredictions function. Since you are using model, and the effect is only triggered when findState.uploadedImage is updated. imageRef is a reference and therefore it doesnt have to be a dependency. "model" has to be since its a chaning state thorughout your flow.

  useEffect(() => {
    async function makePredictions() {
      if (imageRef && model) {
        try {
          const imgTensor = tf.browser
            .fromPixels(imageRef)
            .resizeNearestNeighbor([160, 160])
            .toFloat()
            .sub(127.5)
            .div(127.5)
            .expandDims()

          const y_pred = await model.predict(imgTensor).data()
          const topkPredNames = getTopKPred(y_pred, 5)

          console.log(topkPredNames)
          return topkPredNames
        } catch (e) {
          console.log("Unable to run predictions.", e)
        }
      }
    }
    makePredictions()
  }, [findState.uploadedImage, model]) // added model and remove imageRef

I hope that helps. I didnt fully try to understand exactly whats going on but that caught my eye and usually causes such errors.

React state is always one step behind while making predictions on uploaded image using Tensorflowjs

2 Answers