from gardens.models import *

import torch
import os
# TensorFlow and Keras libraries
import tensorflow as tf
from tensorflow import keras
# used to chech device
from tensorflow.python.client import device_lib
# Other libraries
import numpy as np
from PIL import Image

from django.conf import settings

import shutil

pearModel = os.path.join(settings.STATIC_ROOT, 'MachineLearningModels/PearDetectionModel.pt')

def pearDetectionModel(photo):
    # model loaded for detecting pears
    model = torch.hub.load('ultralytics/yolov5', 'custom', path=pearModel)
    # defining the photo  needed to check
    img = photo.image_uri.path

    results = model(img)
    results.crop()
    print(results)

    for folder in os.listdir('runs/detect'):
        directory = f'runs/detect/{folder}/crops/wine'
        amountOfPears = len([entry for entry in os.listdir(directory) if os.path.isfile(os.path.join(directory, entry))])
        state = TreeState.objects.get(tree_id = photo.tree_id, mission_id = photo.mission_id)
        try:
            state = TreeState.objects.get(tree_id = photo.tree_id, mission_id = photo.mission_id)
            state.yields = amountOfPears
            state.save()
        except TreeState.DoesNotExist:
            TreeState.objects.create(tree_id = photo.tree_id, mission_id = photo.mission_id, yields = amountOfPears, has_scab = True)
        shutil.rmtree(f'runs/detect/{folder}/')
        photo.delete()