import os
import shutil
from fastai.vision.all import *
import gradio as gr

model_path = 'CameraTrapModel.pkl'
learn = load_learner(model_path)

def predict_class(file_path):
    pred, pred_idx, probs = learn.predict(file_path)
    return str(pred)

def classify_image(file):
    # Create your own destination folders
    animal_folder = '/content/gdrive/MyDrive/My_DataSets/With_animal(s)_Rslts'
    no_animal_folder = '/content/gdrive/MyDrive/My_DataSets/No_animal(s)_Rslts'

    os.makedirs(animal_folder, exist_ok=True)
    os.makedirs(no_animal_folder, exist_ok=True)

    # Move the file to the respective folders
    pred_class = predict_class(file.name)
    if pred_class == 'With_Animals':
        shutil.move(file.name, os.path.join(animal_folder, file.name))
    elif pred_class == 'No_Animals':
        shutil.move(file.name, os.path.join(no_animal_folder, file.name))
    
    return 'File classified successfully!'

iface = gr.Interface(fn=classify_image, inputs='file', outputs='text', title='Image Classifier From CameraTrap')
iface.launch()