Baby Yoda Segmentor

Instance segmentation model for detection of the character Baby Yoda, from the Disney TV Show The Mandalorian.

Dataset

This repository depends on baby-yoda-segmentation-dataset, built as a living-dataset.

How to train

Locally

1. Install the project dependencies with poetry
2. Change params, code or data as needed
3. Run

dvc repro auto-train

The model will be saved in models/model.pth

Using Google Colab

1. Open the Colab Notebook. It is also available in src/ColabNotebook.ipynb
2. Run according to the steps

Using the model

Get the model

With DVC

In a DVC repository run:

dvc import https://dagshub.com/simon/baby-yoda-segmentor models/model.pth

Simple download

curl -O https://dagshub.com/Simon/baby-yoda-segmentor/raw/master/models/model.pth

Load and use the model

from PIL import Image
from torchvision.transforms import ToTensor

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

img = Image.open("image.png")
img_t = ToTensor()(img)
model.eval()
with torch.no_grad():
    prediction = model([img_t.to(device)])

How to contribute

1. Fork the repository

git clone <fork-url>
dvc pull -r origin

1. Do your changes
2. Train the model
3. Add a local remote to push your data

dvc remote add --local fork <dagshub-remote-url.dvc>
# Additional commands to set up credentials should appear on you fork homepage

5. Push your code and data

dvc push -r fork
git add .
git commit -m "Changes to dataset"
git push

6. Open a PR