classification.qmd

# 2.3 Image Classification {.unnumbered}

This is a minimal example script showing how to do image
classification using a set of images that are stored on
the same machine where we are running the models. To start,
we will load in a few modules that will be needed for the
task.

```{python}
from os import listdir
from os.path import splitext, join
from transformers import AutoImageProcessor, ResNetForImageClassification
from PIL import Image

import torch
import polars as pl
```

Next, we load the model that we are interested in using. 
There are a large number of image classification algorithms
on HuggingFace; most can be used exactly the same way by simply
changing the name of the model in the function calls below.

```{python}
image_processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
model = ResNetForImageClassification.from_pretrained("microsoft/resnet-50")
```

With the models loaded, the next step is to build a set of
paths to the images that we are interested in using. Here,
we will create a list of all of the image files in the
directory containing the FSA-OWI images. Then, to save time,
we take just the first five images to work with. We could
run over all of the images or use a different collection by
changing the variables below.

```{python}
collection = 'fsaowi'
num_images = 5

paths = sorted(listdir(join('img', collection)))
paths = [x for x in paths if splitext(x)[1] == '.jpg']
paths = paths[:num_images]
```

Now, we will load each of the images and run the model over
it. For each output, we save the top 10 predictions, with
both their names and probabilties.

```{python}
top_n = 10
output = {'path': [], 'label': [], 'prob': []}
for path in paths:
    image = Image.open(join('img', collection, path))
    image = image.convert('RGB')
    inputs = image_processor(image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
    index = logits.argsort(-1, descending=True)
    index = index.tolist()[0][:top_n]
    label = [model.config.id2label[x] for x in index]
    odds = torch.exp(logits[0])
    prob = odds / (1 + odds)
    prob = sorted(prob.tolist(), reverse=True)[:top_n]
    output['path'] += ([path] * top_n)
    output['label'] += label
    output['prob'] += prob
```

The output is constructed such that we can call the `from_dict`
method from **polars** to construct a data frame. If needed, this
can be saved as a CSV file with the `write_csv` method of the
resulting data frame.

```{python}
dt = pl.from_dict(output)
dt
```

Finally, we can visualize the results using the following code:

```{python}
for path in paths:
    image = Image.open(join('img', collection, path))
    image = image.convert('RGB')
    res = dt.filter(pl.col("path") == path)
    print(
        "Predicted to be a {0:1} with probability {1:.04f}:".format(
        res['label'][0],
        res['prob'][0]
    ))
    display(image)
```