-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.swift
executable file
·154 lines (132 loc) · 6.16 KB
/
main.swift
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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
import CoreImage
import CoreGraphics
import Foundation
import AppKit
import CoreServices
import ImageIO
func extractGainMap(inputFilePath: String, outputFilePath: String?) -> (CGImageMetadata?, [String: Any]?) {
let imageURL = URL(fileURLWithPath: inputFilePath)
guard let source = CGImageSourceCreateWithURL(imageURL as CFURL, nil) else {
print("Error creating image source.")
return (nil, nil)
}
let gainmap = CGImageSourceCopyAuxiliaryDataInfoAtIndex(source, 0, kCGImageAuxiliaryDataTypeHDRGainMap)
let gainDict = NSDictionary(dictionary: gainmap ?? [:])
let gainData = gainDict[kCGImageAuxiliaryDataInfoData] as! Data
let gainDescription = gainDict[kCGImageAuxiliaryDataInfoDataDescription] as! [String: Any]
let gainMeta = gainDict[kCGImageAuxiliaryDataInfoMetadata] as! CGImageMetadata
let xmpMetadata = String(data: CGImageMetadataCreateXMPData(gainMeta, [:] as CFDictionary)! as Data, encoding: .utf8)
print("XMP Metadata from kCGImageAuxiliaryDataInfoMetadata: \(xmpMetadata!)")
if (outputFilePath != nil) {
let bitmapRep = NSBitmapImageRep(
bitmapDataPlanes: nil,
pixelsWide: Int(gainDescription["Width"] as! Int32),
pixelsHigh: Int(gainDescription["Height"] as! Int32),
bitsPerSample: 8,
samplesPerPixel: 1,
hasAlpha: false,
isPlanar: false,
colorSpaceName: .deviceWhite,
bytesPerRow: Int(gainDescription["BytesPerRow"] as! Int32),
bitsPerPixel: 8
)
gainData.copyBytes(to: bitmapRep!.bitmapData!, count: gainData.count)
if let pngData = bitmapRep!.representation(using: .png, properties: [:]) {
do {
try pngData.write(to: URL(fileURLWithPath: outputFilePath!))
print("Image saved to \(outputFilePath!)")
} catch {
print("Error saving image: \(error)")
}
} else {
print("Failed to create PNG data.")
}
}
return (gainMeta, gainDescription)
}
func imageDataFromCGImage(cgImage: CGImage) -> Data? {
let width = cgImage.width
let height = cgImage.height
let colorSpace = CGColorSpaceCreateDeviceGray()
let bytesPerPixel = 1
let bytesPerRow = bytesPerPixel * width
let bitsPerComponent = 8
var rawData = [UInt8](repeating: 0, count: height * bytesPerRow)
guard let context = CGContext(data: &rawData,
width: width,
height: height,
bitsPerComponent: bitsPerComponent,
bytesPerRow: bytesPerRow,
space: colorSpace,
bitmapInfo: CGImageAlphaInfo.none.rawValue) else {
return nil
}
context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))
return Data(rawData)
}
func writeAuxiliaryDataToImage(inputImagePath: String, gainMapPath: String, gainMeta: CGImageMetadata, gainDescription: [String: Any], outputImagePath: String) {
// Load the target image
let inputImageURL = URL(fileURLWithPath: inputImagePath)
guard let source = CGImageSourceCreateWithURL(inputImageURL as CFURL, nil) else {
print("Error creating image source.")
return
}
var imageProperties = CGImageSourceCopyPropertiesAtIndex(source, 0, nil) as? [String: Any] ?? [:]
var makerApple = imageProperties[kCGImagePropertyMakerAppleDictionary as String] as? [String: Any] ?? [:]
makerApple["33"] = 0.8
makerApple["48"] = 0.0
imageProperties[kCGImagePropertyMakerAppleDictionary as String] = makerApple
// Load the gain map image
let gainMapURL = URL(fileURLWithPath: gainMapPath)
guard let gainMapSource = CGImageSourceCreateWithURL(gainMapURL as CFURL, nil),
let gainMapImage = CGImageSourceCreateImageAtIndex(gainMapSource, 0, nil) else {
print("Error loading gain map image.")
return
}
// Prepare the auxiliary data
var mutableGainDescription = gainDescription
mutableGainDescription["Width"] = gainMapImage.width
mutableGainDescription["Height"] = gainMapImage.height
mutableGainDescription["BytesPerRow"] = gainMapImage.bytesPerRow
let gainMapImageData = imageDataFromCGImage(cgImage: gainMapImage)!
let auxDataInfo: [String: Any] = [
kCGImageAuxiliaryDataInfoData as String: gainMapImageData,
kCGImageAuxiliaryDataInfoDataDescription as String: mutableGainDescription,
kCGImageAuxiliaryDataInfoMetadata as String: gainMeta
]
// Create a new image destination
let outputImageURL = URL(fileURLWithPath: outputImagePath)
guard let destination = CGImageDestinationCreateWithURL(outputImageURL as CFURL, CGImageSourceGetType(source)!, 1, nil) else {
print("Error creating image destination.")
return
}
// Add the image and auxiliary data
CGImageDestinationAddImageFromSource(destination, source, 0, imageProperties as CFDictionary)
CGImageDestinationAddAuxiliaryDataInfo(destination, kCGImageAuxiliaryDataTypeHDRGainMap, auxDataInfo as CFDictionary)
// Finalize and save the image
if !CGImageDestinationFinalize(destination) {
print("Failed to write image to disk.")
} else {
print("Image successfully saved to \(outputImagePath)")
}
}
// This script uses a reference photo from iPhone to construct data structures
let referenceImagePath = "reference.heic"
// The SDR image
let sdrImagePath = "white.jpg"
let gainMapImagePath = "gain_map.png"
let outputHDRImagePath = "HDR.jpg"
// Extract data structures from a reference image
let (gainMeta, gainDescription) = extractGainMap(inputFilePath: referenceImagePath, outputFilePath: nil)
if let gainMeta = gainMeta, let gainDescription = gainDescription {
// Call the writeAuxiliaryDataToImage function
writeAuxiliaryDataToImage(
inputImagePath: sdrImagePath,
gainMapPath: gainMapImagePath,
gainMeta: gainMeta,
gainDescription: gainDescription,
outputImagePath: outputHDRImagePath
)
} else {
print("Failed to extract gain map metadata or description.")
}