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

appleModel = os.path.join(settings.STATIC_ROOT, 'MachineLearningModels/AppleDetectionModel.pt')

def appleDetectionModel(photo):
    # model loaded for detecting pears
    model = torch.hub.load('ultralytics/yolov5', 'custom', path=appleModel)
    # 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/apples'
        # check if directory exists
        if os.path.exists(directory):
            amountOfPears = len([entry for entry in os.listdir(directory) if os.path.isfile(os.path.join(directory, entry))])
        else:
            amountOfPears = 0
        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()