Search code examples
tensorflowbuttonclassificationdisplaymobilenet

Random Image and Classification with p5.js, TensorFlow and MobileNet


I'm asking for help from the community because I'm stuck with my code. For a project I need to create an image classification program on P5.js using MobileNet and TensorFlow.

I wanted to know how to create a button that randomly displays an image (among the 3 I downloaded) with the MobileNet classification below the image and the information I gave (like "Danger")?

Result wanted

Here is what i've already done. Thank you :)


    var model_mobilenet;
    var loaded;
    
    
    function setup() {
      createCanvas(400, 300);
      mobilenet.load().then(modelLoaded);
    
      military = loadImage('military');
      rifle = loadImage('rifle');
      revolver = loadImage('revolver');
      
      createButton("Take a picture").mousePressed(btnClicked);
    }
    
    function classifyDone(res) {
      print(res);
      if (res[0].className == "rifle")
      {
      print("Danger");
      createP("Danger");
      }
      
      else if (res[0].className == "revolver, six-gun, six-shooter")
      {
      print("Danger");
      createP("Danger");
      }
      else if (res[0].className == "military uniform")
      {
        print("Not sure about this one");
        createP("Not sure about this one");
      }
      else {
        print("Evrything is okay");
        createP("Everything is okay");
      }
    }
    
    function modelLoaded(net) {
      model_mobilenet = net;
      loaded = true;
      print("Model loaded");
    }
    
    function btnClicked() {
      
    image(military, 0, 0, 400, 300);
      
    if (loaded == true)
    {
    model_mobilenet.classify(military.elt).then(classifyDone);
      }
    }


Solution

  •     var model_mobilenet;
    var loaded;
    
    
    function setup() {
      createCanvas(400, 300);
      mobilenet.load().then(modelLoaded);
    
      revolver = loadImage('revolver.jpg');
      baseball = loadImage('baseball.jpg');
      rifle = loadImage('rifle.jpg');
    
      
      createButton("DANGER OR NOT?").mousePressed(btnClicked);
    }
    
    function classifyDone(res) {
      print(res);
      
      createP("Model detected a <b>" + res[0].className + "</b> with a confidence of <b>" + res[0].probability + "</b>");
      if (res[0].className =="assault rifle, assault gun")
      
      { createP("<b>Danger, find a shelter, quickly !</b>");}
      
      else if (res[0].className == "revolver, six-gun, six-shooter")
     
      { createP("<b>Danger, find a shelter, quickly !</b>");}
    
      else {   createP("Everything is okay");  }
    }
    
    /* if ($(".revolver, six-gun, six-shooter").css('font-color: red')){
    $(".assault rifle, assault gun").css('font-color: green'); */
    
    function modelLoaded(net) {
      model_mobilenet = net;
      loaded = true;
      print("Model loaded");
    }
    
    function btnClicked() {
     
    let randomNumber = Math.random()
    if(randomNumber<0.3)
    {
      image(baseball, 0, 0, 400, 300);
    model_mobilenet.classify(baseball.canvas).then(classifyDone);
      }
    else if (randomNumber<0.6)
    {
        image(rifle, 0, 0, 400, 300);
    model_mobilenet.classify(rifle.canvas).then(classifyDone);
      }
      else 
    {
        image(revolver, 0, 0, 400, 300); 
    model_mobilenet.classify(revolver.canvas).then(classifyDone);
      }
    }
    

    Here is the result I obtained. It's not really great but with a little more CSS, it can be satisfying. Feel free to ask me for more details. Here is the result.