ultralytics/yolov5でのAttributeErrorへの対応

ultralytics/yolov5を使用した際、以下のエラーが発生しました。

AttributeError: 'Detections' object has no attribute 'imgs'

これは、以下のissueにあるように、apiが変更されたために発生するようです。

https://github.com/robmarkcole/yolov5-flask/issues/23

一例ですが、以下のようにプログラムを書き換えることで、エラーが解消しました。

results = model(im)  # inference

# new
def getImage(results):
    output_dir = "static"
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    results.save(save_dir=f"{output_dir}/")
    return Image.open(f"{output_dir}/image0.jpg")

# old
def oldGetImage(results):
    results.render()
    return Image.fromarray(results.imgs[0])
    
renderedImg = getImage(results)

同様のことでお困りの方の参考になりましたら幸いです。