from ultralytics import YOLO
from tensorflow.keras.preprocessing.image import img_to_array
import numpy as np
import cv2
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.mobilenet_v3 import preprocess_input
import PIL.Image  as Image
import io
import base64
import os 

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
YOLO_PATH= os.path.join(BASE_DIR,"models",'Yolo.pt')
MOBILENET_PATH= os.path.join(BASE_DIR,"models",'MobileNetV3_rust_classifier.keras')

YOLO_MODEL=YOLO(YOLO_PATH)
MOBILENET_MODEL=load_model(MOBILENET_PATH)



def predict(image_bytes):
    img = Image.open(io.BytesIO(image_bytes)).resize((224, 224))
    x = img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)
    proba = MOBILENET_MODEL.predict(x)[0]
    prediction = np.argmax(proba)
    return prediction,proba[prediction]
    

def detect(image_bytes):
    image  = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    result = YOLO_MODEL(image)[0]
    img_result= Image.fromarray(result.plot()) 
    buffer=io.BytesIO()
    img_result.save(buffer,format="PNG")
    base64_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
    uri =f"data:image/png;base64,{base64_str}"
    return uri