MsChabane/RustDetective

1 2 3from ultralytics import YOLO4from tensorflow.keras.preprocessing.image import img_to_array5import numpy as np6import cv27from tensorflow.keras.models import load_model8from tensorflow.keras.applications.mobilenet_v3 import preprocess_input9import PIL.Image  as Image10import io11import base6412import os 13 14BASE_DIR = os.path.dirname(os.path.abspath(__file__))15YOLO_PATH= os.path.join(BASE_DIR,"models",'Yolo.pt')16MOBILENET_PATH= os.path.join(BASE_DIR,"models",'MobileNetV3_rust_classifier.keras')17 18YOLO_MODEL=YOLO(YOLO_PATH)19MOBILENET_MODEL=load_model(MOBILENET_PATH)20 21 22 23def 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 33def 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 uri42    43 44 45 46 47 48 49