tomc43841/vgg_balanced_smash6

This is how I envision this to work after you've cloned this repository and are making changes inside of it.

from vgg_balanced_smash6.vgg_model import CustomSimpleImageRegressor
from vgg_balanced_smash6.custom_image_processor import CustomImageProcessor


def custom_data_collator_function(processor):
    def return_func(batch):
        images = [item["image"] for item in batch]

        inputs = processor(images, return_tensors="pt")
        targets = torch.tensor([[
            item["Start"],
            item["A"],
            item["B"],
            item["X"],
            item["Y"],
            item["Z"],
            item["DPadUp"],
            item["DPadDown"],
            item["DPadLeft"],
            item["DPadRight"],
            item["L"],
            item["R"],
            item["LPressure"] / 255,
            item["RPressure"] / 255,
            item["XAxis"] / 255,
            item["YAxis"] / 255,
            item["CXAxis"] / 255,
            item["CYAxis"] / 255] for item in batch], dtype=torch.float32)

        return inputs, targets

    return return_func


def do_stuff_with_model():
    custom_data_collator = custom_data_collator_function(CustomImageProcessor())
    train_loader = DataLoader(dataset["train"], batch_size=BATCH_SIZE, collate_fn=custom_data_collator)

    model = CustomSimpleImageRegressor()
    model.update_model_from_checkpoint(6)

    # now it's time to train the model with the train_loader

and there you go you have a model trained on my (admittedly awful) gameplay