<!DOCTYPE html>
<html>
  <head>
    <meta charset="UTF-8">
    <title>Coral USB Accelerator + WebUSB</title>
    <script src="dfu.js"></script>
    <script src="dfuse.js"></script>
    <script src="models.js"></script>
    <script src="tflite.js"></script>
    <script src="interpreter.js"></script>
  </head>
  <body>
    <h1>Coral USB Accelerator Demo</h1>
    <button id="button-firmware">1. Flash Device Firmware</button>&nbsp;<span id="firmware-status"></span><br/><br/>
    <button id="button-init">2. Initialize Interpreter</button>
    <select id="model">
      <option value="mobilenet_cpu">MobileNet V1 (CPU)</option>
      <option value="mobilenet_tpu">MobileNet V1 (TPU)</option>
      <option value="ssd_mobilenet_coco_cpu">SSD MobileNet V2 - Coco (CPU)</option>
      <option value="ssd_mobilenet_coco_tpu">SSD MobileNet V2 - Coco (TPU)</option>
      <option value="ssd_mobilenet_face_cpu">SSD MobileNet V2 - Face (CPU)</option>
      <option value="ssd_mobilenet_face_tpu">SSD MobileNet V2 - Face (TPU)</option>
    </select><br/><br/>
    <button id="button-image" onclick="document.getElementById('file').click();" disabled>3. Choose Image File</button><br/><br/>
    <input type="file" id="file" style="display:none;"/><br/>
    <h2 id="result"></h2>
    <canvas id="canvas"></canvas>
    <br/><br/>

    <!--button id="test">Inference Test</button><br/><br/-->
    <script>
      async function loadLocalImage(file) {
        return new Promise(resolve => {
          var reader = new FileReader();
          reader.onload = function() {
            var img = new Image();
            img.onload = function() {
              resolve(img);
            };
            img.src = reader.result;
          };
          reader.readAsDataURL(file);
        });
      }

      async function getDFUDescriptorProperties(device) {
        let data = await device.readConfigurationDescriptor(0);
        let configDesc = dfu.parseConfigurationDescriptor(data);
        let funcDesc = null;
        if (configDesc.bConfigurationValue == device.settings.configuration.configurationValue) {
            for (let desc of configDesc.descriptors) {
                if (desc.bDescriptorType == 0x21 && desc.hasOwnProperty("bcdDFUVersion")) {
                    funcDesc = desc;
                    break;
                }
            }
        }

        if (!funcDesc) return {};

        return {
            WillDetach:            ((funcDesc.bmAttributes & 0x08) != 0),
            ManifestationTolerant: ((funcDesc.bmAttributes & 0x04) != 0),
            CanUpload:             ((funcDesc.bmAttributes & 0x02) != 0),
            CanDnload:             ((funcDesc.bmAttributes & 0x01) != 0),
            TransferSize:          funcDesc.wTransferSize,
            DetachTimeOut:         funcDesc.wDetachTimeOut,
            DFUVersion:            funcDesc.bcdDFUVersion
        };
      }

      async function connect(device) {
          await device.open();

          let desc = await getDFUDescriptorProperties(device);
          if (desc && Object.keys(desc).length > 0)
            device.properties = desc;

          device.logDebug = console.log;
          device.logInfo = console.log;
          device.logWarning = console.log;
          device.logError = console.log;
          device.logProgress = console.log;

          return device;
      }

      async function openDevice() {
        try {
          let device = await navigator.usb.requestDevice({ 'filters': [{'vendorId': 0x1a6e, 'productId': 0x089a}] });
          let interfaces = dfu.findDeviceDfuInterfaces(device);
          if (interfaces.length != 1) return null;
          return await connect(new dfu.Device(device, interfaces[0]));
        } catch (error) {
          return null;
        }
      }

      async function loadFile(url) {
        return new Promise(resolve => {
          let req = new XMLHttpRequest();
          req.open('GET', url, true);
          req.responseType = 'arraybuffer';
          req.onload = function (event) { resolve(req.response); };
          req.send(null);
        });
      }

      document.addEventListener('DOMContentLoaded', event => {
        let downloadButton = document.querySelector("#button-firmware");
        downloadButton.addEventListener('click', async function(event) {
          let firmwareFile = await loadFile('/firmware.bin');
          let device = await openDevice();
          if (device) {
            console.log('Device: ', device.properties);
          } else {
            console.log('Cannot open device.');
          }

          if (device && firmwareFile) {
            document.getElementById("firmware-status").textContent = "Flashing...";

            try {
              let status = await device.getStatus();
              if (status.state == dfu.dfuERROR) {
                await device.clearStatus();
              }
            } catch (error) {
              device.logWarning("Failed to clear status");
            }

            try {
              let manifestationTolerant = device.properties.ManifestationTolerant;
              await device.do_download(device.properties.TransferSize, firmwareFile, manifestationTolerant);
              if (!manifestationTolerant)
                await device.waitDisconnected(5000);
              document.getElementById("firmware-status").textContent = "Ready :)";
            } catch (error) {
              document.getElementById("firmware-status").textContent = "Failed :(";
            }
          }
        });
      });

      Module['onRuntimeInitialized'] = () => {
        let inference_test = Module.cwrap("inference_test",  null, [], { async: true });
        let testButton = document.querySelector("#test");
        if (testButton) {
          testButton.addEventListener('click', async () => {
            inference_test();
          });
        }

        let interpreter;
        let model;

        document.querySelector("#button-init").addEventListener('click', async () => {
          model = TFLITE_MODELS[document.getElementById('model').value];
          console.log(model);
          interpreter = new tflite.Interpreter();
          if (await interpreter.createFromBuffer(await loadFile(model.url))) {
            document.getElementById("button-firmware").disabled = true;
            document.getElementById("button-init").disabled = true;
            document.getElementById("model").disabled = true;
            document.getElementById("button-image").disabled = false;
          }
        });

        document.querySelector("#file").addEventListener('change', async () => {
          const input = document.getElementById("file");
          const file = input.files[0];
          if (!file) return;

          [_, height, width, _] = interpreter.inputShape(0);
          console.log('SHAPE:', width, 'x', height);

          const img = await loadLocalImage(file);
          const c = document.getElementById("canvas");
          c.width = width;
          c.height = height;
          const ctx = c.getContext('2d');
          const color = 'red';
          ctx.strokeStyle = color;
          ctx.fillStyle = color;
          ctx.textBaseline = 'top';

          switch (model.type) {
            case 'classification':
              ctx.drawImage(img, 0, 0, img.width, img.height,
                                 0, 0, width, height);
              break;
            case 'detection':
              const alpha = Math.min(width / img.width, height / img.height);
              ctx.drawImage(img, 0, 0, img.width, img.height,
                                 0, 0, alpha * img.width, alpha * img.height);
              break;
          }

          const imageData = ctx.getImageData(0, 0, width, height);
          tflite.setRgbaInput(interpreter, imageData.data);
          document.getElementById("result").textContent = "Recognizing...";
          const inferenceStart = Date.now();
          await interpreter.invoke();
          const inferenceTime = Date.now() - inferenceStart;

          let label = null;
          switch (model.type) {
            case 'classification':
              const maxIndex = tflite.getClassificationOutput(interpreter);
              console.log('Class:', maxIndex);
              label = model.labels[maxIndex];
              break;
            case 'detection':
              const objects = tflite.getDetectionOutput(interpreter, threshold=0.5);
              console.log('Detected objects:', objects);
              for (obj of objects) {
                console.log(obj);
                const x = obj.bbox.xmin * width;
                const y = obj.bbox.ymin * height;
                const w = obj.bbox.xmax * width - x;
                const h = obj.bbox.ymax * height - y;
                ctx.strokeRect(x, y, w, h);
                ctx.fillText(model.labels[obj.id], x + 5, y + 5);
              }
              label = `${objects.length} ${objects.length == 1 ? 'object' : 'objects'}`;
              break;
          }
          document.getElementById("result").textContent = `${label}: ${inferenceTime} ms`;
        });
      };
      Module['print'] = txt => console.log(txt);
    </script>
  </body>
</html>