1+ import { useState , useRef , useEffect } from "react" ;
2+ import { useDispatch , useSelector } from "react-redux" ;
3+ import type { AppDispatch } from "../shared/redux/store" ;
4+ import { setImageUrl , setPreprocessedImageUrl } from "../shared/redux/stringArtSlice" ;
5+
6+
7+ const TARGET_SIZE = 1080 ;
8+
9+ declare global {
10+ interface Window {
11+ SelfieSegmentation ?: any ;
12+ }
13+ }
14+
15+ async function loadSelfieSegmentation ( ) : Promise < any > {
16+ if ( ( window as any ) . SelfieSegmentation ) return ( window as any ) . SelfieSegmentation ;
17+ return new Promise ( ( resolve , reject ) => {
18+ const script = document . createElement ( "script" ) ;
19+ script . src = "https://cdn.jsdelivr.net/npm/@mediapipe/selfie_segmentation/selfie_segmentation.js" ;
20+ script . async = true ;
21+ script . onload = ( ) => {
22+ resolve ( ( window as any ) . SelfieSegmentation ) ;
23+ } ;
24+ script . onerror = reject ;
25+ document . head . appendChild ( script ) ;
26+ } ) ;
27+ }
28+
29+
30+ export const usePreProcess = ( ) => {
31+ const dispatch = useDispatch < AppDispatch > ( ) ;
32+ const imageUrl = useSelector ( ( state : any ) => state . stringArt . imageUrl || state . stringArt . preprocessedImageUrl ) ;
33+ const [ loading , setLoading ] = useState ( false ) ;
34+ const [ preview , setPreview ] = useState < string | null > ( null ) ;
35+ const selfieRef = useRef < any > ( null ) ;
36+
37+ useEffect ( ( ) => {
38+ // Lazy load SelfieSegmentation (cdn) — we'll use it after we have a crop
39+ loadSelfieSegmentation ( )
40+ . then ( ( SelfieSegmentation ) => {
41+ selfieRef . current = new SelfieSegmentation ( {
42+ locateFile : ( file : string ) =>
43+ `https://cdn.jsdelivr.net/npm/@mediapipe/selfie_segmentation/${ file } ` ,
44+ } ) ;
45+ selfieRef . current . setOptions ( { modelSelection : 1 } ) ;
46+ } )
47+ . catch ( ( err ) => {
48+ console . warn ( "Failed to load SelfieSegmentation" , err ) ;
49+ } ) ;
50+ // FaceDetector (tasks-vision) is relatively heavy and may already be used elsewhere;
51+ // for simplicity we'll dynamically import it from the same CDN used in the demo if needed.
52+ // If project has a central copy, refactor to reuse it.
53+ // We do not eagerly instantiate tasks-vision FaceDetector here to keep initial load small.
54+ } , [ ] ) ;
55+
56+ async function processImage ( ) {
57+ if ( ! imageUrl ) return ;
58+ setLoading ( true ) ;
59+ setPreview ( null ) ;
60+
61+ // Load image element
62+ const img = new Image ( ) ;
63+ img . crossOrigin = "anonymous" ;
64+ img . src = imageUrl ;
65+ await new Promise ( ( res , rej ) => {
66+ img . onload = res ;
67+ img . onerror = rej ;
68+ } ) ;
69+
70+ const cropCanvas = document . createElement ( "canvas" ) ;
71+ cropCanvas . width = TARGET_SIZE ;
72+ cropCanvas . height = TARGET_SIZE ;
73+ const ctx = cropCanvas . getContext ( "2d" ) ;
74+ if ( ! ctx ) {
75+ setLoading ( false ) ;
76+ return ;
77+ }
78+ // draw full image to canvas preserving aspect ratio (letterbox)
79+ const scale = Math . min ( TARGET_SIZE / img . width , TARGET_SIZE / img . height ) ;
80+ const drawWidth = Math . round ( img . width * scale ) ;
81+ const drawHeight = Math . round ( img . height * scale ) ;
82+ const dx = Math . round ( ( TARGET_SIZE - drawWidth ) / 2 ) ;
83+ const dy = Math . round ( ( TARGET_SIZE - drawHeight ) / 2 ) ;
84+ // fill background white to avoid transparent borders
85+ ctx . fillStyle = "#ffffff" ;
86+ ctx . fillRect ( 0 , 0 , TARGET_SIZE , TARGET_SIZE ) ;
87+ ctx . drawImage ( img , 0 , 0 , img . width , img . height , dx , dy , drawWidth , drawHeight ) ;
88+
89+ // Run selfie segmentation on the cropped canvas
90+ if ( ! selfieRef . current ) {
91+ console . warn ( "SelfieSegmentation not loaded; returning plain crop on white" ) ;
92+ // Composite crop on white and return
93+ const whiteCanvas = document . createElement ( "canvas" ) ;
94+ whiteCanvas . width = TARGET_SIZE ;
95+ whiteCanvas . height = TARGET_SIZE ;
96+ const wctx = whiteCanvas . getContext ( "2d" ) ;
97+ if ( ! wctx ) {
98+ setLoading ( false ) ;
99+ return ;
100+ }
101+ wctx . fillStyle = "#ffffff" ;
102+ wctx . fillRect ( 0 , 0 , TARGET_SIZE , TARGET_SIZE ) ;
103+ wctx . drawImage ( cropCanvas , 0 , 0 ) ;
104+ const dataUrl = whiteCanvas . toDataURL ( "image/png" ) ;
105+ setPreview ( dataUrl ) ;
106+ setLoading ( false ) ;
107+ return ;
108+ }
109+
110+ // SelfieSegmentation in browser: use onResults pattern; we will send the crop canvas as an ImageBitmap
111+ const segmentationResult = await new Promise < any > ( ( resolve ) => {
112+ ( async ( ) => {
113+ selfieRef . current . onResults ( ( results : any ) => {
114+ resolve ( results ) ;
115+ } ) ;
116+ // selfieRef expects an HTMLVideoElement/Image/Canvas; pass the cropCanvas
117+ await selfieRef . current . send ( { image : cropCanvas } ) ;
118+ } ) ( ) ;
119+ } ) ;
120+
121+ // segmentationResult.segmentationMask is an HTMLCanvas or ImageData depending on implementation;
122+ // To be robust, draw the mask onto a temp canvas.
123+ const maskCanvas = document . createElement ( "canvas" ) ;
124+ maskCanvas . width = TARGET_SIZE ;
125+ maskCanvas . height = TARGET_SIZE ;
126+ const mctx = maskCanvas . getContext ( "2d" ) ;
127+ if ( ! mctx ) {
128+ setLoading ( false ) ;
129+ return ;
130+ }
131+
132+ if ( segmentationResult . segmentationMask ) {
133+ // draw mask to canvas
134+ mctx . drawImage ( segmentationResult . segmentationMask , 0 , 0 , TARGET_SIZE , TARGET_SIZE ) ;
135+ } else if ( segmentationResult . multiSegmentationMask ) {
136+ mctx . drawImage ( segmentationResult . multiSegmentationMask [ 0 ] , 0 , 0 , TARGET_SIZE , TARGET_SIZE ) ;
137+ } else if ( segmentationResult . mask ) {
138+ // mask as float array: convert to ImageData
139+ const imgData = mctx . createImageData ( TARGET_SIZE , TARGET_SIZE ) ;
140+ const mask = segmentationResult . mask ; // assume Float32Array or array of [0..1]
141+ for ( let i = 0 ; i < mask . length && i * 4 < imgData . data . length ; i ++ ) {
142+ const a = Math . round ( Math . min ( 1 , Math . max ( 0 , mask [ i ] ) ) * 255 ) ;
143+ imgData . data [ i * 4 + 0 ] = 255 ;
144+ imgData . data [ i * 4 + 1 ] = 255 ;
145+ imgData . data [ i * 4 + 2 ] = 255 ;
146+ imgData . data [ i * 4 + 3 ] = a ;
147+ }
148+ mctx . putImageData ( imgData , 0 , 0 ) ;
149+ } else {
150+ // unknown mask format, fallback
151+ mctx . fillStyle = "#ffffff" ;
152+ mctx . fillRect ( 0 , 0 , TARGET_SIZE , TARGET_SIZE ) ;
153+ }
154+
155+ // Optional: feather edges by applying a small blur filter to mask
156+ // create final canvas: draw white background, then draw image masked by alpha
157+ const finalCanvas = document . createElement ( "canvas" ) ;
158+ finalCanvas . width = TARGET_SIZE ;
159+ finalCanvas . height = TARGET_SIZE ;
160+ const fctx = finalCanvas . getContext ( "2d" ) ;
161+ if ( ! fctx ) {
162+ setLoading ( false ) ;
163+ return ;
164+ }
165+
166+ // Composite: draw white background and keep only foreground using the mask
167+ // Paint white background first
168+ fctx . fillStyle = "#ffffff" ;
169+ fctx . fillRect ( 0 , 0 , TARGET_SIZE , TARGET_SIZE ) ;
170+ // Draw the crop image on top
171+ fctx . drawImage ( cropCanvas , 0 , 0 , TARGET_SIZE , TARGET_SIZE ) ;
172+ // Use mask to keep only foreground pixels (mask white = foreground)
173+ fctx . globalCompositeOperation = 'destination-in' ;
174+ fctx . drawImage ( maskCanvas , 0 , 0 , TARGET_SIZE , TARGET_SIZE ) ;
175+ // Restore default composite mode
176+ fctx . globalCompositeOperation = 'source-over' ;
177+
178+ // Convert final canvas to data URL
179+ const finalDataUrl = finalCanvas . toDataURL ( "image/png" ) ;
180+ setPreview ( finalDataUrl ) ;
181+ // Automatically set the preprocessed image as the image to process
182+ dispatch ( setImageUrl ( finalDataUrl ) ) ;
183+ dispatch ( setPreprocessedImageUrl ( finalDataUrl ) ) ;
184+ setLoading ( false ) ;
185+ }
186+
187+ useEffect ( ( ) => {
188+ if ( imageUrl && ! preview && ! loading ) {
189+ processImage ( ) ;
190+ }
191+ // eslint-disable-next-line
192+ } , [ imageUrl ] ) ;
193+
194+
195+ return preview ;
196+ }
0 commit comments