Merge pull request #16 from Kawaeee/v1.3

Refactor v1.3

.dockerignore

@@ -0,0 +1,6 @@
+*.h5
+.DS_Store
+__pycache__
+datasets/
+.python-version
+*.git

.gitignore

@@ -1,4 +1,5 @@
 *.h5
 .DS_Store
 __pycache__
-datasets/
+datasets/
+.python-version

Dockerfile

@@ -1,41 +1,19 @@
-FROM ubuntu:bionic
-LABEL maintainer="[email protected]"
+FROM python:3.9
 
 ENV LANG=C.UTF-8
-ENV DEBIAN_FRONTEND=noninteractive
+ENV TZ=Asia/Bangkok
+RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
 
-# Update apt package
-RUN apt update --fix-missing
+# Allow statements and log messages to immediately appear in the logs
+ENV PYTHONUNBUFFERED=1
 
-# Install required dependencies by default
-RUN apt install -y wget curl git htop nano
+WORKDIR /app
 
-# miniconda3 - Python 3.7
-WORKDIR /opt/
-ARG HOME="/opt"
-RUN wget https://repo.anaconda.com/miniconda/Miniconda3-py37_4.10.3-Linux-x86_64.sh \
-    && bash Miniconda3-py37_4.10.3-Linux-x86_64.sh -b \
-    && rm -r Miniconda3-py37_4.10.3-Linux-x86_64.sh
-ENV PATH="/opt/miniconda3/bin:${PATH}"
+COPY requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
-# butt_or_bread repository
-RUN git clone https://github.com/Kawaeee/butt_or_bread.git
+COPY . .
 
-# Download model
-RUN cd /opt/butt_or_bread/
-RUN wget https://github.com/Kawaeee/butt_or_bread/releases/download/v1.2/buttbread_resnet152_3.h5
+RUN wget https://github.com/Kawaeee/butt_or_bread/releases/download/v1.3/buttbread_resnet152_3.h5
 
-# Install python packages
-RUN pip install --upgrade pip
-RUN pip install -r /opt/butt_or_bread/requirements.txt --no-cache-dir
-
-# House-keeping
-RUN conda clean -a -y
-RUN pip cache purge
-RUN rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
-RUN apt autoclean
-RUN apt autoremove
-
-ENV HOME="/root"
-
-ENTRYPOINT ["streamlit", "run", "/opt/butt_or_bread/streamlit_app.py"]
+ENTRYPOINT ["streamlit", "run", "/app/streamlit_app.py"]

README.md

@@ -1,6 +1,6 @@
 # Corgi butt or loaf of bread?
-[![GitHub Release](https://img.shields.io/github/v/release/Kawaeee/butt_or_bread)](https://github.com/Kawaeee/butt_or_bread/releases/tag/v1.2)
-![Implemented in](https://upload.wikimedia.org/wikipedia/commons/f/fc/Blue_Python_3.7_Shield_Badge.svg)
+[![GitHub Release](https://img.shields.io/github/v/release/Kawaeee/butt_or_bread)](https://github.com/Kawaeee/butt_or_bread/releases/tag/v1.3)
+![Implemented in](https://upload.wikimedia.org/wikipedia/commons/1/1b/Blue_Python_3.9_Shield_Badge.svg)
 ![Visitor Badge](https://visitor-badge.glitch.me/badge?page_id=Kawaeee.butt_or_bread.visitor-badge)
 [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io/kawaeee/butt_or_bread/)
 
@@ -41,7 +41,7 @@
 
 * We already know that in order to benchmark our model performance, we can't just use `accuracy` and `validation_loss` value as the only acceptable metrics.
 
-#### You can download our model weight here: [v1.2](https://github.com/Kawaeee/butt_or_bread/releases/download/v1.2/buttbread_resnet152_3.h5)
+#### You can download our model weight here: [v1.3](https://github.com/Kawaeee/butt_or_bread/releases/download/v1.3/buttbread_resnet152_3.h5)
 
 ## Hyperparameters and configurations
 
@@ -55,7 +55,7 @@
  * To reproduce the model, requires our datasets. You can send me an e-mail at `[email protected]` if you are interested.
 
  - Initial datasets/ directory structure
-   ```Bash
+   ```bash
       └───datasets/
       │     butt/
       │     bread/
@@ -72,7 +72,7 @@
    ```
 
 -  Ready-to-go datasets/ directory structure
-   ```Bash
+   ```bash
       └───datasets/
       │   │
       │   └───train
@@ -94,12 +94,12 @@
    ```
 
  - Install dependencies
-   ```Bash
+   ```bash
    pip install -r requirements.txt
    ```
 
  - Run the `train.py` python script
-   ```Bash
+   ```bash
    python train.py --dataset-path datasets/ --model-path buttbread_resnet152_3.h5
    ```
 
@@ -113,25 +113,25 @@
    ```
 
  - Install dependencies
-   ```Bash
+   ```bash
    pip install -r requirements.txt
    ```
 
  - Run the streamlit
-   ```Bash
+   ```bash
    streamlit run streamlit_app.py
    ```
 
- - Streamlit web application will be hosted on http://localhost:8501
+ > Streamlit web application will be hosted on http://localhost:8501
 
  ## Streamlit Docker Reproduction
-
- - Build Docker image from Dockerfile
-   ```Bash
-   docker build -t butt_or_bread -f Dockerfile .
-   ```
-
- - Run Docker with exposed port 8501
-   ```Bash
-   docker run -p 8501:8501 butt_or_bread
-   ```
+ - Following instructions below
+  ```bash
+    # Directly build and run
+    docker build -t butt-bread-image .
+    docker run --rm --name=butt-bread-container -p 0.0.0.0:8501:8501 butt-bread-image
+
+    # Serve with docker compose
+    docker-compose build
+    docker-compose up
+  ```

butt_or_bread/core.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import requests
 
@@ -6,8 +7,26 @@
 
 
 class ButtBreadClassifier:
+    """
+    A classifier for predicting whether an image is of a corgi butt or a loaf of bread.
+
+    Attributes:
+        model_url (str): The URL where the trained model weights can be downloaded.
+        model_name (str): The filename of the trained model weights.
+        device (torch.device): The device (CPU or GPU) where the model will be loaded.
+        model (torchvision.models.ResNet): The ResNet152 model with a custom fully connected layer for classifying images.
+        preprocessor (torchvision.transforms.Compose): The image preprocessing pipeline used to normalize and resize images.
+        logger (logging.Logger): The logger object used for logging messages during model download.
+
+    Methods:
+        initialize(): Initializes the model by loading the pretrained ResNet152 model, adding a custom fully connected layer with two output classes, loading the model weights, and putting the model on the CPU or GPU depending on availability. Returns the initialized model.
+        download(): Downloads the model weights from the specified URL if the weights are not already present in the server. If the download is successful, returns True, otherwise raises an error.
+        predict(image): Takes an image as input, preprocesses it using standard image transformations, runs the preprocessed image through the loaded model, and returns a JSON object containing the image metadata, the predicted labels (corgi butt or loaf of bread), and the corresponding probabilities. Returns the JSON object.
+    """
+
     def __init__(self, model_url):
         self.model_url = model_url
+        self.model_name = "buttbread_resnet152_3.h5"
         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         self.model = None
         self.preprocessor = transforms.Compose(
@@ -20,37 +39,52 @@ def __init__(self, model_url):
                 ),
             ]
         )
+        self.logger = logging.getLogger(__name__)
 
     def initialize(self):
-        """Retrieves the butt_bread trained model and maps it to the CPU by default, can also specify GPU here."""
-
-        self.model = models.resnet152(pretrained=False).to(self.device)
+        """
+        Initializes the model by loading the pretrained ResNet152 model, adding a custom fully connected layer with two output classes, loading the model weights, and putting the model on the CPU or GPU depending on availability. Returns the initialized model.
+        """
+        self.model = models.resnet152(weights=None).to(self.device)
         self.model.fc = torch.nn.Sequential(
             torch.nn.Linear(2048, 128),
             torch.nn.ReLU(inplace=True),
             torch.nn.Linear(128, 2),
         ).to(self.device)
 
-        self.model.load_state_dict(torch.load("buttbread_resnet152_3.h5", map_location=self.device))
+        self.model.load_state_dict(torch.load(self.model_name, map_location=self.device))
         self.model.eval()
 
         return self.model
 
     def download(self):
-        """Download model weight, if model does not exist in Streamlit server."""
-
-        if os.path.isfile("buttbread_resnet152_3.h5") is False:
-            print("Downloading butt_bread model !!")
-            req = requests.get(self.model_url, allow_redirects=True)
-            open("buttbread_resnet152_3.h5", "wb").write(req.content)
-            req = None
-            return True
+        """
+        Downloads the model weights from the specified URL if the weights are not already present in the server. If the download is successful, returns True, otherwise raises an error.
+        """
+        try:
+            if not os.path.isfile(self.model_name):
+                self.logger.info("Downloading butt_or_bread model !!")
+                model_weights_request = requests.get(self.model_url, allow_redirects=True)
+                with open(self.model_name, "wb") as f:
+                    f.write(model_weights_request.content)
+                return True
+
+        except Exception as e:
+            self.logger.error(f"Failed to download model. Error: {e}")
+            raise
 
         return False
 
     def predict(self, image):
-        """Make a prediction on a single image"""
+        """
+        Takes an image as input, preprocesses it using standard image transformations, runs the preprocessed image through the loaded model, and returns a JSON object containing the image metadata, the predicted labels (corgi butt or loaf of bread), and the corresponding probabilities. Returns the JSON object.
+
+        Args:
+            image (PIL.Image): The input image to make a prediction on.
 
+        Returns:
+            dict: A JSON object containing the image metadata, the predicted labels (corgi butt or loaf of bread), and the corresponding probabilities.
+        """
         input_image = self.preprocessor(image)
         input_image = input_image.unsqueeze(0)
 

butt_or_bread/utils.py

@@ -1,17 +1,23 @@
 import psutil
 
 
-def health_check():
-    """Check CPU/Memory/Disk usage of deployed machine"""
+def health_check() -> str:
+    """
+    Check the CPU, memory, and disk usage of the deployed machine.
 
+    Returns:
+        A string with information about the current usage levels for each resource, formatted as follows:
+        "CPU Usage: [percent]%" | "Memory usage: [used memory]GB/[total memory]GB" | "Disk usage: [used disk]GB/[total disk]GB"
+
+    Uses the `psutil` module to obtain information about the system resources.
+    """
+    vm = psutil.virtual_memory()
+    du = psutil.disk_usage("/")
     cpu_percent = psutil.cpu_percent(0.15)
-    total_memory = psutil.virtual_memory().total / float(1 << 30)
-    used_memory = psutil.virtual_memory().used / float(1 << 30)
-    total_disk = psutil.disk_usage("/").total / float(1 << 30)
-    used_disk = psutil.disk_usage("/").used / float(1 << 30)
 
-    cpu_usage = f"CPU Usage: {cpu_percent:.2f}%"
-    memory_usage = f"Memory usage: {used_memory:,.2f}G/{total_memory:,.2f}G"
-    disk_usage = f"Disk usage: {used_disk:,.2f}G/{total_disk:,.2f}G"
+    total_memory = vm.total / 1024**3
+    used_memory = vm.used / 1024**3
+    total_disk = du.total / 1024**3
+    used_disk = du.used / 1024**3
 
-    return " | ".join([cpu_usage, memory_usage, disk_usage])
+    return f"CPU Usage: {cpu_percent:.2f}% | Memory usage: {used_memory:.2f}GB/{total_memory:.2f}GB | Disk usage: {used_disk:.2f}GB/{total_disk:.2f}GB"

docker-compose.yml

@@ -0,0 +1,6 @@
+version: "3.9"
+services:
+  container:
+    build: .
+    ports:
+      - "8501:8501"

requirements.txt

@@ -1,13 +1,7 @@
-ImageHash==4.2.0
-matplotlib==3.4.1
-numpy==1.22.0
-pillow==9.3.0
-psutil==5.9.0
-requests==2.27.1
-streamlit==1.11.1
-# [STREAMLIT-DEPLOYMENT-GUIDE] Add "+cpu" on torch and torchvision to get CPU inference support
-# NOTE: On macOS system, it might be unable to install with +cpu flag
---find-links https://download.pytorch.org/whl/torch_stable.html
-torch==1.13.1
-torchvision==0.9.1+cpu
-tqdm==4.60.0
+streamlit==1.21.0
+psutil==5.9.5
+requests==2.28.2
+torch==2.0.0
+torchvision==0.15.1
+Pillow==9.5.0
+tqdm==4.65.0

streamlit_app.json

@@ -4,7 +4,7 @@
       "icon":"assets/icon/corgi-icon.png",
       "markdown":{
          "star":"[![GitHub Star](https://img.shields.io/github/stars/Kawaeee/butt_or_bread)](https://github.com/Kawaeee/butt_or_bread)",
-         "release":"[![GitHub Release](https://img.shields.io/github/v/release/Kawaeee/butt_or_bread)](https://github.com/Kawaeee/butt_or_bread/releases/tag/v1.2)",
+         "release":"[![GitHub Release](https://img.shields.io/github/v/release/Kawaeee/butt_or_bread)](https://github.com/Kawaeee/butt_or_bread/releases/tag/v1.3)",
          "visitor":"![Visitor Badge](https://visitor-badge.glitch.me/badge?page_id=Kawaeee.butt_or_bread.visitor-badge)"
       },
       "mode":{
@@ -20,7 +20,7 @@
       }
    },
    "model":{
-      "url":"https://github.com/Kawaeee/butt_or_bread/releases/download/v1.2/buttbread_resnet152_3.h5",
+      "url":"https://github.com/Kawaeee/butt_or_bread/releases/download/v1.3/buttbread_resnet152_3.h5",
       "label":{
          "corgi":"Corgi butt 🐕",
          "bread":"Loaf of bread 🍞"