Skip to content

Commit 862b77d

Browse files
authored
[RNE Rewrite] feat: add style transfer task implementation (#1278)
## Description Adds style transfer task pipeline implementation and the corresponding example screen in the computer vision app. ### Introduces a breaking change? - [ ] Yes - [x] No ### Type of change - [ ] Bug fix (change which fixes an issue) - [x] New feature (change which adds functionality) - [ ] Documentation update (improves or adds clarity to existing documentation) - [ ] Other (chores, tests, code style improvements etc.) ### Tested on - [x] iOS - [x] Android ### Testing instructions - [ ] Build the computer-vision example app on both Android and iOS - [ ] Test the added style transfer screen. ### Screenshots <!-- Add screenshots here, if applicable --> ### Related issues Closes #1243 ### Checklist - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have updated the documentation accordingly - [x] My changes generate no new warnings ### Additional notes <!-- Include any additional information, assumptions, or context that reviewers might need to understand this PR. -->
1 parent bd2636c commit 862b77d

7 files changed

Lines changed: 504 additions & 0 deletions

File tree

‎apps/computer-vision/app/_layout.tsx‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,13 @@ export default function Layout() {
2727
title: 'Classification',
2828
}}
2929
/>
30+
<Drawer.Screen
31+
name="styleTransfer/index"
32+
options={{
33+
drawerLabel: 'Style Transfer',
34+
title: 'Style Transfer',
35+
}}
36+
/>
3037
<Drawer.Screen
3138
name="segmentation/index"
3239
options={{

‎apps/computer-vision/app/index.tsx‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@ export default function Home() {
1414
<TouchableOpacity style={styles.button} onPress={() => router.navigate('classification/')}>
1515
<Text style={styles.buttonText}>Classification</Text>
1616
</TouchableOpacity>
17+
<TouchableOpacity style={styles.button} onPress={() => router.navigate('styleTransfer/')}>
18+
<Text style={styles.buttonText}>Style Transfer</Text>
19+
</TouchableOpacity>
1720
<TouchableOpacity style={styles.button} onPress={() => router.navigate('segmentation/')}>
1821
<Text style={styles.buttonText}>Semantic Segmentation</Text>
1922
</TouchableOpacity>
Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
1+
import React, { useState } from 'react';
2+
import { View, Text, ScrollView, Platform } from 'react-native';
3+
import { useSafeAreaInsets } from 'react-native-safe-area-context';
4+
import { commonStyles, theme } from '../../theme';
5+
import {
6+
useImage,
7+
Skia,
8+
ColorType,
9+
AlphaType,
10+
type SkImage as SkiaImageType,
11+
} from '@shopify/react-native-skia';
12+
import { useStyleTransfer, models } from 'react-native-executorch';
13+
import ScreenWrapper from '../../components/ScreenWrapper';
14+
import { getImage } from '../../utils';
15+
import { ModelPicker, type ModelOption } from '../../components/ModelPicker';
16+
import { ImageViewport } from '../../components/ImageViewport';
17+
import { ModelStatus } from '../../components/ModelStatus';
18+
import { LatencyIndicator } from '../../components/LatencyIndicator';
19+
import { Button } from '../../components/Button';
20+
21+
const MODEL_OPTIONS: ModelOption[] = [
22+
{
23+
label: 'Candy (XNNPACK INT8)',
24+
value: models.styleTransfer.CANDY.XNNPACK_INT8,
25+
},
26+
{
27+
label: 'Candy (CoreML FP16)',
28+
value: models.styleTransfer.CANDY.COREML_FP16,
29+
disabled: Platform.OS !== 'ios',
30+
},
31+
{
32+
label: 'Mosaic (XNNPACK INT8)',
33+
value: models.styleTransfer.MOSAIC.XNNPACK_INT8,
34+
},
35+
{
36+
label: 'Mosaic (CoreML FP16)',
37+
value: models.styleTransfer.MOSAIC.COREML_FP16,
38+
disabled: Platform.OS !== 'ios',
39+
},
40+
{
41+
label: 'Rain Princess (XNNPACK INT8)',
42+
value: models.styleTransfer.RAIN_PRINCESS.XNNPACK_INT8,
43+
},
44+
{
45+
label: 'Rain Princess (CoreML FP16)',
46+
value: models.styleTransfer.RAIN_PRINCESS.COREML_FP16,
47+
disabled: Platform.OS !== 'ios',
48+
},
49+
{
50+
label: 'Udnie (XNNPACK INT8)',
51+
value: models.styleTransfer.UDNIE.XNNPACK_INT8,
52+
},
53+
{
54+
label: 'Udnie (CoreML FP16)',
55+
value: models.styleTransfer.UDNIE.COREML_FP16,
56+
disabled: Platform.OS !== 'ios',
57+
},
58+
];
59+
60+
function StyleTransferContent() {
61+
const insets = useSafeAreaInsets();
62+
const [selectedModel, setSelectedModel] = useState<any>(MODEL_OPTIONS[0].value);
63+
const [imageUri, setImageUri] = useState<string | null>(null);
64+
const [isProcessing, setIsProcessing] = useState(false);
65+
const [styledImage, setStyledImage] = useState<SkiaImageType | null>(null);
66+
const [showOriginal, setShowOriginal] = useState(false);
67+
const [latency, setLatency] = useState<number | null>(null);
68+
const [error, setError] = useState<string | null>(null);
69+
70+
const skiaImage = useImage(imageUri, (err) => setError(err.message || String(err)));
71+
72+
const {
73+
isReady,
74+
downloadProgress,
75+
error: loadError,
76+
transferStyle,
77+
transferStyleWorklet,
78+
} = useStyleTransfer(selectedModel);
79+
80+
const handlePickImage = async (useCamera: boolean) => {
81+
setError(null);
82+
try {
83+
const uri = await getImage(useCamera);
84+
if (uri) {
85+
setImageUri(uri);
86+
setStyledImage(null);
87+
setShowOriginal(false);
88+
setLatency(null);
89+
}
90+
} catch (e: any) {
91+
setError(e.message || String(e));
92+
}
93+
};
94+
95+
const runStyleTransfer = async (sync: boolean) => {
96+
if (!skiaImage || !transferStyle || !transferStyleWorklet) return;
97+
if (!sync) setIsProcessing(true);
98+
setError(null);
99+
try {
100+
const pixels = skiaImage.readPixels();
101+
if (!pixels) {
102+
throw new Error('Failed to read pixels from image');
103+
}
104+
if (!(pixels instanceof Uint8Array)) {
105+
throw new Error('Expected Uint8Array from readPixels');
106+
}
107+
const buffer = {
108+
data: pixels,
109+
width: skiaImage.width(),
110+
height: skiaImage.height(),
111+
format: 'rgba' as const,
112+
layout: 'hwc' as const,
113+
};
114+
115+
const start = Date.now();
116+
const output = sync ? transferStyleWorklet(buffer) : await transferStyle(buffer);
117+
118+
setLatency(Date.now() - start);
119+
120+
const outData = Skia.Data.fromBytes(output.data);
121+
const info = {
122+
width: skiaImage.width(),
123+
height: skiaImage.height(),
124+
colorType: ColorType.RGBA_8888,
125+
alphaType: AlphaType.Premul,
126+
};
127+
const skiaStyled = Skia.Image.MakeImage(info, outData, skiaImage.width() * 4);
128+
if (skiaStyled) {
129+
setStyledImage(skiaStyled);
130+
setShowOriginal(false);
131+
} else {
132+
throw new Error('Failed to decode styled image in Skia');
133+
}
134+
} catch (e: any) {
135+
setError(e.message || String(e));
136+
} finally {
137+
if (!sync) setIsProcessing(false);
138+
}
139+
};
140+
141+
const activeError = loadError ? String(loadError) : error;
142+
const activeImage = showOriginal ? skiaImage : (styledImage ?? skiaImage);
143+
144+
return (
145+
<ScrollView
146+
style={commonStyles.container}
147+
contentContainerStyle={[
148+
commonStyles.contentContainer,
149+
{ paddingBottom: insets.bottom + theme.spacing.large },
150+
]}
151+
>
152+
<Text style={commonStyles.description}>
153+
Upload or capture an image to apply artistic style transfer filters.
154+
</Text>
155+
156+
<ModelPicker
157+
label="Style Model"
158+
options={MODEL_OPTIONS}
159+
selectedValue={selectedModel}
160+
onValueChange={(model) => {
161+
setSelectedModel(model);
162+
setStyledImage(null);
163+
setShowOriginal(false);
164+
setLatency(null);
165+
setError(null);
166+
}}
167+
/>
168+
169+
<ModelStatus
170+
isReady={isReady}
171+
downloadProgress={downloadProgress}
172+
error={activeError}
173+
modelTypeLabel="style transfer model"
174+
/>
175+
176+
<ImageViewport skiaImage={activeImage} onPressPlaceholder={() => handlePickImage(false)} />
177+
178+
<View style={commonStyles.buttonRow}>
179+
<Button title="Gallery" onPress={() => handlePickImage(false)} variant="secondary" />
180+
<Button title="Camera" onPress={() => handlePickImage(true)} variant="secondary" />
181+
</View>
182+
183+
<View style={commonStyles.buttonRow}>
184+
<Button
185+
title="Run Async"
186+
onPress={() => runStyleTransfer(false)}
187+
disabled={!skiaImage || !isReady || isProcessing}
188+
loading={isProcessing}
189+
/>
190+
<Button
191+
title="Run Sync"
192+
onPress={() => runStyleTransfer(true)}
193+
disabled={!skiaImage || !isReady || isProcessing}
194+
variant="accent"
195+
/>
196+
</View>
197+
198+
{styledImage && (
199+
<View style={commonStyles.buttonRow}>
200+
<Button
201+
title={showOriginal ? 'Show Styled' : 'Show Original'}
202+
onPress={() => setShowOriginal(!showOriginal)}
203+
variant="secondary"
204+
/>
205+
</View>
206+
)}
207+
208+
<LatencyIndicator latency={latency} />
209+
</ScrollView>
210+
);
211+
}
212+
213+
export default function StyleTransferScreen() {
214+
return (
215+
<ScreenWrapper>
216+
<StyleTransferContent />
217+
</ScreenWrapper>
218+
);
219+
}
Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
import type { WorkletRuntime } from 'react-native-worklets';
2+
3+
import { tensor } from '../../../core/tensor';
4+
import { loadModel } from '../../../core/model';
5+
import { validateModelSchema, SymbolicTensor } from '../../../core/modelSchema';
6+
import { wrapAsync } from '../../../core/runtime';
7+
8+
import type { ImageBuffer } from '../image';
9+
import { createImagePreprocessor, type ImagePreprocessorOptions } from './preprocessing';
10+
import {
11+
toChannelsLast,
12+
normalize,
13+
cvtColor,
14+
resize,
15+
type InterpolationMethod,
16+
} from '../ops/image';
17+
18+
/**
19+
* Options for configuring the style transfer preprocessor and postprocessor.
20+
* @category Types
21+
*/
22+
export type StyleTransferOptions = Omit<ImagePreprocessorOptions, 'resizeMode'> & {
23+
readonly resizeMode: 'stretch';
24+
readonly outAlpha: number | number[];
25+
readonly outBeta: number | number[];
26+
readonly outInterpolation: InterpolationMethod;
27+
};
28+
29+
/**
30+
* Model configuration required to instantiate a style transfer task runner.
31+
* @category Types
32+
*/
33+
export type StyleTransferModel = {
34+
readonly modelPath: string;
35+
readonly opts: StyleTransferOptions;
36+
};
37+
38+
/**
39+
* Creates an image style transfer runner for executing local style transfer models.
40+
*
41+
* It validates the model inputs and outputs requirements, pre-allocates
42+
* the necessary static execution tensors, sets up an image preprocessor, and
43+
* registers clean disposal hooks to clear all native memory.
44+
* @category Typescript API
45+
* @param config Style transfer task configuration containing path and options.
46+
* @param runtime Optional worklet runtime thread on which to run the model execution.
47+
* @returns A promise resolving to an object containing style transfer and disposal controls.
48+
*/
49+
export async function createStyleTransfer(
50+
config: StyleTransferModel,
51+
runtime?: WorkletRuntime
52+
): Promise<{
53+
/**
54+
* Releases all allocated native resources.
55+
*/
56+
dispose: () => void;
57+
/**
58+
* Performs asynchronous image style transfer on the given input image.
59+
* @param input The input image buffer.
60+
* @returns A promise resolving to the styled image buffer.
61+
*/
62+
transferStyle: (input: ImageBuffer) => Promise<ImageBuffer>;
63+
/**
64+
* Synchronous version of {@link transferStyle} to be executed directly on the
65+
* caller or worklet thread.
66+
*/
67+
transferStyleWorklet: (input: ImageBuffer) => ImageBuffer;
68+
}> {
69+
const { modelPath, opts } = config;
70+
const model = await wrapAsync(loadModel, runtime)(modelPath);
71+
72+
const meta = validateModelSchema(
73+
model,
74+
'forward',
75+
[SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])],
76+
[SymbolicTensor('float32', [1, 3, 'H', 'W'], [3, 'H', 'W'])]
77+
);
78+
const inpShape = meta.inputTensorMeta[0]!.shape;
79+
const outShape = meta.outputTensorMeta[0]!.shape;
80+
81+
const targetH = outShape.at(-2)!;
82+
const targetW = outShape.at(-1)!;
83+
84+
const tensors = [
85+
tensor('float32', outShape),
86+
tensor('float32', [3, targetH, targetW]),
87+
tensor('uint8', [3, targetH, targetW]),
88+
tensor('uint8', [targetH, targetW, 3]),
89+
tensor('uint8', [targetH, targetW, 4]),
90+
] as const;
91+
92+
const [tOutput, tReshape, tUint8, tChanLast, tRgba] = tensors;
93+
const preprocessor = createImagePreprocessor(opts, inpShape);
94+
95+
const dispose = () => {
96+
tensors.forEach((t) => t.dispose());
97+
preprocessor.dispose();
98+
model.dispose();
99+
};
100+
101+
const transferStyleWorklet = (input: ImageBuffer): ImageBuffer => {
102+
'worklet';
103+
const tInput = preprocessor.process(input);
104+
model.execute('forward', [tInput], [tOutput]);
105+
106+
const data = new Uint8Array(input.width * input.height * 4);
107+
const tResize = tensor('uint8', [input.height, input.width, 4]);
108+
try {
109+
tOutput
110+
.copyTo(tReshape)
111+
.through(normalize, tUint8, { alpha: opts.outAlpha, beta: opts.outBeta })
112+
.through(toChannelsLast, tChanLast)
113+
.through(cvtColor, tRgba, 'RGB2RGBA')
114+
.through(resize, tResize, { mode: 'stretch', interpolation: opts.outInterpolation })
115+
.getData(data);
116+
} finally {
117+
tResize.dispose();
118+
}
119+
120+
return { data, width: input.width, height: input.height, format: 'rgba', layout: 'hwc' };
121+
};
122+
123+
const transferStyle = wrapAsync(transferStyleWorklet, runtime);
124+
125+
return { transferStyle, transferStyleWorklet, dispose };
126+
}

0 commit comments

Comments
 (0)