Use new ML model

backend/Dockerfile

@@ -21,13 +21,9 @@ RUN pip --default-timeout=300 install --upgrade pip \
     && rm -r /root/.cache
 
 ARG VERSION
-ARG MODEL="EffB7_2023-03-06_08"
 ENV SSL_CERT_FILE=$CACERT_LOCATION
-RUN curl -o model.pth https://storage.gra.cloud.ovh.net/v1/AUTH_df731a99a3264215b973b3dee70a57af/basegun-models/${MODEL}.pth
 COPY src/ src/
-RUN mkdir -p src/weights \
-    && mv model.pth src/weights/model.pth \
-    && echo '{"app": "'${VERSION}'", "model": "'${MODEL}'"}' > versions.json
+COPY model.pt .
 
 FROM base as dev
 COPY tests/ tests/

backend/requirements.txt

@@ -10,6 +10,7 @@ user-agents==2.2.0
 boto3==1.28.39
 torch==2.1.1
 torchvision==0.16.1
+ultralytics==8.1.2
 # Dev
 pytest==7.4.3
 coverage==7.3.2

backend/src/main.py

@@ -159,7 +159,7 @@ async def add_owasp_middleware(request: Request, call_next):
 logger = setup_logs(PATH_LOGS)
 
 # Load model
-MODEL_PATH = os.path.join(CURRENT_DIR, "weights/model.pth")
+MODEL_PATH = os.path.join(CURRENT_DIR, "../model.pt")
 model = None
 if os.path.exists(MODEL_PATH):
     model = load_model_inference(MODEL_PATH)
@@ -242,9 +242,9 @@ async def imageupload(
         extras_logging["bg_label"] = label
         extras_logging["bg_confidence"] = confidence
         extras_logging["bg_model_time"] = round(time.time() - start, 2)
-        if confidence < 46:
+        if confidence < 0.46:
             extras_logging["bg_confidence_level"] = "low"
-        elif confidence < 76:
+        elif confidence < 0.76:
             extras_logging["bg_confidence_level"] = "medium"
         else:
             extras_logging["bg_confidence_level"] = "high"

backend/src/model.py

@@ -2,10 +2,8 @@
 from typing import Union
 
 import numpy as np
-import torch
-import torchvision.models as Model
 from PIL import Image
-from torchvision import transforms
+from ultralytics import YOLO
 
 CLASSES = [
     "autre_pistolet",
@@ -22,144 +20,21 @@
     "semi_auto_style_militaire_autre",
 ]
 
-MODEL_TORCH = Model.efficientnet_b7
-INPUT_SIZE = 600
-device = torch.device("cpu")
 
-
-class ConvertRgb(object):
-    """Converts an image to RGB"""
-
-    def __init__(self):
-        pass
-
-    def __call__(self, image):
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-        return image
-
-
-class Rescale(object):
-    """Rescale the image in a sample to a given size while keeping ratio
-
-    Args:
-        output_size (int): Desired output size. The largest of image edges is matched
-            to output_size keeping aspect ratio the same.
-    """
-
-    def __init__(self, output_size):
-        assert isinstance(output_size, int)
-        self.output_size = output_size
-
-    def __call__(self, image):
-        w, h = image.size
-        if w > h:
-            new_h, new_w = self.output_size * h / w, self.output_size
-        else:
-            new_h, new_w = self.output_size, self.output_size * w / h
-        new_h, new_w = int(new_h), int(new_w)
-        return transforms.functional.resize(image, (new_h, new_w))
-
-
-class RandomPad(object):
-    """Pad an image to reach a given size
-
-    Args:
-        output_size (int): Desired output size. We pad all edges
-                        symmetrically to reach a size x size image.
-    """
-
-    def __init__(self, output_size):
-        assert isinstance(output_size, int)
-        self.output_size = output_size
-
-    def __call__(self, image):
-        w, h = image.size
-        pads = {
-            "horiz": [self.output_size - w, 0, 0],
-            "vert": [self.output_size - h, 0, 0],
-        }
-        if pads["horiz"][0] >= 0 and pads["vert"][0] >= 0:
-            for direction in ["horiz", "vert"]:
-                pads[direction][1] = pads[direction][0] // 2
-                if (
-                    pads[direction][0] % 2 == 1
-                ):  # if the size to pad is odd, add a random +1 on one side
-                    pads[direction][1] += np.random.randint(0, 1)
-                pads[direction][2] = pads[direction][0] - pads[direction][1]
-
-            return transforms.functional.pad(
-                image,
-                [pads["horiz"][1], pads["vert"][1], pads["horiz"][2], pads["vert"][2]],
-                fill=int(np.random.choice([0, 255])),  # border randomly white or black
-            )
-        else:
-            return image
-
-
-def build_model(model: Model) -> Model:
-    """Create the model structure
-
-    Args:
-        model (Model): raw torchvision model
-
-    Returns:
-        Model: modified model with classification layer size len(CLASSES)
-    """
-    # freeze all layers except classification - very important
-    for param in model.parameters():
-        param.requires_grad = False
-    # replace last layer of model for our number of classes
-    num_ftrs = model.classifier[1].in_features
-    model.classifier[1] = torch.nn.Linear(num_ftrs, len(CLASSES))
-    model = model.to(device)
-    return model
-
-
-def load_model_inference(state_dict_path: str) -> Model:
+def load_model_inference(model_path: str):
     """Load model structure and weights
 
     Args:
-        state_dict_path (str): path to model (.pth file)
+        model_path (str): path to model (.pt file)
 
     Returns:
         Model: loaded model ready for prediction
     """
-    model = build_model(MODEL_TORCH())
-    # Initialize model with the pretrained weights
-    model.load_state_dict(
-        torch.load(state_dict_path, map_location=device)["model_state_dict"]
-    )
-    model.to(device)
-    # set the model to inference mode
-    model.eval()
+    model = YOLO(model_path)
     return model
 
 
-def prepare_input(image: Image) -> torch.Tensor:
-    """Convert a PIL Image to model-compatible input
-
-    Args:
-        image (Image): input image
-
-    Returns:
-        torch.Tensor: converted image
-        (shape (1, 3, size, size), normalized on ImageNet)
-    """
-    loader = transforms.Compose(
-        [
-            ConvertRgb(),
-            Rescale(INPUT_SIZE),
-            RandomPad(INPUT_SIZE),
-            transforms.ToTensor(),
-            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-        ]
-    )
-    image = loader(image).float()
-    return image.unsqueeze(0).to(device)
-
-
-def predict_image(model: Model, img: bytes) -> Union[str, float]:
+def predict_image(model, img: bytes) -> Union[str, float]:
     """Run the model prediction on an image
 
     Args:
@@ -170,9 +45,8 @@ def predict_image(model: Model, img: bytes) -> Union[str, float]:
         Union[str, float]: (label, confidence) of best class predicted
     """
     im = Image.open(BytesIO(img))
-    image = prepare_input(im)
-    output = model(image)
-    probs = torch.nn.functional.softmax(output, dim=1).detach().numpy()[0]
-    res = [(CLASSES[i], round(probs[i] * 100, 2)) for i in range(len(CLASSES))]
-    res.sort(key=lambda x: x[1], reverse=True)
-    return res[0]
+    results = model(im, verbose=False)
+    predicted_class = results[0].probs.top5[0]
+    label = CLASSES[predicted_class]
+    confidence = float(results[0].probs.top5conf[0])
+    return (label, confidence)

backend/tests/test_api.py

@@ -81,7 +81,7 @@ def test_upload(self):
 
         # checks that the json result is as expected
         assert res["label"] == "revolver"
-        assert res["confidence"] == pytest.approx(98.43, 0.1)
+        assert res["confidence"] == pytest.approx(1, 0.1)
         assert res["confidence_level"] == "high"
 
     def test_feedback_and_logs(self):

backend/tests/test_model.py

@@ -1,55 +1,18 @@
 import os
 
-import numpy as np
 import pytest
-from PIL import Image
-from src.model import (
-    CLASSES,
-    INPUT_SIZE,
-    load_model_inference,
-    predict_image,
-    prepare_input,
-)
-from torch import Tensor
+from src.model import CLASSES, load_model_inference, predict_image
 
 
 class TestModel:
-    model_path = os.path.join(
-        os.path.dirname(os.path.abspath(__file__)), "../src/weights/model.pth"
-    )
+    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../model.pt")
     assert os.path.exists(model_path)
     model = load_model_inference(model_path)
 
-    def test_is_efficientnet(self):
-        """Checks that the loaded model is an efficientnet"""
-        assert "efficientnet" in self.model.__class__.__name__.lower()
-        assert len(self.model.features) == 9
-        assert len(self.model.classifier) == 2
-
-    def test_model_correctly_built(self):
-        """Checks that the model is correctly built for prediction"""
-        # check number of classes
-        assert self.model.classifier[1].out_features == len(CLASSES)
-        # check model in "eval" mode
-        assert self.model.training is False
-        # check model on cpu
-        assert next(self.model.parameters()).is_cuda is False
-
-    def test_prepare_input(self):
-        """Checks prepare_input works properly"""
-        # create random RGBA image
-        image = Image.fromarray((np.random.rand(100, 200, 4) * 255).astype("uint8"))
-        image = prepare_input(image)
-        assert type(image) == Tensor
-        # checks converted to 3 channels
-        assert image.size(dim=1) == 3
-        # checks image resized to INPUT_SIZE x INPUT_SIZE
-        assert image.size() == (1, 3, INPUT_SIZE, INPUT_SIZE)
-
     def test_predict_image(self):
         """Checks the prediction of an image by the model"""
         path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "revolver.jpg")
         with open(path, "rb") as f:
             res = predict_image(self.model, f.read())
         assert res[0] == "revolver"
-        assert res[1] == pytest.approx(98.43, 0.1)
+        assert res[1] == pytest.approx(1, 0.1)