1+ import React , { useState , useEffect } from 'react'
2+ import { Alert } from 'react-native'
3+ // @ts -ignore
4+ import Tflite from 'tflite-react-native'
5+
6+ import { ChooseFile , TakePicture , Translate } from '~/services'
7+
8+ import { Button , ButtonContainer , ImagePreview , Recognition , Title } from '~/components'
9+
10+ import { RECOGNITION } from '~/components/Recognition'
11+
12+
13+ type MODELS = {
14+ ssdMobileNet : ( ) => any ,
15+ yolo : ( ) => any
16+ }
17+
18+
19+ const tflite = new Tflite ( )
20+
21+
22+ function ObjectDetection ( ) {
23+ const [ chosenModel , setChosenModel ] = useState < keyof MODELS | null > ( null )
24+ const [ fileUri , setFileUri ] = useState ( '' )
25+ const [ recognition , setRecognition ] = useState < RECOGNITION [ ] > ( [ ] )
26+
27+
28+ const loadModel : MODELS = {
29+ ssdMobileNet : ( ) => tflite . loadModel ( {
30+ model : 'models/ssd_mobilenet.tflite' ,
31+ labels : 'models/ssd_mobilenet.txt' ,
32+ numThreads : 1
33+ } ,
34+ ( error : string ) => {
35+ if ( error ) {
36+ Alert . alert ( Translate ( 'error' ) , Translate ( 'errorLoadModel' ) )
37+ }
38+ } ) ,
39+ yolo : ( ) => tflite . loadModel ( {
40+ model : 'models/yolov2_tiny.tflite' ,
41+ labels : 'models/yolov2_tiny.txt' ,
42+ numThreads : 1
43+ } ,
44+ ( error : string ) => {
45+ if ( error ) {
46+ Alert . alert ( Translate ( 'error' ) , Translate ( 'errorLoadModel' ) )
47+ }
48+ } )
49+ }
50+
51+
52+ const performObjectDetection : MODELS = {
53+ ssdMobileNet : ( ) => tflite . detectObjectOnImage ( {
54+ path : fileUri ,
55+ model : 'SSDMobileNet' ,
56+ imageMean : 127.5 ,
57+ imageStd : 127.5 ,
58+ threshold : 0.3 ,
59+ numResultsPerClass : 2
60+ } ,
61+ ( error : string , response : RECOGNITION [ ] ) => {
62+ if ( error ) {
63+ Alert . alert ( Translate ( 'error' ) , Translate ( 'errorProcessingTflite' ) )
64+ return
65+ }
66+ setRecognition ( response )
67+ } ) ,
68+ yolo : ( ) => tflite . detectObjectOnImage ( {
69+ path : fileUri ,
70+ model : 'YOLO' ,
71+ imageMean : 0.0 ,
72+ imageStd : 255.0 ,
73+ threshold : 0.3 ,
74+ numResultsPerClass : 2
75+ } ,
76+ ( error : string , response : RECOGNITION [ ] ) => {
77+ if ( error ) {
78+ Alert . alert ( Translate ( 'error' ) , Translate ( 'errorProcessingTflite' ) )
79+ return
80+ }
81+ setRecognition ( response )
82+ } )
83+ }
84+
85+
86+ function ChooseModel ( ) {
87+ return (
88+ < >
89+ < Title text = 'chooseModel' />
90+ < ButtonContainer style = { { width : '60%' , maxWidth : 300 } } >
91+ < Button
92+ onPress = { ( ) => setChosenModel ( 'ssdMobileNet' ) }
93+ text = 'SSD MobileNet'
94+ translate = { false }
95+ />
96+ < Button
97+ onPress = { ( ) => setChosenModel ( 'yolo' ) }
98+ text = 'Tiny YOLO'
99+ translate = { false }
100+ />
101+ </ ButtonContainer >
102+ </ >
103+ )
104+ }
105+
106+
107+ useEffect ( ( ) => {
108+ if ( chosenModel && fileUri ) {
109+ performObjectDetection [ chosenModel ] ( )
110+ }
111+ } , [ chosenModel , fileUri ] )
112+
113+
114+ useEffect ( ( ) => {
115+ if ( chosenModel ) {
116+ loadModel [ chosenModel ] ( )
117+ }
118+
119+ return ( ) => {
120+ if ( chosenModel ) {
121+ tflite . close ( )
122+ }
123+ }
124+ } , [ chosenModel ] )
125+
126+
127+ if ( ! chosenModel ) return (
128+ < ChooseModel />
129+ )
130+
131+
132+ return (
133+ < >
134+ < Title text = 'pickImagesCG' />
135+ < ImagePreview uri = { fileUri } >
136+ { /* @ts -ignore */ recognition && (
137+ < Recognition recognition = { recognition } />
138+ ) }
139+ </ ImagePreview >
140+ < ButtonContainer style = { { width : '60%' , maxWidth : 300 } } >
141+ < Button onPress = { ( ) => ChooseFile ( setFileUri ) } text = 'chooseFile' />
142+ < Button onPress = { ( ) => TakePicture ( setFileUri ) } text = 'takePicture' />
143+ </ ButtonContainer >
144+ </ >
145+ )
146+ }
147+
148+
149+ export default ObjectDetection
0 commit comments