# AUTOGENERATED! DO NOT EDIT! File to edit: ModelTester.ipynb.

# %% auto 0
__all__ = ['learn', 'categories', 'examples', 'intf', 'OrdinalRegressionMetric', 'classify_image']

# %% ModelTester.ipynb 1
from fastai.vision.all import *
import gradio as gr
import os
# needed when model was saved/ported on windows
import platform
import pathlib



# %% ModelTester.ipynb 3
from fastai.metrics import Metric

class OrdinalRegressionMetric(Metric):
    def __init__(self):
        super().__init__()
        self.total = 0
        self.count = 0

    def accumulate(self, learn):
        # Get predictions and targets
        preds, targs = learn.pred, learn.y

        # Your custom logic to convert predictions and targets to numeric values
        preds_numeric = torch.argmax(preds, dim=1)
        targs_numeric = targs

        #print("preds_numeric: ",preds_numeric)
        #print("targs_numeric: ",targs_numeric)

        # Calculate the metric (modify this based on your specific needs)
        squared_diff = torch.sum(torch.sqrt((preds_numeric - targs_numeric)**2))

        # Normalize by the maximum possible difference
        max_diff = torch.sqrt((torch.max(targs_numeric) - torch.min(targs_numeric))**2)

        #print("squared_diff: ",squared_diff)
        #print("max_diff: ",max_diff)

        # Update the metric value
        self.total += squared_diff
        #print("self.total: ",self.total)
        self.count += max_diff
        #print("self.count: ",self.count)
    @property
    def value(self):
        if self.count == 0:
            return 0.0  # or handle this case appropriately
        #print("self.total / self.count: ", (self.total / self.count))
        # Calculate the normalized metric value
        metric_value = 1/(self.total / self.count)
        return metric_value

# %% ModelTester.ipynb 4
plt = platform.system()
if plt == 'Linux': pathlib.WindowsPath = pathlib.PosixPath
learn = load_learner("newmodel.pk1")

# %% ModelTester.ipynb 6
categories = ("1","1-2","2","2-3","3","3-4","4","4-5","5")

#HF_TOKEN = os.getenv('HF_TOKEN')

#hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "FJ-flagging")

def classify_image(img):
    pred, idx, probs = learn.predict(img)
    return dict(zip(categories, map(float, probs)))

# %% ModelTester.ipynb 8
examples = ['2-3.jpg','1.jpg','4.jpg']

intf = gr.Interface(fn=classify_image, 
                    inputs='image', 
                    outputs='label', 
                    examples=examples, 
                    allow_flagging="manual", 
                    flagging_options=["1","1-2","2","2-3","3","3-4","4","4-5","5"])
                    #flagging_callback=hf_writer)
intf.launch(inline=False)