Spaces:
Sleeping
Sleeping
Update app/process.py
Browse files- 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 |
+
|