from fastai.vision.all import *
from fastai import *
from fastai.vision import *
from fastai.vision.widgets import *

# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    counter = 0
    for filename in filenames:
        counter = counter + 1
        #print(os.path.join(dirname, filename))

print(f"Loaded all files from the directory: {dirname}")
print(f"Number of files loaded:{counter}")
    

# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session

dir_name = '../input/proton-collision-image-set/Proton Collision 13TeV/Train'

path = Path(dir_name)

path.ls()

dls = DataBlock(
    blocks=(ImageBlock, CategoryBlock), 
    get_items=get_image_files, 
    splitter=RandomSplitter(valid_pct=0.2, seed=42),
    get_y=parent_label,
    item_tfms=[Resize(128, method='squish')]
).dataloaders(path, bs=32)

dls.show_batch(max_n=9)

#learn.fine_tune(3)

learn = vision_learner(dls, resnet18, metrics = error_rate)
learn.lr_find()

learn.fine_tune(4, 4.7e-3)

img = PILImage.create('../input/proton-collision-image-set/Proton Collision 13TeV/Test/TTbar/ttbar_lepFilter_13TeV_1087_1934.png')
what_proton,_,probs = learn.predict(img)
idx = 0
if what_proton == "TTbar":
    idx = 1
elif what_proton == "WJets":
    idx = 2
print(f"This collision belongs to category of: {what_proton}")
print(f"Probability it belongs to that category: {probs[idx]:.4f}")
img

interp = ClassificationInterpretation.from_learner(learn)
interp.plot_confusion_matrix()

learn.show_results()

interp.plot_top_losses(5, nrows=1)

cleaner = ImageClassifierCleaner(learn)
cleaner