(Swift) Applying Style Transfer on iPhone
Published in
18 min readOct 1, 2023
目的：將三種不同風格的 Pytorch 模型轉換成 CoreML，並在 iphone 上進行 inference
Fast-Neural-Style-Transfer-Mac
Usage
git clone https://github.com/jasonyen1009/Fast-Neural-Style-Transfer-Mac.git
cd Fast-Neural-Style-Transfer-MacInference on Mac
python test_on_image.py --image_path food.jpg --checkpoint_model pre-trained/cuphead_10000.pth 
Convert model to CoreML
python modeltocoreml.py --checkpoint_model pre-trained/cuphead_10000.pth You can also convert the model to fp16 or int8
python modeltocoreml.py --checkpoint_model pre-trained/cuphead_10000.pth --int8
python modeltocoreml.py --checkpoint_model pre-trained/cuphead_10000.pth --fp16Converting to fp16 or int8 has little impact on the model and can reduce the space by 4 times.
Inference on Iphone
在上述文章中，可以觀察到將模型轉換成 int8 ，不會對模型的精度產生太大的影響。因此，在這次的專案中，我計劃將這三種模型都轉換為 int8格式。
新增 Xcode 專案
在 Xcode 新增專案後，進行簡單的佈局，右上角兩個都是 UIBarButtonItem ，左邊的是用來打開相簿，右邊的則是能將相片進行保存，下方就是由 UIButton 與 UIImage 所組成的三個風格轉換按鈕。
將三個轉換後的模型拖入到專案中
對模型進行簡單的測試，可以先選擇任一模型，然後點擊 Preview，接著將照片拖放到預覽介面中。
給予權限
在這次的專案中，會使用到相簿中的照片，還有將風格轉換後的照片保存到相簿中，需要新增下面的兩個權限：
Privacy — Photo Library Additions Usage Description
Privacy — Photo Library Usage Description
⭐️若是沒有添加 Privacy — Photo Library Additions Usage Description 會發現使用 UIActivityViewController 保存照片時，Save Image 欄位會消失。
import Vision
這次我將使用 Vision 來處理 CoreML 模型，由於我會使用三種不同的模型，我會首先建立一個名為 StyleModel 的列舉，以列出專案中的所有模型。
enum StyleModel {
case starryNight
case cuphead
case mosaic
}
// 選擇使用的模型
func initializeModel(_ style: StyleModel) {
let configuration = MLModelConfiguration()
switch style {
case .starryNight:
model = try? VNCoreMLModel(for: starry_night_int8(configuration: configuration).model)
case .cuphead:
model = try? VNCoreMLModel(for: cuphead_int8(configuration: configuration).model)
case .mosaic:
model = try? VNCoreMLModel(for: mosaic_int8(configuration: configuration).model)
}
}開啟相簿
擴充 ViewController，實作 UIImagePickerControllerDelegate 與UINavigationControllerDelegate 協議
extension ViewController: UIImagePickerControllerDelegate, UINavigationControllerDelegate {
func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) {
if let image = info[.originalImage] as? UIImage {
imageView.image = image
// 保留一份原始相片
originalImage = image
}
picker.dismiss(animated: true, completion: nil)
}
}為了避免照片不斷被重複覆蓋風格，必須透過 originalImage 保留最原始從相簿抓到的相片。
進行風格轉換
將 image 轉成 CVPixelBuffer
extension UIImage {
// 將UIImage轉換為CVPixelBuffer
func toPixelBuffer(pixelFormatType: OSType, width: Int, height: Int) -> CVPixelBuffer? {
var pixelBuffer: CVPixelBuffer?
let attrs: [String: NSNumber] = [
kCVPixelBufferCGImageCompatibilityKey as String: NSNumber(booleanLiteral: true),
kCVPixelBufferCGBitmapContextCompatibilityKey as String: NSNumber(booleanLiteral: true)
]
// 創建CVPixelBuffer
let status = CVPixelBufferCreate(kCFAllocatorDefault, width, height, pixelFormatType, attrs as CFDictionary, &pixelBuffer)
guard status == kCVReturnSuccess else {
return nil
}
// 鎖定CVPixelBuffer 的基地址
CVPixelBufferLockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
let pixelData = CVPixelBufferGetBaseAddress(pixelBuffer!)
// 創建CGContext
let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
let context = CGContext(data: pixelData, width: width, height: height, bitsPerComponent: 8, bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer!), space: rgbColorSpace, bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue)
// 調整座標系
context?.translateBy(x: 0, y: CGFloat(height))
context?.scaleBy(x: 1.0, y: -1.0)
// 繪製圖像
UIGraphicsPushContext(context!)
draw(in: CGRect(x: 0, y: 0, width: CGFloat(width), height: CGFloat(height)))
UIGraphicsPopContext()
// 解鎖基地址並返回CVPixelBuffer
CVPixelBufferUnlockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
return pixelBuffer
}
}將轉換後的圖片轉回原本的尺寸
// 將風格轉換後的圖片變回原本的尺寸
func resizeImage(uiImage: UIImage, newSize: CGSize) -> UIImage {
UIGraphicsBeginImageContextWithOptions(newSize, true, uiImage.scale)
uiImage.draw(in: CGRect(x: 0, y: 0, width: newSize.width, height: newSize.height))
let resizedImage = UIGraphicsGetImageFromCurrentImageContext()!
UIGraphicsEndImageContext()
return resizedImage
}風格轉換的函式
// 風格轉換
func styleTransfer(image: UIImage, model: VNCoreMLModel) -> UIImage? {
// 將圖片轉成模型輸入的尺寸
UIGraphicsBeginImageContextWithOptions(CGSize(width: 640, height: 960), true, 2.0)
image.draw(in: CGRect(x: 0, y: 0, width: 640, height: 960))
let newImage = UIGraphicsGetImageFromCurrentImageContext()!
UIGraphicsEndImageContext()
// 將處理後的圖片轉換為像素緩衝區，以供模型輸入使用
guard let pixelBuffer = newImage.toPixelBuffer(pixelFormatType: kCVPixelFormatType_32ARGB, width: 640, height: 960) else {
return nil
}
let request = VNCoreMLRequest(model: model)
let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer, options: [:])
do {
try handler.perform([request])
guard let result = request.results?.first as? VNPixelBufferObservation else {
return nil
}
let styledImage = UIImage(ciImage: CIImage(cvPixelBuffer: result.pixelBuffer))
// 將經過風格轉換後的照片尺寸，變回原本相片尺寸
return resizeImage(uiImage: styledImage, newSize: image.size)
} catch {
print("Error occurred during style transfer: \(error)")
return nil
}
}風格轉換的函式的步驟與之前做影像辨識的步驟都差不多，只差在最後處理的一些細節有點不同罷了。此外，在最後需要將圖片變回原本尺寸也很好理解，因為 pre-trained model 的輸入尺寸為 640x960。為了避免照片輸出都是直立的，必須透過 resizeImage 函式將風格轉換後的圖片變回原本的尺寸。
分別對下面三個 UIButton 拉 @IBAction 之後，添加下列的程式碼
// 風格轉換
@IBAction func cupheadTransfer(_ sender: UIButton) {
// 選擇使用的模型
initializeModel(.cuphead)
// 使用 originalImage 確保圖片不被重複覆蓋特效
guard let image = originalImage else {return}
// 進行風格轉換
if let styledImage = styleTransfer(image: image, model: model) {
imageView.image = styledImage
}
}做完上述的幾個步驟後，你就能進行風格轉換了
保存照片
接下來我會透過 IOS 中 SDK 的 UIActivityViewController 實現相片保存，很常看到大家會透過三段程式碼就將照片進行保存：
@IBAction func shareImage(_ sender: UIBarButtonItem) {
// 如果 imageView 上沒有圖像，則無法分享或儲存
guard let imageToShare = imageView.image else {return}
let activityViewController = UIActivityViewController(activityItems: [imageToShare], applicationActivities: nil)
present(activityViewController, animated: true, completion: nil)
}
僅靠這三段程式碼會出現左上方的情況，縮圖的照片會是空白的，且會沒有相片的資訊，右上方為從檔案保存相片時出現的狀況。
正確的作法是：先將照片檔案儲存到本機檔案系統，並將檔案的 URL 用作你的 activityItem。
// 將照片進行保存
@IBAction func shareImage(_ sender: UIBarButtonItem) {
// 如果 imageView 上沒有圖像，則無法分享或儲存
guard let imageToShare = imageView.image else {return}
// 產生一個唯一的檔名，以便保存圖像
let timestamp = Int(Date().timeIntervalSince1970)
let fileName = "shared_image_\(timestamp).png"
// 取得檔案的本機 URL 路徑
guard let fileURL = saveImageToDocumentsDirectory(image: imageToShare, fileName: fileName) else {return}
// 建立一個帶有圖像和文件 URL 的活動項目
let itemsToShare: [Any] = [fileURL]
let activityViewController = UIActivityViewController(activityItems: itemsToShare, applicationActivities: nil)
present(activityViewController, animated: true, completion: nil)
}
// 儲存圖片到應用程式的文件目錄並傳回檔案的本機 URL
func saveImageToDocumentsDirectory(image: UIImage, fileName: String) -> URL? {
guard let documentsDirectory = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first else {
return nil
}
// 建置檔案的本機 URL，將檔案名稱新增至文件目錄路徑中
let fileURL = documentsDirectory.appendingPathComponent(fileName)
// 將圖片轉換為 PNG 數據
if let imageData = image.pngData() {
do {
try imageData.write(to: fileURL)
return fileURL
} catch {
print("Error occurred during saving the image：\(error.localizedDescription)")
return nil
}
}
return nil
}
在完成上述步驟後，就能成功解決了儲存照片時左上角縮圖和照片資訊消失的問題，並且可以順利將經過風格轉換後的照片儲存到相簿中。
