Spaces:

competitions

/

aiornot

like 137

Runtime error

Change paths as required. Data needs to downloaded separately.

import os

import albumentations
import numpy as np
import pandas as pd
import timm
import torch
import torch.nn as nn
from sklearn import metrics
from sklearn.model_selection import train_test_split
from tez import Tez, TezConfig
from tez.callbacks import EarlyStopping
from tez.datasets import ImageDataset

TRAIN_BATCH_SIZE = 8
VALID_BATCH_SIZE = 8
EPOCHS = 10
IMAGE_SIZE = 256


class AIdentifierModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = timm.create_model("resnet18", pretrained=True)
        n_features = self.model.fc.in_features
        self.model.fc = nn.Linear(n_features, 2)

    def monitor_metrics(self, outputs, targets):
        device = targets.device.type
        outputs = torch.argmax(outputs, dim=1).cpu().detach().numpy()
        targets = targets.cpu().detach().numpy()
        f1 = metrics.f1_score(targets, outputs, average="macro")
        return {"f1": torch.tensor(f1, device=device)}

    def optimizer_scheduler(self):
        opt = torch.optim.Adam(self.parameters(), lr=1e-3)
        sch = torch.optim.lr_scheduler.ReduceLROnPlateau(
            opt,
            factor=0.5,
            patience=2,
            verbose=True,
            mode="max",
            threshold=1e-4,
        )
        return opt, sch

    def forward(self, image, targets=None):
        outputs = self.model(image)
        if targets is not None:
            loss = nn.CrossEntropyLoss()(outputs, targets)
            metrics = self.monitor_metrics(outputs, targets)
            return outputs, loss, metrics
        return outputs, 0, {}


if __name__ == "__main__":
    train_aug = albumentations.Compose(
        [
            albumentations.RandomResizedCrop(IMAGE_SIZE, IMAGE_SIZE, p=1.0),
            albumentations.Transpose(p=0.5),
            albumentations.HorizontalFlip(p=0.5),
            albumentations.VerticalFlip(p=0.5),
            albumentations.ShiftScaleRotate(p=0.5),
            albumentations.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),
            albumentations.RandomBrightnessContrast(brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1), p=0.5),
            albumentations.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225],
                max_pixel_value=255.0,
                p=1.0,
            ),
        ],
        p=1.0,
    )

    valid_aug = albumentations.Compose(
        [
            albumentations.Resize(IMAGE_SIZE, IMAGE_SIZE, p=1.0),
            albumentations.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225],
                max_pixel_value=255.0,
                p=1.0,
            ),
        ],
        p=1.0,
    )

    test_aug = albumentations.Compose(
        [
            albumentations.Resize(IMAGE_SIZE, IMAGE_SIZE, p=1.0),
            albumentations.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225],
                max_pixel_value=255.0,
                p=1.0,
            ),
        ],
        p=1.0,
    )

    df = pd.read_csv("train.csv")
    train_df, valid_df = train_test_split(
        df,
        test_size=0.1,
        random_state=42,
        stratify=df.label.values,
    )
    train_path = "train"
    train_image_paths = train_df.id.values.tolist()
    train_image_paths = [os.path.join(train_path, i) for i in train_image_paths]
    train_targets = train_df.label.values

    valid_image_paths = valid_df.id.values.tolist()
    valid_image_paths = [os.path.join(train_path, i) for i in valid_image_paths]
    valid_targets = valid_df.label.values

    df_test = pd.read_csv("sample_submission.csv")
    test_path = "test"
    test_image_paths = df_test.id.values.tolist()
    test_image_paths = [os.path.join(test_path, i) for i in test_image_paths]
    test_targets = df_test.label.values
    test_ids = df_test.id.values

    train_dataset = ImageDataset(
        image_paths=train_image_paths,
        targets=train_targets,
        augmentations=train_aug,
    )

    valid_dataset = ImageDataset(
        image_paths=valid_image_paths,
        targets=valid_targets,
        augmentations=valid_aug,
    )

    test_dataset = ImageDataset(
        image_paths=test_image_paths,
        targets=test_targets,
        augmentations=valid_aug,
    )
    es = EarlyStopping(
        monitor="valid_f1",
        model_path="model.bin",
        patience=7,
        mode="max",
        save_weights_only=True,
    )

    model = AIdentifierModel()
    model = Tez(model)
    config = TezConfig(
        training_batch_size=TRAIN_BATCH_SIZE,
        validation_batch_size=VALID_BATCH_SIZE,
        epochs=EPOCHS,
        step_scheduler_after="epoch",
        step_scheduler_metric="valid_f1",
    )
    model.fit(
        train_dataset,
        valid_dataset=valid_dataset,
        config=config,
        callbacks=[es],
    )

    model.load("model.bin", weights_only=True)

    preds_iter = model.predict(test_dataset)
    final_preds = []

    for preds in preds_iter:
        final_preds.append(preds)

    final_preds = np.vstack(final_preds)
    final_preds = final_preds[: len(test_dataset), 1]

    df = pd.DataFrame(
        {
            "id": test_ids,
            "label": final_preds,
        }
    )
    df.to_csv("submission.csv", index=False)