import ultralytics
ultralytics.checks()
from render_results import annotate_image_prediction
from PIL import Image

if __name__=='__main__':
    # load model
    model = ultralytics.YOLO('../model_v3_202402021.pt')
    
    # set parameters
    image_path = '../images/1688033955437.jpg'
    saving_folder = '../images'
    
    # infer
    image = Image.open(image_path)
    results = model.predict(image_path)
    
    # plot
    hex_class_colors = {'green_cherry': '#9CF09A',
                    'yellow_cherry': '#F3C63D',
                    'red_cherry': '#F44336',
                    'dark_brown_cherry': '#C36105',
                    'low_visibility_unsure': '#02D5FA'}
    class_dic = {1: 'green_cherry', 3: 'yellow_cherry', 2: 'red_cherry', 0: 'dark_brown_cherry'}
    render = annotate_image_prediction(
        image_path=image_path,
        yolo_boxes=results[0].boxes,
        class_dic=class_dic,
        saving_folder=saving_folder,
        hex_class_colors=hex_class_colors,
        show=True,
        font_scale=5,
        font_thickness=10,
        )