|
import numpy as np |
|
import pandas as pd |
|
import PIL |
|
import pytest |
|
import torch |
|
from pytorch_lightning import Trainer |
|
|
|
from .data import ImageDataset |
|
from .data_lightning import InferenceDataModel |
|
from .model_lightning import LinearClassifier as LightningLinearClassifier |
|
from .submission import make_submission |
|
|
|
|
|
class TestingInferenceDataModel(InferenceDataModel): |
|
def train_dataloader(self): |
|
for batch in self.predict_dataloader(): |
|
|
|
batch["label"] = torch.randint(0, 10, (batch["features"].shape[0],)) |
|
yield batch |
|
|
|
|
|
@pytest.fixture |
|
def images_root(tmp_path): |
|
images_root = tmp_path / "images" |
|
images_root.mkdir() |
|
for i in range(10): |
|
img = PIL.Image.fromarray( |
|
np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) |
|
) |
|
img.save(images_root / f"{i}.jpg") |
|
return images_root |
|
|
|
|
|
@pytest.fixture |
|
def metadata(tmp_path, images_root): |
|
res = [] |
|
for i, img in enumerate(images_root.glob("*.jpg")): |
|
res.append({"filename": img.name, "observation_id": i}) |
|
df = pd.DataFrame(res) |
|
df.to_csv(tmp_path / "metadata.csv", index=False) |
|
return tmp_path / "metadata.csv" |
|
|
|
|
|
@pytest.fixture |
|
def model_checkpoint(tmp_path, metadata, images_root): |
|
model_checkpoint = tmp_path / "model.ckpt" |
|
model = LightningLinearClassifier(768, 10) |
|
trainer = Trainer(max_epochs=1, fast_dev_run=True) |
|
dm = TestingInferenceDataModel(metadata, images_root) |
|
trainer.fit(model, dm) |
|
trainer.save_checkpoint(model_checkpoint) |
|
return model_checkpoint |
|
|
|
|
|
def test_image_dataset(images_root, metadata): |
|
dataset = ImageDataset(metadata, images_root) |
|
assert len(dataset) == 10 |
|
for i in range(10): |
|
assert dataset[i]["features"].shape == torch.Size([3, 100, 100]) |
|
|
|
|
|
def test_inference_datamodel(images_root, metadata): |
|
batch_size = 5 |
|
model = InferenceDataModel(metadata, images_root, batch_size=batch_size) |
|
model.setup() |
|
assert len(model.dataloader) == 2 |
|
for batch in model.predict_dataloader(): |
|
assert set(batch.keys()) == {"features", "observation_id"} |
|
assert batch["features"].shape == torch.Size([batch_size, 768]) |
|
|
|
|
|
def test_model_checkpoint(model_checkpoint): |
|
model = LightningLinearClassifier.load_from_checkpoint(model_checkpoint) |
|
assert model |
|
|
|
|
|
def test_make_submission(model_checkpoint, metadata, images_root, tmp_path): |
|
output_csv_path = tmp_path / "submission.csv" |
|
make_submission(metadata, model_checkpoint, output_csv_path, images_root) |
|
submission_df = pd.read_csv(output_csv_path) |
|
assert len(submission_df) == 10 |
|
assert set(submission_df.columns) == {"observation_id", "class_id"} |
|
assert submission_df["class_id"].isin(range(10)).all() |
|
|
