diff --git a/AOT-GAN-for-Inpainting-CoreML.xcodeproj/project.pbxproj b/AOT-GAN-for-Inpainting-CoreML.xcodeproj/project.pbxproj index 073d9d7..c425bab 100644 --- a/AOT-GAN-for-Inpainting-CoreML.xcodeproj/project.pbxproj +++ b/AOT-GAN-for-Inpainting-CoreML.xcodeproj/project.pbxproj @@ -16,6 +16,11 @@ 734FA4F9298C59DF007627D9 /* AOT_GAN_for_Inpainting_CoreMLTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 734FA4F8298C59DF007627D9 /* AOT_GAN_for_Inpainting_CoreMLTests.swift */; }; 734FA503298C59DF007627D9 /* AOT_GAN_for_Inpainting_CoreMLUITests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 734FA502298C59DF007627D9 /* AOT_GAN_for_Inpainting_CoreMLUITests.swift */; }; 734FA505298C59DF007627D9 /* AOT_GAN_for_Inpainting_CoreMLUITestsLaunchTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 734FA504298C59DF007627D9 /* AOT_GAN_for_Inpainting_CoreMLUITestsLaunchTests.swift */; }; + 734FA512298C5A69007627D9 /* aotgan.mlmodel in Sources */ = {isa = PBXBuildFile; fileRef = 734FA511298C5A69007627D9 /* aotgan.mlmodel */; }; + 734FA514298C5BCF007627D9 /* CoreMLHelpers.swift in Sources */ = {isa = PBXBuildFile; fileRef = 734FA513298C5BCE007627D9 /* CoreMLHelpers.swift */; }; + 734FA55A298E0D13007627D9 /* Drawsana in Frameworks */ = {isa = PBXBuildFile; productRef = 734FA559298E0D13007627D9 /* Drawsana */; }; + 734FA59C29920691007627D9 /* Utils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 734FA59B29920691007627D9 /* Utils.swift */; }; + 734FA59E29920E08007627D9 /* realesrgangeneral512.mlmodel in Sources */ = {isa = PBXBuildFile; fileRef = 734FA59D29920E08007627D9 /* realesrgangeneral512.mlmodel */; }; /* End PBXBuildFile section */ /* Begin PBXContainerItemProxy section */ @@ -49,6 +54,10 @@ 734FA4FE298C59DF007627D9 /* AOT-GAN-for-Inpainting-CoreMLUITests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = "AOT-GAN-for-Inpainting-CoreMLUITests.xctest"; sourceTree = BUILT_PRODUCTS_DIR; }; 734FA502298C59DF007627D9 /* AOT_GAN_for_Inpainting_CoreMLUITests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AOT_GAN_for_Inpainting_CoreMLUITests.swift; sourceTree = ""; }; 734FA504298C59DF007627D9 /* AOT_GAN_for_Inpainting_CoreMLUITestsLaunchTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AOT_GAN_for_Inpainting_CoreMLUITestsLaunchTests.swift; sourceTree = ""; }; + 734FA511298C5A69007627D9 /* aotgan.mlmodel */ = {isa = PBXFileReference; lastKnownFileType = file.mlmodel; path = aotgan.mlmodel; sourceTree = ""; }; + 734FA513298C5BCE007627D9 /* CoreMLHelpers.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = CoreMLHelpers.swift; sourceTree = ""; }; + 734FA59B29920691007627D9 /* Utils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Utils.swift; sourceTree = ""; }; + 734FA59D29920E08007627D9 /* realesrgangeneral512.mlmodel */ = {isa = PBXFileReference; lastKnownFileType = file.mlmodel; path = realesrgangeneral512.mlmodel; sourceTree = ""; }; /* End PBXFileReference section */ /* Begin PBXFrameworksBuildPhase section */ @@ -56,6 +65,7 @@ isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( + 734FA55A298E0D13007627D9 /* Drawsana in Frameworks */, ); runOnlyForDeploymentPostprocessing = 0; }; @@ -102,7 +112,11 @@ 734FA4E1298C59DD007627D9 /* AppDelegate.swift */, 734FA4E3298C59DD007627D9 /* SceneDelegate.swift */, 734FA4E5298C59DD007627D9 /* ViewController.swift */, + 734FA59B29920691007627D9 /* Utils.swift */, + 734FA513298C5BCE007627D9 /* CoreMLHelpers.swift */, 734FA4E7298C59DD007627D9 /* Main.storyboard */, + 734FA511298C5A69007627D9 /* aotgan.mlmodel */, + 734FA59D29920E08007627D9 /* realesrgangeneral512.mlmodel */, 734FA4EA298C59DF007627D9 /* Assets.xcassets */, 734FA4EC298C59DF007627D9 /* LaunchScreen.storyboard */, 734FA4EF298C59DF007627D9 /* Info.plist */, @@ -143,6 +157,9 @@ dependencies = ( ); name = "AOT-GAN-for-Inpainting-CoreML"; + packageProductDependencies = ( + 734FA559298E0D13007627D9 /* Drawsana */, + ); productName = "AOT-GAN-for-Inpainting-CoreML"; productReference = 734FA4DE298C59DD007627D9 /* AOT-GAN-for-Inpainting-CoreML.app */; productType = "com.apple.product-type.application"; @@ -215,6 +232,9 @@ Base, ); mainGroup = 734FA4D5298C59DD007627D9; + packageReferences = ( + 734FA558298E0D13007627D9 /* XCRemoteSwiftPackageReference "Drawsana" */, + ); productRefGroup = 734FA4DF298C59DD007627D9 /* Products */; projectDirPath = ""; projectRoot = ""; @@ -260,7 +280,11 @@ files = ( 734FA4E6298C59DD007627D9 /* ViewController.swift in Sources */, 734FA4E2298C59DD007627D9 /* AppDelegate.swift in Sources */, + 734FA514298C5BCF007627D9 /* CoreMLHelpers.swift in Sources */, + 734FA59E29920E08007627D9 /* realesrgangeneral512.mlmodel in Sources */, 734FA4E4298C59DD007627D9 /* SceneDelegate.swift in Sources */, + 734FA512298C5A69007627D9 /* aotgan.mlmodel in Sources */, + 734FA59C29920691007627D9 /* Utils.swift in Sources */, ); runOnlyForDeploymentPostprocessing = 0; }; @@ -436,10 +460,12 @@ ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; CODE_SIGN_STYLE = Automatic; - CURRENT_PROJECT_VERSION = 1; + CURRENT_PROJECT_VERSION = 1.2; DEVELOPMENT_TEAM = MFN25KNUGJ; GENERATE_INFOPLIST_FILE = YES; INFOPLIST_FILE = "AOT-GAN-for-Inpainting-CoreML/Info.plist"; + INFOPLIST_KEY_CFBundleDisplayName = Inpainting; + INFOPLIST_KEY_NSPhotoLibraryUsageDescription = "$(MARKETING_VERSION)"; INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES; INFOPLIST_KEY_UILaunchStoryboardName = LaunchScreen; INFOPLIST_KEY_UIMainStoryboardFile = Main; @@ -464,10 +490,12 @@ ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; CODE_SIGN_STYLE = Automatic; - CURRENT_PROJECT_VERSION = 1; + CURRENT_PROJECT_VERSION = 1.2; DEVELOPMENT_TEAM = MFN25KNUGJ; GENERATE_INFOPLIST_FILE = YES; INFOPLIST_FILE = "AOT-GAN-for-Inpainting-CoreML/Info.plist"; + INFOPLIST_KEY_CFBundleDisplayName = Inpainting; + INFOPLIST_KEY_NSPhotoLibraryUsageDescription = "$(MARKETING_VERSION)"; INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES; INFOPLIST_KEY_UILaunchStoryboardName = LaunchScreen; INFOPLIST_KEY_UIMainStoryboardFile = Main; @@ -602,6 +630,25 @@ defaultConfigurationName = Release; }; /* End XCConfigurationList section */ + +/* Begin XCRemoteSwiftPackageReference section */ + 734FA558298E0D13007627D9 /* XCRemoteSwiftPackageReference "Drawsana" */ = { + isa = XCRemoteSwiftPackageReference; + repositoryURL = "https://github.com/Asana/Drawsana"; + requirement = { + kind = upToNextMajorVersion; + minimumVersion = 0.9.4; + }; + }; +/* End XCRemoteSwiftPackageReference section */ + +/* Begin XCSwiftPackageProductDependency section */ + 734FA559298E0D13007627D9 /* Drawsana */ = { + isa = XCSwiftPackageProductDependency; + package = 734FA558298E0D13007627D9 /* XCRemoteSwiftPackageReference "Drawsana" */; + productName = Drawsana; + }; +/* End XCSwiftPackageProductDependency section */ }; rootObject = 734FA4D6298C59DD007627D9 /* Project object */; } diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/AppIcon.appiconset/Contents.json b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/AppIcon.appiconset/Contents.json index 13613e3..c926a5d 100644 --- a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/AppIcon.appiconset/Contents.json +++ b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/AppIcon.appiconset/Contents.json @@ -1,6 +1,7 @@ { "images" : [ { + "filename" : "inpintingIcon.png", "idiom" : "universal", "platform" : "ios", "size" : "1024x1024" diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/AppIcon.appiconset/inpintingIcon.png b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/AppIcon.appiconset/inpintingIcon.png new file mode 100644 index 0000000..52c0d27 Binary files /dev/null and b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/AppIcon.appiconset/inpintingIcon.png differ diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/alcove_00003905.imageset/Contents.json b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/alcove_00003905.imageset/Contents.json new file mode 100644 index 0000000..c83397f --- /dev/null +++ b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/alcove_00003905.imageset/Contents.json @@ -0,0 +1,21 @@ +{ + "images" : [ + { + "filename" : "alcove_00003905.png", + "idiom" : "universal", + "scale" : "1x" + }, + { + "idiom" : "universal", + "scale" : "2x" + }, + { + "idiom" : "universal", + "scale" : "3x" + } + ], + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/alcove_00003905.imageset/alcove_00003905.png b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/alcove_00003905.imageset/alcove_00003905.png new file mode 100644 index 0000000..cb9d6fe Binary files /dev/null and b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/alcove_00003905.imageset/alcove_00003905.png differ diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/image_masked.imageset/Contents.json b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/image_masked.imageset/Contents.json new file mode 100644 index 0000000..15dfb09 --- /dev/null +++ b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/image_masked.imageset/Contents.json @@ -0,0 +1,21 @@ +{ + "images" : [ + { + "filename" : "image_masked.png", + "idiom" : "universal", + "scale" : "1x" + }, + { + "idiom" : "universal", + "scale" : "2x" + }, + { + "idiom" : "universal", + "scale" : "3x" + } + ], + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/image_masked.imageset/image_masked.png b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/image_masked.imageset/image_masked.png new file mode 100644 index 0000000..44a067a Binary files /dev/null and b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/image_masked.imageset/image_masked.png differ diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/input.imageset/Contents.json b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/input.imageset/Contents.json new file mode 100644 index 0000000..c83397f --- /dev/null +++ b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/input.imageset/Contents.json @@ -0,0 +1,21 @@ +{ + "images" : [ + { + "filename" : "alcove_00003905.png", + "idiom" : "universal", + "scale" : "1x" + }, + { + "idiom" : "universal", + "scale" : "2x" + }, + { + "idiom" : "universal", + "scale" : "3x" + } + ], + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/input.imageset/alcove_00003905.png b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/input.imageset/alcove_00003905.png new file mode 100644 index 0000000..cd4f002 Binary files /dev/null and b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/input.imageset/alcove_00003905.png differ diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/mask.imageset/Contents.json b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/mask.imageset/Contents.json new file mode 100644 index 0000000..f86de16 --- /dev/null +++ b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/mask.imageset/Contents.json @@ -0,0 +1,21 @@ +{ + "images" : [ + { + "filename" : "mask_.png", + "idiom" : "universal", + "scale" : "1x" + }, + { + "idiom" : "universal", + "scale" : "2x" + }, + { + "idiom" : "universal", + "scale" : "3x" + } + ], + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/mask.imageset/mask_.png b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/mask.imageset/mask_.png new file mode 100644 index 0000000..1199c08 Binary files /dev/null and b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/mask.imageset/mask_.png differ diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/maskBackGround.imageset/Contents.json b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/maskBackGround.imageset/Contents.json new file mode 100644 index 0000000..f670ed3 --- /dev/null +++ b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/maskBackGround.imageset/Contents.json @@ -0,0 +1,21 @@ +{ + "images" : [ + { + "filename" : "maskBackGround.jpg", + "idiom" : "universal", + "scale" : "1x" + }, + { + "idiom" : "universal", + "scale" : "2x" + }, + { + "idiom" : "universal", + "scale" : "3x" + } + ], + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/maskBackGround.imageset/maskBackGround.jpg b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/maskBackGround.imageset/maskBackGround.jpg new file mode 100644 index 0000000..ead1c70 Binary files /dev/null and b/AOT-GAN-for-Inpainting-CoreML/Assets.xcassets/maskBackGround.imageset/maskBackGround.jpg differ diff --git a/AOT-GAN-for-Inpainting-CoreML/CoreMLHelpers.swift b/AOT-GAN-for-Inpainting-CoreML/CoreMLHelpers.swift new file mode 100644 index 0000000..195efb3 --- /dev/null +++ b/AOT-GAN-for-Inpainting-CoreML/CoreMLHelpers.swift @@ -0,0 +1,546 @@ +// +// CoreMLHelpers.swift +// Inpating +// +// Created by 間嶋大輔 on 2023/01/12. +// + +import Accelerate +import CoreML + +public protocol MultiArrayType: Comparable { + static var multiArrayDataType: MLMultiArrayDataType { get } + static func +(lhs: Self, rhs: Self) -> Self + static func -(lhs: Self, rhs: Self) -> Self + static func *(lhs: Self, rhs: Self) -> Self + static func /(lhs: Self, rhs: Self) -> Self + init(_: Int) + var toUInt8: UInt8 { get } +} + +extension Double: MultiArrayType { + public static var multiArrayDataType: MLMultiArrayDataType { return .double } + public var toUInt8: UInt8 { return UInt8(self) } +} + +extension Float: MultiArrayType { + public static var multiArrayDataType: MLMultiArrayDataType { return .float32 } + public var toUInt8: UInt8 { return UInt8(self) } +} + +extension Int32: MultiArrayType { + public static var multiArrayDataType: MLMultiArrayDataType { return .int32 } + public var toUInt8: UInt8 { return UInt8(self) } +} + +extension MLMultiArray { + /** + Converts the multi-array to a CGImage. + The multi-array must have at least 2 dimensions for a grayscale image, or + at least 3 dimensions for a color image. + The default expected shape is (height, width) or (channels, height, width). + However, you can change this using the `axes` parameter. For example, if + the array shape is (1, height, width, channels), use `axes: (3, 1, 2)`. + If `channel` is not nil, only converts that channel to a grayscale image. + This lets you visualize individual channels from a multi-array with more + than 4 channels. + Otherwise, converts all channels. In this case, the number of channels in + the multi-array must be 1 for grayscale, 3 for RGB, or 4 for RGBA. + Use the `min` and `max` parameters to put the values from the array into + the range [0, 255], if not already: + - `min`: should be the smallest value in the data; this will be mapped to 0. + - `max`: should be the largest value in the data; will be mapped to 255. + For example, if the range of the data in the multi-array is [-1, 1], use + `min: -1, max: 1`. If the range is already [0, 255], then use the defaults. + */ + public func cgImage(min: Double = 0, + max: Double = 255, + channel: Int? = nil, + axes: (Int, Int, Int)? = nil) -> CGImage? { + switch self.dataType { + case .double: + return _image(min: min, max: max, channel: channel, axes: axes) + case .float32: + return _image(min: Float(min), max: Float(max), channel: channel, axes: axes) + case .int32: + return _image(min: Int32(min), max: Int32(max), channel: channel, axes: axes) + @unknown default: + fatalError("Unsupported data type \(dataType.rawValue)") + } + } + + /** + Helper function that allows us to use generics. The type of `min` and `max` + is also the dataType of the MLMultiArray. + */ + private func _image(min: T, + max: T, + channel: Int?, + axes: (Int, Int, Int)?) -> CGImage? { + if let (b, w, h, c) = toRawBytes(min: min, max: max, channel: channel, axes: axes) { + if c == 1 { + return CGImage.fromByteArrayGray(b, width: w, height: h) + } else { + return CGImage.fromByteArrayRGBA(b, width: w, height: h) + } + } + return nil + } + + /** + Converts the multi-array into an array of RGBA or grayscale pixels. + - Note: This is not particularly fast, but it is flexible. You can change + the loops to convert the multi-array whichever way you please. + - Note: The type of `min` and `max` must match the dataType of the + MLMultiArray object. + - Returns: tuple containing the RGBA bytes, the dimensions of the image, + and the number of channels in the image (1, 3, or 4). + */ + public func toRawBytes(min: T, + max: T, + channel: Int? = nil, + axes: (Int, Int, Int)? = nil) + -> (bytes: [UInt8], width: Int, height: Int, channels: Int)? { + // MLMultiArray with unsupported shape? + if shape.count < 2 { + print("Cannot convert MLMultiArray of shape \(shape) to image") + return nil + } + + // Figure out which dimensions to use for the channels, height, and width. + let channelAxis: Int + let heightAxis: Int + let widthAxis: Int + if let axes = axes { + channelAxis = axes.0 + heightAxis = axes.1 + widthAxis = axes.2 + guard channelAxis >= 0 && channelAxis < shape.count && + heightAxis >= 0 && heightAxis < shape.count && + widthAxis >= 0 && widthAxis < shape.count else { + print("Invalid axes \(axes) for shape \(shape)") + return nil + } + } else if shape.count == 2 { + // Expected shape for grayscale is (height, width) + heightAxis = 0 + widthAxis = 1 + channelAxis = -1 // Never be used + } else { + // Expected shape for color is (channels, height, width) + channelAxis = 0 + heightAxis = 1 + widthAxis = 2 + } + + let height = self.shape[heightAxis].intValue + let width = self.shape[widthAxis].intValue + let yStride = self.strides[heightAxis].intValue + let xStride = self.strides[widthAxis].intValue + + let channels: Int + let cStride: Int + let bytesPerPixel: Int + let channelOffset: Int + + // MLMultiArray with just two dimensions is always grayscale. (We ignore + // the value of channelAxis here.) + if shape.count == 2 { + channels = 1 + cStride = 0 + bytesPerPixel = 1 + channelOffset = 0 + + // MLMultiArray with more than two dimensions can be color or grayscale. + } else { + let channelDim = self.shape[channelAxis].intValue + if let channel = channel { + if channel < 0 || channel >= channelDim { + print("Channel must be -1, or between 0 and \(channelDim - 1)") + return nil + } + channels = 1 + bytesPerPixel = 1 + channelOffset = channel + } else if channelDim == 1 { + channels = 1 + bytesPerPixel = 1 + channelOffset = 0 + } else { + if channelDim != 3 && channelDim != 4 { + print("Expected channel dimension to have 1, 3, or 4 channels, got \(channelDim)") + return nil + } + channels = channelDim + bytesPerPixel = 4 + channelOffset = 0 + } + cStride = self.strides[channelAxis].intValue + } + + // Allocate storage for the RGBA or grayscale pixels. Set everything to + // 255 so that alpha channel is filled in if only 3 channels. + let count = height * width * bytesPerPixel + var pixels = [UInt8](repeating: 255, count: count) + + // Grab the pointer to MLMultiArray's memory. + var ptr = UnsafeMutablePointer(OpaquePointer(self.dataPointer)) + ptr = ptr.advanced(by: channelOffset * cStride) + + // Loop through all the pixels and all the channels and copy them over. + for c in 0.. CGImage? { + assert(features.dataType == .float32) + assert(features.shape.count == 3) + + let ptr = UnsafeMutablePointer(OpaquePointer(features.dataPointer)) + + let height = features.shape[1].intValue + let width = features.shape[2].intValue + let channelStride = features.strides[0].intValue + let rowStride = features.strides[1].intValue + let srcRowBytes = rowStride * MemoryLayout.stride + + var blueBuffer = vImage_Buffer(data: ptr, + height: vImagePixelCount(height), + width: vImagePixelCount(width), + rowBytes: srcRowBytes) + var greenBuffer = vImage_Buffer(data: ptr.advanced(by: channelStride), + height: vImagePixelCount(height), + width: vImagePixelCount(width), + rowBytes: srcRowBytes) + var redBuffer = vImage_Buffer(data: ptr.advanced(by: channelStride * 2), + height: vImagePixelCount(height), + width: vImagePixelCount(width), + rowBytes: srcRowBytes) + + let destRowBytes = width * 4 + var pixels = [UInt8](repeating: 0, count: height * destRowBytes) + var destBuffer = vImage_Buffer(data: &pixels, + height: vImagePixelCount(height), + width: vImagePixelCount(width), + rowBytes: destRowBytes) + + let error = vImageConvert_PlanarFToBGRX8888(&blueBuffer, + &greenBuffer, + &redBuffer, + Pixel_8(255), + &destBuffer, + [max, max, max], + [min, min, min], + vImage_Flags(0)) + if error == kvImageNoError { + return CGImage.fromByteArrayRGBA(pixels, width: width, height: height) + } else { + return nil + } +} + +#if canImport(UIKit) + +import UIKit + +extension MLMultiArray { + public func image(min: Double = 0, + max: Double = 255, + channel: Int? = nil, + axes: (Int, Int, Int)? = nil) -> UIImage? { + let cgImg = cgImage(min: min, max: max, channel: channel, axes: axes) + return cgImg.map { UIImage(cgImage: $0) } + } +} + +public func createUIImage(fromFloatArray features: MLMultiArray, + min: Float = 0, + max: Float = 255) -> UIImage? { + let cgImg = createCGImage(fromFloatArray: features, min: min, max: max) + return cgImg.map { UIImage(cgImage: $0) } +} + +#endif + +public func clamp(_ x: T, min: T, max: T) -> T { + if x < min { return min } + if x > max { return max } + return x +} + +import CoreGraphics + +extension CGImage { + /** + Converts the image into an array of RGBA bytes. + */ + @nonobjc public func toByteArrayRGBA() -> [UInt8] { + var bytes = [UInt8](repeating: 0, count: width * height * 4) + bytes.withUnsafeMutableBytes { ptr in + if let colorSpace = colorSpace, + let context = CGContext( + data: ptr.baseAddress, + width: width, + height: height, + bitsPerComponent: bitsPerComponent, + bytesPerRow: bytesPerRow, + space: colorSpace, + bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) { + let rect = CGRect(x: 0, y: 0, width: width, height: height) + context.draw(self, in: rect) + } + } + return bytes + } + + /** + Creates a new CGImage from an array of RGBA bytes. + */ + @nonobjc public class func fromByteArrayRGBA(_ bytes: [UInt8], + width: Int, + height: Int) -> CGImage? { + return fromByteArray(bytes, width: width, height: height, + bytesPerRow: width * 4, + colorSpace: CGColorSpaceCreateDeviceRGB(), + alphaInfo: .premultipliedLast) + } + + /** + Creates a new CGImage from an array of grayscale bytes. + */ + @nonobjc public class func fromByteArrayGray(_ bytes: [UInt8], + width: Int, + height: Int) -> CGImage? { + return fromByteArray(bytes, width: width, height: height, + bytesPerRow: width, + colorSpace: CGColorSpaceCreateDeviceGray(), + alphaInfo: .none) + } + + @nonobjc class func fromByteArray(_ bytes: [UInt8], + width: Int, + height: Int, + bytesPerRow: Int, + colorSpace: CGColorSpace, + alphaInfo: CGImageAlphaInfo) -> CGImage? { + return bytes.withUnsafeBytes { ptr in + let context = CGContext(data: UnsafeMutableRawPointer(mutating: ptr.baseAddress!), + width: width, + height: height, + bitsPerComponent: 8, + bytesPerRow: bytesPerRow, + space: colorSpace, + bitmapInfo: alphaInfo.rawValue) + return context?.makeImage() + } + } +} + + +extension CGImage { + var frame: CGRect { + return CGRect(x: 0, y: 0, width: self.width, height: self.height) + } + + func toBGR()->CGImage{ + let ciImage = CIImage(cgImage: self) + let kernelStr: String = """ + kernel vec4 swapRedAndGreenAmount(__sample s) { + return s.bgra; + } + """ + let ctx = CIContext(options: nil) + let swapKernel = CIColorKernel( source: + "kernel vec4 swapRedAndGreenAmount(__sample s) {" + + "return s.bgra;" + + "}" + ) + let ciOutput = swapKernel?.apply(extent: (ciImage.extent), arguments: [ciImage as Any]) + let cgOut:CGImage = ctx.createCGImage(ciOutput!, from: ciOutput!.extent)! + return cgOut + + } +} + + +extension UIImage { + func mlMultiArray(scale preprocessScale:Double=1/255, rBias preprocessRBias:Double=0, gBias preprocessGBias:Double=0, bBias preprocessBBias:Double=0) -> MLMultiArray { + let imagePixel = self.getPixelRgb(scale: preprocessScale, rBias: preprocessRBias, gBias: preprocessGBias, bBias: preprocessBBias) +// let size = self.size + let imagePointer : UnsafePointer = UnsafePointer(imagePixel) + let mlArray = try! MLMultiArray(shape: [1,3, NSNumber(value: Float(512)), NSNumber(value: Float(512))], dataType: MLMultiArrayDataType.double) + mlArray.dataPointer.initializeMemory(as: Double.self, from: imagePointer, count: imagePixel.count) + + return mlArray + } + + func mlMultiArrayGrayScale(scale preprocessScale:Double=1/255,bias preprocessBias:Double=0) -> MLMultiArray { + let imagePixel = self.getPixelGrayScale(scale: preprocessScale, bias: preprocessBias) +// let size = self.size + let imagePointer : UnsafePointer = UnsafePointer(imagePixel) + let mlArray = try! MLMultiArray(shape: [1,1, NSNumber(value: Float(512)), NSNumber(value: Float(512))], dataType: MLMultiArrayDataType.double) + mlArray.dataPointer.initializeMemory(as: Double.self, from: imagePointer, count: imagePixel.count) + return mlArray + } + + func mlMultiArrayComposite(outImage out:UIImage, inputImage input:UIImage, maskImage mask: UIImage, scale preprocessScale:Double=1/255, rBias preprocessRBias:Double=0, gBias preprocessGBias:Double=0, bBias preprocessBBias:Double=0) -> MLMultiArray { + let imagePixel = self.getMaskedPixelRgb(out: out, input: input, mask: mask) +// let size = self.size + let imagePointer : UnsafePointer = UnsafePointer(imagePixel) + let mlArray = try! MLMultiArray(shape: [1,3, NSNumber(value: Float(512)), NSNumber(value: Float(512))], dataType: MLMultiArrayDataType.double) + mlArray.dataPointer.initializeMemory(as: Double.self, from: imagePointer, count: imagePixel.count) + + return mlArray + } + + func getMaskedPixelRgb(out: UIImage,input: UIImage, mask:UIImage, scale preprocessScale:Double=1, rBias preprocessRBias:Double=0, gBias preprocessGBias:Double=0, bBias preprocessBBias:Double=0) -> [Double] + { + guard let outCGImage = out.cgImage?.resize(size: CGSize(width: 512, height: 512)) else { + return [] + } + let outbytesPerRow = outCGImage.bytesPerRow + let outwidth = outCGImage.width + let outheight = outCGImage.height + let outbytesPerPixel = 4 + let outpixelData = outCGImage.dataProvider!.data! as Data + + guard let inputCGImage = input.cgImage?.resize(size: CGSize(width: 512, height: 512)) else { + return [] + } + let inputpixelData = inputCGImage.dataProvider!.data! as Data + + guard let maskCgImage = mask.cgImage?.resize(size: CGSize(width: 512, height: 512)) else { + return [] + } + let maskBytesPerRow = maskCgImage.bytesPerRow + let maskBytesPerPixel = 4 + let maskPixelData = maskCgImage.dataProvider!.data! as Data + + var r_buf : [Double] = [] + var g_buf : [Double] = [] + var b_buf : [Double] = [] + + for j in 0.. 0 { + r_buf.append(Double(r*preprocessScale)+preprocessRBias) + g_buf.append(Double(g*preprocessScale)+preprocessGBias) + b_buf.append(Double(b*preprocessScale)+preprocessBBias) + } else { + r_buf.append(Double(bgr*preprocessScale)+preprocessRBias) + g_buf.append(Double(bgg*preprocessScale)+preprocessGBias) + b_buf.append(Double(bgb*preprocessScale)+preprocessBBias) + + } + } + } + + return ((r_buf + g_buf) + b_buf) + } + + func getPixelRgb(scale preprocessScale:Double=1/255, rBias preprocessRBias:Double=0, gBias preprocessGBias:Double=0, bBias preprocessBBias:Double=0) -> [Double] + { + guard let cgImage = self.cgImage?.resize(size: CGSize(width: 512, height: 512)) else { + return [] + } + let bytesPerRow = cgImage.bytesPerRow + let width = cgImage.width + let height = cgImage.height + let bytesPerPixel = 4 + let pixelData = cgImage.dataProvider!.data! as Data + + var r_buf : [Double] = [] + var g_buf : [Double] = [] + var b_buf : [Double] = [] + + for j in 0.. [Double] + { + guard let cgImage = self.cgImage?.resize(size: CGSize(width: 512, height: 512)) else { + return [] + } + let bytesPerRow = cgImage.bytesPerRow + let width = cgImage.width + let height = cgImage.height + let bytesPerPixel = 4 + let pixelData = cgImage.dataProvider!.data! as Data + + var buf : [Double] = [] + + for j in 0.. CGImage? { + let width: Int = Int(size.width) + let height: Int = Int(size.height) + + let bytesPerPixel = self.bitsPerPixel / self.bitsPerComponent + let destBytesPerRow = width * bytesPerPixel + + + guard let colorSpace = self.colorSpace else { return nil } + guard let context = CGContext(data: nil, width: width, height: height, bitsPerComponent: self.bitsPerComponent, bytesPerRow: destBytesPerRow, space: colorSpace, bitmapInfo: self.alphaInfo.rawValue) else { return nil } + + context.interpolationQuality = .high + context.draw(self, in: CGRect(x: 0, y: 0, width: width, height: height)) + + return context.makeImage() + } +} + diff --git a/AOT-GAN-for-Inpainting-CoreML/Utils.swift b/AOT-GAN-for-Inpainting-CoreML/Utils.swift new file mode 100644 index 0000000..01ed753 --- /dev/null +++ b/AOT-GAN-for-Inpainting-CoreML/Utils.swift @@ -0,0 +1,8 @@ +// +// Utils.swift +// AOT-GAN-for-Inpainting-CoreML +// +// Created by 間嶋大輔 on 2023/02/07. +// + +import Foundation diff --git a/AOT-GAN-for-Inpainting-CoreML/ViewController.swift b/AOT-GAN-for-Inpainting-CoreML/ViewController.swift index 7fd5c31..9f01f0d 100644 --- a/AOT-GAN-for-Inpainting-CoreML/ViewController.swift +++ b/AOT-GAN-for-Inpainting-CoreML/ViewController.swift @@ -1,19 +1,269 @@ -// -// ViewController.swift -// AOT-GAN-for-Inpainting-CoreML -// -// Created by 間嶋大輔 on 2023/02/03. -// - import UIKit +import Vision +import CoreML +import Drawsana +import PhotosUI -class ViewController: UIViewController { - +class ViewController: UIViewController,PHPickerViewControllerDelegate, UIPickerViewDelegate { + + var imageView:UIImageView! + let drawingView = DrawsanaView() + let penTool = PenTool() + let undoButton = UIButton() + let runButton = UIButton() + let selectPhotoButton = UIButton() + let superResolutionButton = UIButton() + let brushButton = UIButton() + var inputImage: UIImage? + let ciContext = CIContext() + lazy var maskBackGroundImage: UIImage = { + guard let image = UIImage(named: "maskBackGround") else { fatalError("Please set black image ") } + return image + }() + + lazy var model: MLModel? = { + do { + let config = MLModelConfiguration() + config.computeUnits = .cpuAndGPU + let model = try aotgan(configuration: config).model + return model + } catch let error { + print(error) + fatalError("model initialize error") + } + }() + + lazy var srRequest: VNCoreMLRequest = { + do { + let config = MLModelConfiguration() + config.computeUnits = .cpuAndGPU + let model = try realesrgangeneral512(configuration: config).model + let vnModel = try VNCoreMLModel(for: model) + let request = VNCoreMLRequest(model: vnModel) + request.imageCropAndScaleOption = .scaleFill + return request + } catch let error { + print(error) + fatalError("model initialize error") + } + }() + override func viewDidLoad() { super.viewDidLoad() - // Do any additional setup after loading the view. + setupView() } + + + func inference(maskedImage inputImage:UIImage, maskImage mask:UIImage) { + guard let model = model else { fatalError("Model initialize error.") } + do { + let mask_1 = mask.mlMultiArrayGrayScale(scale: 1/255) + + // input + + let x_1 = inputImage.mlMultiArray(scale: 1/127.5,rBias: -1,gBias: -1, bBias: -1) + + let inputTensor = MLMultiArray(concatenating: [x_1, mask_1], + axis: 1, + dataType: .float32) + + let input = aotganInput(x_1: inputTensor) + + // run + let start = Date() + let out = try model.prediction(from: input) + let timeElapsed = -start.timeIntervalSinceNow + print(timeElapsed) + let outArray = out.featureValue(for: "var_915")?.multiArrayValue + let outImage = outArray!.cgImage(min: -1,max: 1, axes: (1,2,3)) + let uiOut = UIImage(cgImage: outImage!) + let originalSize = self.inputImage?.size + let comp = uiOut.mlMultiArrayComposite(outImage: uiOut, inputImage: inputImage, maskImage: mask).cgImage(axes: (1,2,3))?.resize(size: originalSize!) + let final = UIImage(cgImage: comp!) + + self.inputImage = final + + DispatchQueue.main.async { + self.resetDrawingView() + self.imageView.image = final + } + print("Done") + } catch let error { + print(error) + } + } + + func resetDrawingView() { + if drawingView.drawing.shapes.count != 0 { + for _ in 0...drawingView.drawing.shapes.count-1 { + drawingView.operationStack.undo() + } + } + } + + func setupView() { + imageView = UIImageView(frame: view.bounds) + imageView.contentMode = .scaleAspectFit + view.addSubview(imageView) + inputImage = UIImage(named: "input") + imageView.image = inputImage + + drawingView.set(tool: penTool) + drawingView.userSettings.strokeWidth = 20 + drawingView.userSettings.strokeColor = .white + drawingView.userSettings.fillColor = .black + + drawingView.userSettings.fontSize = 24 + drawingView.userSettings.fontName = "Marker Felt" + drawingView.frame = imageView.frame + view.addSubview(drawingView) + let buttonWidth = view.bounds.width*0.3 + undoButton.frame = CGRect(x: view.bounds.width*0.025, y: 50, width: buttonWidth, height: 50) + superResolutionButton.frame = CGRect(x: view.bounds.maxX-view.bounds.width*0.025-buttonWidth, y: 50, width: buttonWidth, height: 50) + brushButton.frame = CGRect(x: undoButton.frame.maxX + view.bounds.width * 0.025, y: 50, width: buttonWidth, height: 50) + selectPhotoButton.frame = CGRect(x: view.bounds.width*0.1, y: view.bounds.maxY - 100, width: buttonWidth, height: 50) + runButton.frame = CGRect(x: view.bounds.maxX - view.bounds.width*0.1 - buttonWidth, y: view.bounds.maxY - 100, width: buttonWidth, height: 50) + undoButton.setTitle("undoDraw", for: .normal) + selectPhotoButton.setTitle("select Photo", for: .normal) + superResolutionButton.setTitle("SR", for: .normal) + brushButton.setTitle("brushWidth", for: .normal) + runButton.setTitle("run", for: .normal) + undoButton.backgroundColor = .gray + undoButton.setTitleColor(.white, for: .normal) + selectPhotoButton.backgroundColor = .gray + selectPhotoButton.setTitleColor(.white, for: .normal) + superResolutionButton.backgroundColor = .gray + superResolutionButton.setTitleColor(.white, for: .normal) + brushButton.backgroundColor = .gray + brushButton.setTitleColor(.white, for: .normal) + runButton.backgroundColor = .gray + runButton.setTitleColor(.white, for: .normal) + selectPhotoButton.addTarget(self, action: #selector(presentPhPicker), for: .touchUpInside) + superResolutionButton.addTarget(self, action: #selector(sr), for: .touchUpInside) + brushButton.addTarget(self, action: #selector(brushWidth), for: .touchUpInside) + + undoButton.addTarget(self, action: #selector(undo), for: .touchUpInside) + runButton.addTarget(self, action: #selector(run), for: .touchUpInside) + view.addSubview(selectPhotoButton) + view.addSubview(superResolutionButton) + view.addSubview(undoButton) + view.addSubview(brushButton) + + view.addSubview(runButton) + adjustDrawingViewSize() + } + + @objc func brushWidth(){ + let ac = UIAlertController(title: "SelectBrush", message: "", preferredStyle: .alert) + ac.addAction(UIAlertAction(title: "thin", style: .default,handler: { action in + self.drawingView.userSettings.strokeWidth = 5 + })) + ac.addAction(UIAlertAction(title: "medium", style: .default,handler: { action in + self.drawingView.userSettings.strokeWidth = 20 + })) + ac.addAction(UIAlertAction(title: "thin", style: .default,handler: { action in + self.drawingView.userSettings.strokeWidth = 40 + })) + present(ac, animated: true) + } + + @objc func undo() { + drawingView.operationStack.undo() + } + + @objc func run() { + guard let overlapImage:UIImage = drawingView.render(over:inputImage), + let maskImage:UIImage = drawingView.render() + else { fatalError("Mask overlap error") } + inference(maskedImage: overlapImage, maskImage: maskImage) + } + + @objc func presentPhPicker(){ + var configuration = PHPickerConfiguration() + configuration.selectionLimit = 1 + configuration.filter = .images + let picker = PHPickerViewController(configuration: configuration) + picker.delegate = self + present(picker, animated: true) + } + + @objc func sr() { + let handler = VNImageRequestHandler(ciImage: CIImage(image: inputImage!)!) + do { + try handler.perform([srRequest]) + guard let result = srRequest.results?.first as? VNPixelBufferObservation else { + return + } + let srCIImage = CIImage(cvPixelBuffer: result.pixelBuffer) + let resizedCGImage = ciContext.createCGImage(srCIImage, from: srCIImage.extent)?.resize(size: CGSize(width: inputImage!.size.width, height: inputImage!.size.height)) + let srUIImage = UIImage(cgImage: resizedCGImage!) + inputImage = srUIImage + DispatchQueue.main.async { + self.imageView.image = srUIImage + self.adjustDrawingViewSize() + } + } catch let error { + print(error) + } + } + + func picker(_ picker: PHPickerViewController, didFinishPicking results: [PHPickerResult]) { + picker.dismiss(animated: true) + guard let result = results.first else { return } + if result.itemProvider.canLoadObject(ofClass: UIImage.self) { + result.itemProvider.loadObject(ofClass: UIImage.self) { [weak self] image, error in + if let image = image as? UIImage, let safeSelf = self { + let correctOrientImage = safeSelf.getCorrectOrientationUIImage(uiImage: image) + safeSelf.inputImage = correctOrientImage + DispatchQueue.main.async { + safeSelf.resetDrawingView() + safeSelf.imageView.image = correctOrientImage + safeSelf.adjustDrawingViewSize() + } + } + } + } + } + + func getCorrectOrientationUIImage(uiImage:UIImage) -> UIImage { + var newImage = UIImage() + let ciContext = CIContext() + switch uiImage.imageOrientation.rawValue { + case 1: + guard let orientedCIImage = CIImage(image: uiImage)?.oriented(CGImagePropertyOrientation.down), + let cgImage = ciContext.createCGImage(orientedCIImage, from: orientedCIImage.extent) else { return uiImage} + + newImage = UIImage(cgImage: cgImage) + case 3: + guard let orientedCIImage = CIImage(image: uiImage)?.oriented(CGImagePropertyOrientation.right), + let cgImage = ciContext.createCGImage(orientedCIImage, from: orientedCIImage.extent) else { return uiImage} + newImage = UIImage(cgImage: cgImage) + default: + newImage = uiImage + } + return newImage + } + + func adjustDrawingViewSize() { + let displayAspect = imageView.frame.height / imageView.frame.width + let imageSize = imageView.image!.size + let imageAspect = imageSize.height / imageSize.width + if imageAspect <= displayAspect { + // fit to width + let minX = imageView.frame.minX + let minY = imageView.center.y - (imageView.frame.width * imageAspect / 2) + let width = imageView.frame.width + let height = imageView.frame.width * imageAspect + drawingView.frame = CGRect(x: minX, y: minY, width: width, height: height) + + } else { + // fit to height + let aspect = imageSize.width / imageSize.height + drawingView.frame = CGRect(x: imageView.center.x - (imageView.frame.height * aspect / 2), y: imageView.frame.minY, width: imageView.frame.height * aspect, height: imageView.frame.height) + } + + } } diff --git a/AOT-GAN-for-Inpainting-CoreML/aotgan.mlmodel b/AOT-GAN-for-Inpainting-CoreML/aotgan.mlmodel new file mode 100644 index 0000000..0e91677 Binary files /dev/null and b/AOT-GAN-for-Inpainting-CoreML/aotgan.mlmodel differ diff --git a/AOT-GAN-for-Inpainting-CoreML/realesrgangeneral512.mlmodel b/AOT-GAN-for-Inpainting-CoreML/realesrgangeneral512.mlmodel new file mode 100644 index 0000000..a8ce183 Binary files /dev/null and b/AOT-GAN-for-Inpainting-CoreML/realesrgangeneral512.mlmodel differ