CHGenerator/Sources/Generator/Processing/ImageGenerator.swift
2022-12-04 19:15:22 +01:00

286 lines
10 KiB
Swift

import Foundation
import AppKit
import CryptoKit
import Darwin.C
final class ImageGenerator {
private let imageOptimSupportedFileExtensions: Set<String> = ["jpg", "png", "svg"]
private let imageOptimizationBatchSize = 50
/**
The path to the input folder.
*/
private let input: URL
/**
The path to the output folder
*/
private let output: URL
private let imageReader: ImageReader
/**
All warnings produced for images during generation
*/
private var imageWarnings: Set<String> = []
/**
All images modified or created during this generator run.
*/
private var generatedImages: Set<String> = []
/**
The images optimized by ImageOptim
*/
private var optimizedImages: Set<String> = []
private var numberOfGeneratedImages = 0
private let numberOfTotalImages: Int
private lazy var numberImagesToCreate = jobs.reduce(0) { $0 + $1.images.count } + multiJobs.count * 2
private var numberOfImagesToOptimize = 0
private var numberOfOptimizedImages = 0
private let images: ImageData
private lazy var jobs: [(source: String, images: [ImageJob])] = images.jobs
.sorted { $0.key < $1.key }
.map { (source: $0.key, images: $0.value) }
.filter {
// Only load image if required
let imageHasChanged = imageReader.imageHasChanged(at: $0.source)
return imageHasChanged || $0.images.contains { job in
job.alwaysGenerate || !output.appendingPathComponent(job.destination).exists
}
}
private lazy var multiJobs: [String : String] = {
let imagesToGenerate: Set<String> = jobs.reduce([]) { $0.union($1.images.map { $0.destination }) }
return images.multiJobs.filter { imagesToGenerate.contains($0.key) }
}()
init(input: URL, output: URL, reader: ImageReader, images: ImageData) {
self.input = input
self.output = output
self.imageReader = reader
self.images = images
self.numberOfTotalImages = images.jobs.reduce(0) { $0 + $1.value.count }
+ images.multiJobs.count * 2
}
func generateImages() {
var notes: [String] = []
func addIfNotZero(_ count: Int, _ name: String) {
guard count > 0 else {
return
}
notes.append("\(count) \(name)")
}
addIfNotZero(images.missing.count, "missing images")
addIfNotZero(images.unreadable.count, "unreadable images")
print(" Changed sources: \(jobs.count)/\(images.jobs.count)")
print(" Total images: \(numberOfTotalImages) (\(numberOfTotalImages - imageReader.numberOfFilesAccessed) versions)")
print(" Warnings: \(images.warnings.count)")
if !notes.isEmpty {
print(" Notes: " + notes.joined(separator: ", "))
}
for (source, jobs) in jobs {
create(images: jobs, from: source)
}
for (baseImage, source) in multiJobs {
createMultiImages(from: baseImage, path: source)
}
print(" Generated images: \(numberOfGeneratedImages)/\(numberImagesToCreate)")
optimizeImages()
print(" Optimized images: \(numberOfOptimizedImages)/\(numberOfImagesToOptimize)")
}
private func create(images: [ImageJob], from source: String) {
guard let image = imageReader.getImage(atPath: source) else {
// TODO: Add to failed images
didGenerateImage(count: images.count)
return
}
let jobs = imageReader.imageHasChanged(at: source) ? images : images.filter(isMissing)
// Update all images
jobs.forEach { job in
// Prevent memory overflow due to repeated NSImage operations
autoreleasepool {
create(job: job, from: image, source: source)
didGenerateImage()
}
}
}
private func isMissing(_ job: ImageJob) -> Bool {
job.alwaysGenerate || !output.appendingPathComponent(job.destination).exists
}
private func create(job: ImageJob, from image: NSImage, source: String) {
let destinationUrl = output.appendingPathComponent(job.destination)
create(job: job, from: image, source: source, at: destinationUrl)
}
private func create(job: ImageJob, from image: NSImage, source: String, at destinationUrl: URL) {
// Ensure that image file is supported
let ext = destinationUrl.pathExtension.lowercased()
guard ImageType(fileExtension: ext) != nil else {
fatalError()
}
let destinationExtension = destinationUrl.pathExtension.lowercased()
guard let type = ImageType(fileExtension: destinationExtension)?.fileType else {
addWarning("Invalid image extension \(destinationExtension)", job: job)
return
}
let desiredWidth = CGFloat(job.width)
let sourceRep = image.representations[0]
let destinationSize = NSSize(width: sourceRep.pixelsWide, height: sourceRep.pixelsHigh)
.scaledDown(to: desiredWidth)
// create NSBitmapRep manually, if using cgImage, the resulting size is wrong
let rep = NSBitmapImageRep(bitmapDataPlanes: nil,
pixelsWide: Int(destinationSize.width),
pixelsHigh: Int(destinationSize.height),
bitsPerSample: 8,
samplesPerPixel: 4,
hasAlpha: true,
isPlanar: false,
colorSpaceName: NSColorSpaceName.deviceRGB,
bytesPerRow: Int(destinationSize.width) * 4,
bitsPerPixel: 32)!
let ctx = NSGraphicsContext(bitmapImageRep: rep)
NSGraphicsContext.saveGraphicsState()
NSGraphicsContext.current = ctx
image.draw(in: NSMakeRect(0, 0, destinationSize.width, destinationSize.height))
ctx?.flushGraphics()
NSGraphicsContext.restoreGraphicsState()
// Get NSData, and save it
guard let data = rep.representation(using: type, properties: [.compressionFactor: NSNumber(value: job.quality)]) else {
addWarning("Failed to get data", job: job)
return
}
do {
try data.createFolderAndWrite(to: destinationUrl)
} catch {
addWarning("Failed to write image (\(error))", job: job)
return
}
generatedImages.insert(job.destination)
}
private func addWarning(_ message: String, destination: String, path: String) {
let warning = " \(destination): \(message) required by \(path)"
imageWarnings.insert(warning)
}
private func addWarning(_ message: String, job: ImageJob) {
addWarning(message, destination: job.destination, path: job.path)
}
private func createMultiImages(from source: String, path: String) {
guard generatedImages.contains(source) else {
didGenerateImage(count: 2)
return
}
let sourceUrl = output.appendingPathComponent(source)
let sourcePath = sourceUrl.path
guard sourceUrl.exists else {
addWarning("No image at path \(sourcePath)", destination: source, path: path)
didGenerateImage(count: 2)
return
}
let avifPath = source.dropAfterLast(".") + ".avif"
createAVIF(at: output.appendingPathComponent(avifPath).path, from: sourcePath)
generatedImages.insert(avifPath)
didGenerateImage()
let webpPath = source.dropAfterLast(".") + ".webp"
createWEBP(at: output.appendingPathComponent(webpPath).path, from: sourcePath)
generatedImages.insert(webpPath)
didGenerateImage()
compress(at: source)
}
private func createAVIF(at destination: String, from source: String, quality: Int = 55, effort: Int = 5) {
let folder = destination.dropAfterLast("/")
let command = "npx avif --input=\(source) --quality=\(quality) --effort=\(effort) --output=\(folder) --overwrite"
do {
_ = try safeShell(command)
} catch {
addWarning("Failed to create AVIF image", destination: destination, path: destination)
}
}
private func createWEBP(at destination: String, from source: String, quality: Int = 75) {
let command = "cwebp \(source) -q \(quality) -o \(destination)"
do {
_ = try safeShell(command)
} catch {
addWarning("Failed to create WEBP image", destination: destination, path: destination)
}
}
private func compress(at destination: String, quality: Int = 70) {
let command = "magick convert \(destination) -quality \(quality)% \(destination)"
do {
_ = try safeShell(command)
} catch {
addWarning("Failed to compress image", destination: destination, path: destination)
}
}
private func optimizeImages() {
let all = generatedImages
.filter { imageOptimSupportedFileExtensions.contains($0.lastComponentAfter(".")) }
.map { output.appendingPathComponent($0).path }
for i in stride(from: 0, to: all.count, by: imageOptimizationBatchSize) {
let endIndex = min(i+imageOptimizationBatchSize, all.count)
let batch = all[i..<endIndex]
if optimizeImageBatch(batch) {
optimizedImages.formUnion(batch)
}
didOptimizeImage(count: batch.count)
}
}
private func optimizeImageBatch(_ batch: ArraySlice<String>) -> Bool {
let command = "imageoptim " + batch.joined(separator: " ")
do {
_ = try safeShell(command)
return true
} catch {
addWarning("Failed to optimize images", destination: "", path: "")
return false
}
}
// MARK: Output
private func didGenerateImage(count: Int = 1) {
numberOfGeneratedImages += count
print(" Generated images: \(numberOfGeneratedImages)/\(numberImagesToCreate) \r", terminator: "")
fflush(stdout)
}
private func didOptimizeImage(count: Int) {
numberOfOptimizedImages += count
print(" Optimized images: \(numberOfOptimizedImages)/\(numberOfImagesToOptimize) \r", terminator: "")
fflush(stdout)
}
}