model.py

# * SPDX-License-Identifier: Apache-2.0
# * © 2023 ETH Zurich and other contributors, see AUTHORS.txt for details

import os
from typing import List

import torch
from PIL import Image
from fastapi import HTTPException
from mtc_api_utils.base_model import MLBaseModel
from mtc_api_utils.init_api import download_if_not_exists
from torchvision import transforms

from aesthetics_model.api_types import AestheticsModelResponse
from aesthetics_model.config import AestheticsModelConfig
from aesthetics_model.model.vision_transformer import vit_large_patch16_224_in21k


class AestheticModel(MLBaseModel):
 model = None
 transform = None
 device = None
 model_checkpoint_path = None

 def init_model(self):
 print("Initializing Aesthetic Model...")
 # Download image archive
 self.model_checkpoint_path = download_if_not_exists(
 artifact_url=os.path.join(AestheticsModelConfig.model_base_url, AestheticsModelConfig.model_url),
 download_dir=AestheticsModelConfig.model_checkpoint_path,
 polybox_auth=AestheticsModelConfig.polybox_credentials,
 is_tar=False,
 )
 print("Model downloaded, loading to device..")

 self.model, self.transform, self.device = self.load_model()
 print("Model initialized")

 def is_ready(self):
 return self.model is not None

 def inference(self, images: List[Image.Image]) -> List[AestheticsModelResponse]:
 if not self.is_ready():
 raise HTTPException(status_code=503, detail="Model is not ready")

 results = []
 for image in images:
 im = self.validate_and_resize_image(image, max_size=AestheticsModelConfig.max_image_size)
 image_tensor = self.transform(im).unsqueeze(0).to(self.device)
 with torch.no_grad():
 image_embedding = self.model.forward_features(image_tensor)
 score = self.model.head(image_embedding).squeeze().to("cpu")
 result: AestheticsModelResponse = AestheticsModelResponse(
 aesthetics_score=float(score),
 aesthetics_embedding=image_embedding.to("cpu").tolist(),
 # TODO: success is going to be always true
 success=True
 )
 results.append(result)
 return results

 def load_model(self):
 # Check if a gpu and cuda is available
 use_cuda = torch.cuda.is_available()
 device = torch.device("cuda:0" if use_cuda else "cpu")
 print(f"Running on {device}")

 model = vit_large_patch16_224_in21k()
 model.reset_classifier(num_classes=1)
 model.load_state_dict(torch.load(self.model_checkpoint_path, map_location=device))

 model.eval()
 model.to(device)

 transform = transforms.Compose(
 [transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]
 )

 return model, transform, device

 def validate_and_resize_image(self, image: Image.Image, max_size: int):
 """Check if an image is to big, if it is too big resize it. Transform to RGB if B&W image.

 Args:
 img : PIL Image
 max_size : int
 """
 if not image.mode == "RGB":
 image = image.convert(mode="RGB")

 size_factor = (max_size + 0.0) / max(image.size)
 if size_factor < 1:
 new_h = round(image.size[0] * size_factor)
 new_w = round(image.size[1] * size_factor)
 image = image.resize((new_w, new_h), resample=Image.HAMMING)
 return image


if __name__ == '__main__':
 model = AestheticModel()
 model.__wait_until_ready__()

 example_images = ["../tests/test_image.png"] # model.get_examples()

 # Perform inference with underlying model
 images = [model.model(Image.open(image)) for image in example_images]

 model_result = model.inference(images=images)
 print(model_result)