MsChabane commited on
Commit
5c23ccf
·
verified ·
1 Parent(s): f11ffdb

Update app/process.py

Browse files
Files changed (1) hide show
  1. app/process.py +48 -48
app/process.py CHANGED
@@ -1,48 +1,48 @@
1
-
2
-
3
- from ultralytics import YOLO
4
- from tensorflow.keras.preprocessing.image import img_to_array
5
- import numpy as np
6
- import cv2
7
- from tensorflow.keras.models import load_model
8
- from tensorflow.keras.applications.mobilenet_v3 import preprocess_input
9
- import PIL.Image as Image
10
- import io
11
- import base64
12
- import os
13
-
14
- BASE_DIR = os.path.dirname(os.path.abspath(__file__))
15
- YOLO_PATH= os.join(BASE_DIR,"models",'Yolo.pt')
16
- MOBILENET_PATH= os.join(BASE_DIR,"models",'MobileNetV3_rust_classifier.keras')
17
-
18
- YOLO_MODEL=YOLO(YOLO_PATH)
19
- MOBILENET_MODEL=load_model(MOBILENET_PATH)
20
-
21
-
22
-
23
- def predict(image_bytes):
24
- img = Image.open(io.BytesIO(image_bytes)).resize((224, 224))
25
- x = img_to_array(img)
26
- x = np.expand_dims(x, axis=0)
27
- x = preprocess_input(x)
28
- proba = MOBILENET_MODEL.predict(x)[0]
29
- prediction = np.argmax(proba)
30
- return prediction,proba[prediction]
31
-
32
-
33
- def detect(image_bytes):
34
- image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
35
- result = YOLO_MODEL(image)[0]
36
- img_result= Image.fromarray(result.plot())
37
- buffer=io.BytesIO()
38
- img_result.save(buffer,format="PNG")
39
- base64_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
40
- uri =f"data:image/png;base64,{base64_str}"
41
- return uri
42
-
43
-
44
-
45
-
46
-
47
-
48
-
 
1
+
2
+
3
+ from ultralytics import YOLO
4
+ from tensorflow.keras.preprocessing.image import img_to_array
5
+ import numpy as np
6
+ import cv2
7
+ from tensorflow.keras.models import load_model
8
+ from tensorflow.keras.applications.mobilenet_v3 import preprocess_input
9
+ import PIL.Image as Image
10
+ import io
11
+ import base64
12
+ import os
13
+
14
+ BASE_DIR = os.path.dirname(os.path.abspath(__file__))
15
+ YOLO_PATH= os.path.join(BASE_DIR,"models",'Yolo.pt')
16
+ MOBILENET_PATH= os.path.join(BASE_DIR,"models",'MobileNetV3_rust_classifier.keras')
17
+
18
+ YOLO_MODEL=YOLO(YOLO_PATH)
19
+ MOBILENET_MODEL=load_model(MOBILENET_PATH)
20
+
21
+
22
+
23
+ def predict(image_bytes):
24
+ img = Image.open(io.BytesIO(image_bytes)).resize((224, 224))
25
+ x = img_to_array(img)
26
+ x = np.expand_dims(x, axis=0)
27
+ x = preprocess_input(x)
28
+ proba = MOBILENET_MODEL.predict(x)[0]
29
+ prediction = np.argmax(proba)
30
+ return prediction,proba[prediction]
31
+
32
+
33
+ def detect(image_bytes):
34
+ image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
35
+ result = YOLO_MODEL(image)[0]
36
+ img_result= Image.fromarray(result.plot())
37
+ buffer=io.BytesIO()
38
+ img_result.save(buffer,format="PNG")
39
+ base64_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
40
+ uri =f"data:image/png;base64,{base64_str}"
41
+ return uri
42
+
43
+
44
+
45
+
46
+
47
+
48
+