Initial commit

.dockerignore

@@ -0,0 +1,9 @@
+.git
+.gitattributes
+.coverage
+.coverage.*
+.env
+.venv
+.aws
+*.log
+*.md

Dockerfile

@@ -0,0 +1,13 @@
+ARG APP_PORT
+FROM python:3.10
+
+# The two following lines are requirements for the Dev Mode to be functional
+# Learn more about the Dev Mode at https://huggingface.co/dev-mode-explorers
+RUN useradd -m -u 1000 user
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user ./src /app
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", $APP_PORT]

requirements.txt

@@ -0,0 +1,8 @@
+accelerate==1.8.1
+fastapi==0.116.0
+pillow==11.3.0
+python-multipart==0.0.20
+torch==2.7.1
+torchvision==0.22.1
+transformers==4.53.1
+uvicorn==0.35.0

src/backup_services.py

@@ -0,0 +1,91 @@
+import logging
+
+from pathlib import Path
+
+import torch
+
+from fastapi import HTTPException, status
+from PIL import Image
+from torchvision import models
+from typing import Tuple
+
+import src.config as config
+
+
+logger = logging.getLogger(__name__)
+
+
+async def classify_mushroom_in_image_svc(img: Image.Image) -> Tuple[str, str, str]:
+    """Service used to classify a mushroom shown in an image.
+    The mushroom is classified to one of many well known mushroom classes/types,
+    as well as according to its toxicity profile (i.e. edible or poisonous).
+    Additionally, a probability is returned showing confidence of classification.
+
+    :param img: the image of the mushroom to be classified
+    :type img: Image.Image
+    :return: mushroom_type, toxicity_profile, classification_confidence
+    :rtype: Tuple[str, str, str]
+    """
+
+    try:
+        # Device agnostic
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        logger.debug("Loading classification model.")
+
+        model_path = config.MODEL_PATH
+
+        # Load saved model checkpoint
+        model_state_dict = torch.load(model_path, map_location=device)
+
+        # Get class_names from saved model checkpoint
+        model_dirname = Path(model_path).resolve().parent
+        with open(model_dirname / "labels.txt", "r") as labels_fp:
+            class_names = [line.strip() for line in labels_fp]
+
+        model = models.get_model(config.BASE_MODEL_NAME, num_classes=len(class_names))
+
+        # Load state_dict of saved model
+        model.load_state_dict(model_state_dict)
+
+        weights_enum = models.get_model_weights(config.BASE_MODEL_NAME)
+
+        # Get the model's default transforms
+        image_transform = weights_enum.DEFAULT.transforms()
+
+        # Make sure the model is on the target device
+        model.to(device)
+
+        # Turn on model evaluation mode and inference mode
+        model.eval()
+        with torch.inference_mode():
+            logger.debug("Adapting input image by applying necessary transforms!")
+            # Transform and add an extra dimension to image (model requires samples in [batch_size, color_channels, height, width])
+            transformed_image = image_transform(img).unsqueeze(dim=0)
+
+            # Make a prediction on image with an extra dimension and send it to the target device
+            target_image_pred = model(transformed_image.to(device))
+
+        logger.debug("Starting classification process...")
+        # Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
+        target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
+
+        # Convert prediction probabilities -> prediction labels
+        target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
+
+        class_name = class_names[target_image_pred_label]
+
+        # Split class_name to mushroom type and toxicity profile
+        class_type, toxicity = class_name.rsplit("_", 1)
+
+        # 4 decimal points precision
+        prob = round(target_image_pred_probs.max().item(), 4)
+
+        return class_type, toxicity, prob
+
+    except Exception as e:
+        logger.error("Classification process error: {e}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Classification process failed due to an internal error. Contact support if this persists.",
+        )

src/config.py

@@ -0,0 +1 @@
+MODEL_ID = "blasisd/musheff"

src/dependencies.py

@@ -0,0 +1,15 @@
+from fastapi import Request, HTTPException
+from transformers import PreTrainedModel
+from transformers.image_processing_utils import BaseImageProcessor
+
+
+def get_model(request: Request) -> PreTrainedModel:
+    if not hasattr(request.app.state, "model"):
+        raise HTTPException(status_code=500, detail="Model not loaded")
+    return request.app.state.model
+
+
+def get_preprocessor(request: Request) -> BaseImageProcessor:
+    if not hasattr(request.app.state, "preprocessor"):
+        raise HTTPException(status_code=500, detail="Preprocessor not loaded")
+    return request.app.state.preprocessor

src/main.py

@@ -0,0 +1,47 @@
+from contextlib import asynccontextmanager
+from fastapi import FastAPI
+from transformers import (
+    AutoImageProcessor,
+    AutoModel,
+)
+
+import src.config as config
+
+from src.router import router
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Load models during startup
+
+    app.state.model = AutoModel.from_pretrained(
+        config.MODEL_ID,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,  # Activates memory-efficient loading
+        device_map="auto",  # Distributes layers across devices
+    )
+
+    app.state.preprocessor = AutoImageProcessor.from_pretrained(
+        config.MODEL_ID,
+        trust_remote_code=True,
+        use_fast=True,
+    )
+
+    yield
+
+    # Cleanup during shutdown (e.g., GPU memory)
+    del app.state.model
+    del app.state.preprocessor
+
+
+app = FastAPI(
+    description="Mushrooms Classification API", version="0.1.0", lifespan=lifespan
+)
+
+
+@app.get("/")
+async def root():
+    return {"message": "Welcome to Mushrooms Classification API 🍄"}
+
+
+app.include_router(router)

src/router.py

@@ -0,0 +1,59 @@
+import io
+import logging
+
+from fastapi import APIRouter, HTTPException, status, UploadFile, Depends
+from PIL import Image
+
+from src.dependencies import get_model, get_preprocessor
+from src.schema import MushroomClassification
+from src.services import classify_mushroom_in_image_svc
+
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+
+@router.post(
+    "/classify", response_model=MushroomClassification, status_code=status.HTTP_200_OK
+)
+async def classify_mushroom_in_image(
+    image_file: UploadFile,
+    model=Depends(get_model),
+    preprocessor=Depends(get_preprocessor),
+):
+    """Open uploaded image file and call mushroom classification
+    service.
+
+    :param image_file: the uploaded image file
+    :type image_file: UploadFile
+    :param model: the pretrained model, defaults to Depends(get_model)
+    :type model: PreTrainedModel, optional
+    :param preprocessor: the preprocessor for image input transforms, defaults to Depends(get_preprocessor)
+    :type preprocessor: BaseImageProcessor, optional
+    :raises HTTPException: Internal Server Error in case of model/preprocessor loading failure or some uknown error,
+    or Bad Request Error in case of corrupted or invalid uploaded file
+    :return: mushroom_type, toxicity_profile, classification_confidence
+    :rtype: MushroomClassification
+    """
+    logger.info(f"Classify image: {image_file.filename}")
+
+    try:
+        request_object_content = await image_file.read()
+        img = Image.open(io.BytesIO(request_object_content))
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+    except Exception as e:
+        logger.error(f"Error reading file: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Failed to read uploaded file. The file may be corrupted or invalid.",
+        )
+
+    # Get class_name predicted and prediction probability
+    class_name, toxicity, confidence = await classify_mushroom_in_image_svc(
+        img, model, preprocessor
+    )
+    return MushroomClassification(
+        mushroom_type=class_name, toxicity_profile=toxicity, confidence=confidence
+    )

src/schema.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+
+class MushroomClassification(BaseModel):
+    mushroom_type: str
+    toxicity_profile: str
+    confidence: float

src/services.py

@@ -0,0 +1,72 @@
+import logging
+
+from typing import Tuple
+
+import torch
+
+from fastapi import HTTPException, status
+from PIL import Image
+from transformers import PreTrainedModel
+from transformers.image_processing_utils import BaseImageProcessor
+
+
+logger = logging.getLogger(__name__)
+
+
+async def classify_mushroom_in_image_svc(
+    img: Image.Image, model: PreTrainedModel, preprocessor: BaseImageProcessor
+) -> Tuple[str, str, str]:
+    """Service used to classify a mushroom shown in an image.
+    The mushroom is classified to one of many well known mushroom classes/types,
+    as well as according to its toxicity profile (i.e. edible or poisonous).
+    Additionally, a probability is returned showing confidence of classification.
+
+    :param img: the input image of the mushroom to be classified
+    :type img: Image.Image
+    :param model: the pretrained model
+    :type model: PretrainedModel
+    :param preprocessor: the auto preprocessor for image transforms (rescales, crops, normalizations etc.)
+    :type preprocessor: BaseImageProcessor
+    :raises HTTPException: Internal Server Error
+    :return: mushroom_type, toxicity_profile, classification_confidence
+    :rtype: Tuple[str, str, float]
+    """
+
+    try:
+
+        logger.debug("Loading classification model.")
+
+        inputs = preprocessor(img, return_tensors="pt").to(model.device)
+
+        # Turn on model evaluation mode and inference mode
+        model.eval()
+        with torch.inference_mode():
+            logger.debug("Starting classification process...")
+
+            # Make a prediction on image with an extra dimension and send it to the target device
+            target_image_pred = model(inputs["pixel_values"])
+
+        # Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
+        target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
+
+        # model predicts one of the 12 potential mushroom classes
+        predicted_label = target_image_pred.argmax(dim=1).item()
+
+        # Get the label/class name of the prediction made using id2label
+        class_name = model.config.id2label[predicted_label]
+
+        # Split class_name to mushroom type and toxicity profile
+        class_type, toxicity = class_name.rsplit("_", 1)
+
+        # 4 decimal points precision
+        prob = round(target_image_pred_probs.max().item(), 4)
+
+        logger.debug("Finished classification process...")
+        return class_type, toxicity, prob
+
+    except Exception as e:
+        logger.error(f"Classification process error: {e}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Classification process failed due to an internal error. Contact support if this persists.",
+        )