import gradio as gr
from ultralytics import YOLO
import cv2

# Laden Sie das YOLO-Modell
model = YOLO('yolov8m-seg.pt')

# Funktion für die Vorhersage und Rückgabe des bearbeiteten Bildes
def predict_and_visualize(image):
    # Führen Sie die Vorhersage durch
    prediction = model.predict(image, imgsz=(1024, 1024), show=False, save=True)
    # Nehmen Sie das generierte Bild, das von YOLO gespeichert wurde
  
    processed_image_path = prediction.__getitem__(0).save_dir + "/" +prediction.__getitem__(0).path
    print(processed_image_path)
    #print("Image PATH: ",processed_image_path)

    # Laden Sie das Bild und geben Sie es zurück
    processed_image = cv2.imread(processed_image_path)
    return processed_image

# Gradio-Schnittstelle für die Bildbearbeitungsfunktion
iface = gr.Interface(fn=predict_and_visualize, inputs="image", outputs="image")
iface.launch()