# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Object Detection
# In this notebook we will have a look at the fridge dataset using different dashboards.

# Imports
from icevision.all import *
import icedata
from icevision_dashboards.data import *

# ## Setup data
#

path = icedata.fridge.load_data()
class_map = icedata.fridge.class_map()
parser = icedata.fridge.parser(data_dir=path)
train_records, valid_records = parser.parse()



record_ds = RecordDataset(valid_records, class_map)

record_ds.name = "record_dataset"

record_ds.save("test_data")





# ## Train model for data generation

# Define transforms
train_tfms = tfms.A.Adapter(
    [*tfms.A.aug_tfms(size=384, presize=512), tfms.A.Normalize()]
)
valid_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(384), tfms.A.Normalize()])

# Create datasets for training
train_ds = Dataset(train_records, train_tfms)
valid_ds = Dataset(valid_records, valid_tfms)

model_type = models.mmdet.faster_rcnn
backbone = model_type.backbones.resnet50_fpn_1x(pretrained=True)

# + tags=[]
model = model_type.model(backbone=backbone, num_classes=len(class_map))
metrics = [COCOMetric(metric_type=COCOMetricType.bbox)]
train_dl = model_type.train_dl(train_ds, batch_size=1, num_workers=4, shuffle=True)
valid_dl = model_type.valid_dl(valid_ds, batch_size=1, num_workers=4, shuffle=False)

learn = model_type.fastai.learner(dls=[train_dl, valid_dl], model=model, metrics=metrics)

learn.fine_tune(5, freeze_epochs=1, base_lr=1e-4)
# -

# ## Create preds and sampels files

_ = model.to("cuda:0")

samples, losses_stats = model_type.interp.get_losses(model, valid_ds)

dl = model_type.interp.infer_dl(valid_ds, batch_size=2)
preds = model_type.interp.predict_from_dl(model=model, infer_dl=dl, keep_images=True)


def get_compnent(components, component_type):
    for component in components:
        if isinstance(component, component_type):
            return component
    else:
        return None


def remove_image(components):
    for entry in list(components):
        if isinstance(entry, FilepathRecordComponent) or isinstance(entry, ImageRecordComponent):
            entry.img = None


def cleanup_preds(preds):
    new_preds = deepcopy(preds)
    for pred in new_preds:
        remove_image(pred.ground_truth.common.components)
        remove_image(pred.pred.common.components)
    return new_preds


# remove all the data not required (image) and convert the masks to a smaller data fromat
clean_preds = cleanup_preds(preds)
pickle.dump(clean_preds, open("test_data/object_detection_preds.pkl", "wb"))
preds = pickle.load(open("test_data/object_detection_preds.pkl", "rb"))


def cleanup_samples(samples):
    new_samples = deepcopy(samples)
    for sample in new_samples:
        remove_image(sample.components)
    return new_samples


# remove all the data not required (image)
clean_samples = cleanup_samples(samples)
pickle.dump(clean_samples, open("test_data/object_detection_samples.pkl", "wb"))
samples = pickle.load(open("test_data/object_detection_samples.pkl", "rb"))

valid_result_ds = ObjectDetectionResultsDataset.init_from_preds_and_samples(preds, samples, class_map=class_map)
valid_result_ds.save("test_data/object_detection_result_ds.dat")

# ## Cleanup

shutil.rmtree("checkpoints/")



StarcoderdataJupyterScriptsDedupFiltered