Grogu Segmentor

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

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/grogu-segmentor models/model.pth

Simple download

curl -O https://dagshub.com/Simon/grogu-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

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