CoreML conversion/export and usage (non-max suppression)

[pytholic]

Hi and thank you for the repository.

I trained my model on custom dataset and now trying to export it as coreml model to use in iOS. However I am facing some difficulties.

1. When I convert my best.pt model using the provided script, I see a lot of "Adding op" like this:

scikit-learn version 0.20.0 is not supported. Minimum required version: 0.17. Maximum required version: 0.19.2. Disabling scikit-learn conversion API.
TensorFlow version 2.6.0 detected. Last version known to be fully compatible is 2.3.1 .
Keras version 2.6.0 detected. Last version known to be fully compatible of Keras is 2.2.4 .
Adding op 'pow' of type pow
Adding op 'pow_y_0' of type const
Adding op 'mul_1' of type mul
Adding op 'mul_1_x_0' of type const
Adding op 'add' of type add
Adding op 'mul_2' of type mul
Adding op 'mul_2_x_0' of type const
Adding op 'tanh' of type tanh
Adding op 'add_1' of type add
Adding op 'add_1_x_0' of type const
Adding op 'mul' of type mul
Adding op 'mul_x_0' of type const
Adding op 'mul_3' of type mul
Adding op 'pow' of type pow
Adding op 'pow_y_1' of type const
Adding op 'mul_1' of type mul
Adding op 'mul_1_x_1' of type const
.
.
.

Is it normal?

2. In the converted model, I lose the concatenation part (combing output of three levels) and also NMS. Then I manually implement NMS in swift but I really want it to be integrated in the model itself. Can you help with this? I think I am missing something during conversion step maybe.

3. When I test my code on iOS, in the end I have 10 boxes with highest scores. However, some of these boxes have exactly same score (even up to 7 decimal points) which I think should not be possible because boxes are different. So then when I apply NMS to keep only the best box, it doesn't really work because the some boxes have same score. Can you give some idea regarding this?

BoundingBox(classIndex: 0, score: 0.91258967, rect: (133.61363220214844, 171.43484497070312, 181.71127319335938, 233.848876953125))
BoundingBox(classIndex: 0, score: 0.91258967, rect: (218.47067260742188, 217.3438262939453, 75.99720764160156, 142.03091430664062))
BoundingBox(classIndex: 0, score: 0.91258967, rect: (260.2138977050781, 256.07952880859375, 56.5107421875, 64.55950164794922))
BoundingBox(classIndex: 0, score: 0.9090115, rect: (259.60150146484375, 224.28857421875, 57.72898483276367, 65.09613800048828))
BoundingBox(classIndex: 0, score: 0.8489735, rect: (188.3462371826172, 214.0371551513672, 73.1979751586914, 148.64938354492188))
BoundingBox(classIndex: 0, score: 0.8489735, rect: (229.7305908203125, 254.57789611816406, 54.429264068603516, 67.56790161132812))
BoundingBox(classIndex: 0, score: 0.81404513, rect: (218.50527954101562, 264.7089538574219, 107.66510009765625, 47.08483123779297))
BoundingBox(classIndex: 0, score: 0.7050861, rect: (290.6436462402344, 100.86282348632812, 106.90132141113281, 38.55246353149414))
BoundingBox(classIndex: 0, score: 0.7050861, rect: (326.1788330078125, 94.99613952636719, 51.8309440612793, 50.285823822021484))
BoundingBox(classIndex: 0, score: 0.7050861, rect: (343.8971252441406, 109.24378967285156, 32.39434051513672, 21.790523529052734))

I have also attached converted mlmodel image with properties. Thank you in advance!!

[github-actions]

👋 Hello @rajahaseeb147, thank you for your interest in YOLOv5 🚀! Please visit our ⭐️ Tutorials to get started, where you can find quickstart guides for simple tasks like Custom Data Training all the way to advanced concepts like Hyperparameter Evolution.

If this is a 🐛 Bug Report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset images, training logs, screenshots, and a public link to online W&B logging if available.

For business inquiries or professional support requests please visit https://ultralytics.com or email Glenn Jocher at glenn.jocher@ultralytics.com.

Requirements

Python>=3.6.0 with all requirements.txt installed including PyTorch>=1.7. To get started:

$ git clone https://github.com/ultralytics/yolov5
$ cd yolov5
$ pip install -r requirements.txt

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Google Colab and Kaggle notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training (train.py), validation (val.py), inference (detect.py) and export (export.py) on MacOS, Windows, and Ubuntu every 24 hours and on every commit.

I found that the CoreML export worked when using CoreMLTools 5.0b3 and set the model to eval in export.py. This will give four outputs, three for the detectors and one for the concatenated output (with 640px that's a 1 x 25200 x 85 MLMultiArray). No NMS though, I wrote that myself. With this setup I get 15FPS inference time on an A13 CPU (SE II), 45-50 FPS with 320px.

[pytholic]

@Softmatic Hi and thank you for your reply. I used 5.0b3 and now I have the concat layer. Did you upload your code to any repo? Would appreciate if I can take a look at your NMS part and how to process the concatenated output array. Thanks again!

[pytholic]

@Softmatic on a side note, i worked with this https://github.com/dbsystel/yolov5-coreml-tools repo and my final model includes NMS now!

No repo, but here's a snippet that I use for testing (775 is the concat layer in my model with yolov5s, yours will be different, use Netron to find it):

if let output = try? cml.prediction(image: pixelBuffer){
  let confidenceThreshold : Float = 0.3
  let output1D: [Float] = try Array(UnsafeBufferPointer<Float>(output._775))
  let rows = output._775.shape[1].intValue // 25200 @ 640x640 
  
  for i in 0..<rows{
    let confidence = output1D[(i * 85) + 4]
    if(confidence > confidenceThreshold){
      let row = Array(output1D[(i * 85)..<(i + 1) * 85])
      let classes = Array(row.dropFirst(5))
      let classIndex : Int = classes.firstIndex(of: classes.max() ?? 0) ?? 0
      let detection: [Float] = [row[0] - row[2]/2, row[1] - row[3]/2, row[2], row[3], confidence, Float(classIndex)]
      if detections[String(classIndex)] != nil{
        var existing = detections[String(classIndex)]!
        existing.append(detection)
        detections[String(classIndex)] = existing
      }
      else{
        detections[String(classIndex)] = [detection]
      }
    }
  }
}

I put the detections in a dict with the class no. as the key. That makes it easier to do NMS later for each class. Obviously, going through the rows one by one like this is sub-optimal. Should be possible to speed this up with either vDSP or a compute shader.

[pytholic]

@Softmatic Thanks a lot for sharing this and explaining. Appreciate it!!

As a follow up, some experimenting shows that the loop is considerably faster than using compactMap or map to filter the rows with a confidence above the threshold, e.g.:

let output1D: [Float] = try Array(UnsafeBufferPointer<Float>(output._775))

let candidates : [Int] = output1D.enumerated().compactMap { index, element in
           (index % 85) != 4 ? nil : element > confidenceThreshold ? (index - 4) / 85 : nil
          }

[pytholic]

@Softmatic Thanks a lot, I will check this out as well. I made my model work with the repo that I mentioned above. Now my model has NMS integrated into it. I will also use concatenated model with your NMS script and then compare results from both models. Appreciated it :))

Got it up to 60FPS (320x320) on an SE II with some optimization. Video.

[abhimanyuchadha96]

@Softmatic Great work, can you elaborate on your optimization?

[pytholic]

@Softmatic That looks pretty good! I still haven't tested my model on video feed, just images for now. May I ask what was the key optimization step?

@rajahaseeb147 Thanks! I had forked the TF Lite iOS sample as a starting point but found that their way of resizing the pixel buffer before inference took almost 20ms. Swapping this part for a CI-based solution reduced this to 2ms. I also do the overlays with a CALayer rather than a overlay UIView, that's another speedup. Finally, I made the size of the video capture depend on the model size; no point in capturing Full HD if you scale it down to 320x320 anyway.

BTW, I also tried CoreML INT8 quantization to see if inference is faster but to my surprise it didn't make any difference (just a smaller model). This is very different behaviour from TF Lite where you get a 2-3x speedup.