-
Notifications
You must be signed in to change notification settings - Fork 110
/
detection_utils.go
40 lines (34 loc) · 1.21 KB
/
detection_utils.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
package objectdetection
import (
"fmt"
"image"
"image/color"
"github.com/fogleman/gg"
"github.com/pkg/errors"
"go.viam.com/rdk/rimage"
)
// Overlay returns a color image with the bounding boxes overlaid on the original image.
func Overlay(img image.Image, dets []Detection) (image.Image, error) {
gimg := gg.NewContextForImage(img)
for _, det := range dets {
if !det.BoundingBox().In(img.Bounds()) {
return nil, errors.Errorf("bounding box (%v) does not fit in image (%v)", det.BoundingBox(), img.Bounds())
}
drawDetection(gimg, det)
}
return gimg.Image(), nil
}
// drawDetection overlays text of the image label and score in the upper left hand of the bounding box.
func drawDetection(img *gg.Context, d Detection) {
red := &color.NRGBA{255, 0, 0, 255}
box := d.BoundingBox()
rimage.DrawRectangleEmpty(img, *box, red, 2.0)
text := fmt.Sprintf("%s: %.2f", d.Label(), d.Score())
rimage.DrawString(img, text, image.Point{box.Min.X, box.Min.Y}, red, 30)
}
// OverlayText writes a string in the top of the image.
func OverlayText(img image.Image, text string) image.Image {
gimg := gg.NewContextForImage(img)
rimage.DrawString(gimg, text, image.Point{30, 30}, color.NRGBA{255, 0, 0, 255}, 30)
return gimg.Image()
}