Add a fitted notification badge to a UIImage

It is common to use a notification badge to alert the user to unread or new features, accessible when pressing a button. Here, I show a method of adding this badge (e.g. a white text number inside a red circle) for any image. The size of the badge circle is defined relative to the original image, and a looping font fitting approach is used to ensure the text fits inside the circle.

static func addNotificationBadgeToImage(source: UIImage, text: String) -> UIImage? {

    //define the badge as half the size of the original image, and located in the top-right corner
    let iconSizeFraction: CGFloat = 0.5
    let iconOffset = CGPoint(x: 0.5, y: 0.0)

    let diameter: CGFloat = iconSizeFraction * source.size.width
    let iconSize = CGSize(width: diameter, height: diameter)

    // define the new size of the original image
    let newSize = CGSize(
        width: max(source.size.width, iconOffset.x + iconSize.width) - min(0.0, iconOffset.x),
        height: max(source.size.height, iconOffset.y + iconSize.height) - min(0.0, iconOffset.y)
    )

    // start drawing
    UIGraphicsBeginImageContextWithOptions(newSize, false, 0.0)
    let context = UIGraphicsGetCurrentContext()!
    context.saveGState()

    // put the origin of the coordinate system at the top left
    if let cgImage = source.cgImage {
        context.translateBy(x: 0, y: source.size.height)
        context.scaleBy(x: 1.0, y: -1.0)
        let rect = CGRect(x: 0.0, y: 0.0, width: source.size.width, height: source.size.height)
        context.draw(cgImage, in: rect)
    }
    context.restoreGState()

    // define the rectangle in which the badge should be drawn
    let badgeRect = CGRect(
        x: source.size.width * iconOffset.x,
        y: source.size.height * iconOffset.y,
        width: iconSize.width,
        height: iconSize.height
    )

    // draw the red badge circle
    if let circle = drawCircle(diameter: diameter, color: UIColor.red)?.cgImage {
        context.draw(circle, in: badgeRect)
    }

    let image = UIGraphicsGetImageFromCurrentImageContext()!
    UIGraphicsEndImageContext()

    // add the label to this new image
    if let imageWIthText = DrawClass.textToImage(
        text: text,
        image: image,
        rect: badgeRect) { return imageWIthText }
    return image
}

A simple function to draw a circle with solid color fill with a specified diameter

static func drawCircle(diameter: CGFloat, color: UIColor) -> UIImage? {
    UIGraphicsBeginImageContextWithOptions(CGSize(width: diameter, height: diameter), false, 0.0)
    let ctx = UIGraphicsGetCurrentContext()!
    ctx.saveGState()
    ctx.setFillColor(color.cgColor)
    let rect = CGRect(x: 0, y: 0, width: diameter, height: diameter)
    ctx.fillEllipse(in: rect)
    ctx.restoreGState()
    let image = UIGraphicsGetImageFromCurrentImageContext()!
    UIGraphicsEndImageContext()
    return image
}

This function draws the text in the specified rect.

static func textToImage(text: String, image: UIImage, rect: CGRect) -> UIImage? {

    // initialize new image with source image
    let scale = UIScreen.main.scale
    UIGraphicsBeginImageContextWithOptions(image.size, false, scale)
    image.draw(in: CGRect(origin: CGPoint.zero, size: image.size))

    // set limits on text area and find matching font
    let padding = 0.075 * rect.size.width
    let textWidth = rect.size.width - 2 * padding

    // use a UIFont extension to find the optimal font for the given size constraints
    // e.g. https://qiita.com/cayozin/items/f574fa803eeeb6b3310a
    guard let font = UIFont.createFittedFont(maxWidth: textWidth, maxHeight: textWidth, text: text, fontName: "Helvetica Bold") else {
        return nil
    }

    // find the equivalent size in pixels and use the text height to center it in the circle
    let size = (text as NSString).size(withAttributes: [NSAttributedStringKey.font : font])
    let xOffset = padding
    let yOffset = 0.5 * rect.size.height - 0.5 * size.height
    let rect = CGRect(
        origin: CGPoint(x: rect.origin.x + xOffset, y: rect.origin.y + yOffset),
        size: CGSize(width: textWidth, height: font.lineHeight)
    )

    // define properties of the text to be drawn
    let textColor = UIColor.white
    let paragraph = NSMutableParagraphStyle()
    paragraph.alignment = .center
    let textFontAttributes = [
        NSAttributedStringKey.font: font,
        NSAttributedStringKey.paragraphStyle: paragraph,
        NSAttributedStringKey.foregroundColor: textColor,
        ] as [NSAttributedStringKey : Any]
    text.draw(in: rect, withAttributes: textFontAttributes)

    let newImage = UIGraphicsGetImageFromCurrentImageContext()
    UIGraphicsEndImageContext()
    return newImage
}

Finally, we need to calculate the correct font size to ensure the text fits inside the specified rect. I recommend using my looping font fitting extension to UIFont to calculate this: