|
|
from typing import List |
|
|
from urllib.request import urlopen |
|
|
from PIL import Image |
|
|
from .data.model_data import ModelData |
|
|
from .models.mobilenet_v3 import MobilenetV3 |
|
|
from .models.clip_vit import ClipVit |
|
|
from .models.google_vit import GoogleVit |
|
|
from .models.resnet_50 import Resnet50 |
|
|
|
|
|
from .data.classification_result import ClassificationResult |
|
|
|
|
|
class ClassificationModel: |
|
|
""" |
|
|
Base class for all classification models. |
|
|
""" |
|
|
|
|
|
def __init__(self): |
|
|
self.load_model() |
|
|
|
|
|
def get_model_names(self): |
|
|
return [model.name for model in self.models] |
|
|
|
|
|
def get_model_data(self, model_name): |
|
|
for model in self.models: |
|
|
if model.name == model_name: |
|
|
return model |
|
|
raise Exception(f'Model {model_name} not found') |
|
|
|
|
|
def load_model(self): |
|
|
self.models = [ |
|
|
ModelData('clip-vit-base-patch32', model_class=ClipVit()), |
|
|
ModelData('mobilenet_v3', model_class=MobilenetV3()), |
|
|
ModelData('google-vit-base-patch16-224', model_class=GoogleVit()), |
|
|
ModelData('microsoft/resnet-50', model_class=Resnet50()) |
|
|
] |
|
|
|
|
|
def classify(self, model_name, image) -> List[ClassificationResult]: |
|
|
|
|
|
print('>> image type -->',type(image)) |
|
|
|
|
|
|
|
|
img = self.image_to_pil(image) |
|
|
|
|
|
model = self.get_model_data(model_name) |
|
|
return model.model_class.classify_image(img) |
|
|
|
|
|
def image_to_pil(self, image): |
|
|
|
|
|
if image.startswith('https'): |
|
|
return Image.open(urlopen(image)) |
|
|
return Image.open(image) |
|
|
|
