Stop versioning the model checkpoints, now they are downloaded from huggingface. Add env vars

.env.example

@@ -0,0 +1,4 @@
+CLAS_FILENAME = "classifier_model_name.pkl"
+DET_FILENAME = "detector_model_name.pth.tar"
+HF_DET_REPO_NAME = "org/repo_name"
+HF_CLAS_REPO_NAME = "org/repo_name"

.gitignore

@@ -3,3 +3,7 @@ __pycache__
 *.jpg
 *.png
 *.jpeg
+*.tar
+*.pkl
+.env
+*.pth

README.md

@@ -32,6 +32,17 @@ python -m venv venv-waste-classifier
 source venv-waste-classifier/bin/activate  # On Windows, use 'venv-waste-classifier\Scripts\activate'
 pip install -r requirements.txt
 ```
+Create a `.env` file and set the following properties:
+
+- CLAS_FILENAME --> The name of the waste classificator model checkpoint
+- DET_FILENAME --> The name of the waste detector model checkpoint
+- HF_DET_REPO_NAME --> The huggingface repository name of the detector model
+  - For example: [rootstrap-org/waste-detector](https://huggingface.co/rootstrap-org/waste-detector)
+- HF_CLAS_REPO_NAME --> The huggingface repository name of the classifier model
+  - For example: [rootstrap-org/waste-classifier](https://huggingface.co/rootstrap-org/waste-classifier)
+
+You can ommit setting the last two properties if you download the models manually and put them under the `models` directory.
+
 
 ### Running the App
 

app.py

@@ -4,14 +4,12 @@ Streamlit app
 import sys
 
 import streamlit as st
+from PIL import Image
 
-from constants import (CLAS_FILEPATH, CLAS_THRESHOLD, CLASSES, DET_FILEPATH,
-                       DET_NAME, DET_THRESHOLD, DEVICE, OUTPUT_IMG_FILEPATH)
+from constants import CLASSES, OUTPUT_IMG_FILEPATH
 
 sys.path.append("./efficientdet")
 
-from PIL import Image
-
 from efficientdet.efficientdet import plot_results
 from trash_detector import detect_trash
 
@@ -57,15 +55,7 @@ def render():
             with col2:
                 with st.spinner(text="Classifying the trash..."):
                     img = Image.open(uploaded_file).convert("RGB")
-                    cls_prob, bboxes_final = detect_trash(
-                        img,
-                        DET_NAME,
-                        DET_FILEPATH,
-                        CLAS_FILEPATH,
-                        DEVICE,
-                        DET_THRESHOLD,
-                        CLAS_THRESHOLD,
-                    )
+                    cls_prob, bboxes_final = detect_trash(img)
                     # plot and save demo image
                     plot_results(
                         img, cls_prob, bboxes_final, CLASSES, OUTPUT_IMG_FILEPATH

constants.py

@@ -1,8 +1,24 @@
-CLAS_FILEPATH = "models/resnet50-classifier.pkl"
-DET_FILEPATH = "models/efficientdet-d2-detector.pth.tar"
+import os
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Model checkpoints and repo names
+CLAS_FILENAME = os.getenv("CLAS_FILENAME")
+DET_FILENAME = os.getenv("DET_FILENAME")
+HF_CLAS_REPO_NAME = os.getenv("HF_CLAS_REPO_NAME")
+HF_DET_REPO_NAME = os.getenv("HF_DET_REPO_NAME")
+
+# Models paths
+MODELS_PATH = "models"
+CLAS_FILEPATH = f"{MODELS_PATH}/{CLAS_FILENAME}"
+DET_FILEPATH = f"{MODELS_PATH}/{DET_FILENAME}"
+
+# Other constants
 CLASSES = ["cardboard", "compost", "glass", "metal", "paper", "plastic", "trash"]
-DET_NAME = "tf_efficientdet_d2"
 CLAS_THRESHOLD = 0.5
+DET_NAME = "tf_efficientdet_d2"
 DET_THRESHOLD = 0.17
 DEVICE = "cpu"
 OUTPUT_IMG_FILEPATH = "classified_image.jpg"

models/efficientdet-d2-detector.pth.tar

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:499a3f0c75e13669d69be25854e980812e2f6b50e618ba2b2e90b25f193e7fd9
-size 97791163

models/resnet50-classifier.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8d2c0667090f996cbe4bab8585300528b8896071e70b1edfdbe671015a074e85
-size 102980821

requirements.txt

@@ -2,6 +2,7 @@ albumentations>=0.5.2
 efficientnet_pytorch
 fastai==2.7.13
 funcy==1.15
+huggingface_hub
 iterative-stratification==0.1.6
 matplotlib==3.8.2
 numpy==1.26.2
@@ -9,6 +10,7 @@ omegaconf>=2.0
 opencv-python==4.8.1.78
 opencv-python-headless==4.8.1.78
 pycocotools>=2.0.0
+python-dotenv
 pytorch_lightning
 pyyaml
 rembg==2.0.53

trash_detector.py

@@ -1,29 +1,41 @@
+import os
+
 import numpy as np
 import torch
 from fastai.vision.all import load_learner
+from huggingface_hub import hf_hub_download
 
+from constants import (CLAS_FILENAME, CLAS_FILEPATH, CLAS_THRESHOLD,
+                       DET_FILENAME, DET_FILEPATH, DET_NAME, DET_THRESHOLD,
+                       DEVICE, HF_CLAS_REPO_NAME, HF_DET_REPO_NAME,
+                       MODELS_PATH)
 from efficientdet.efficientdet import get_transforms, rescale_bboxes, set_model
 
 
-def localize_trash(im, det_name, det_checkpoint, device, prob_threshold):
-    # detector
-    detector = set_model(det_name, 1, det_checkpoint, device)
+def localize_trash(im):
+    # detector, if checkpoint doesn't exist then download from hf
+    if not os.path.exists(DET_FILEPATH):
+        hf_hub_download(HF_DET_REPO_NAME, DET_FILENAME, local_dir=MODELS_PATH)
+    detector = set_model(DET_NAME, 1, DET_FILEPATH, DEVICE)
+
     detector.eval()
     # mean-std normalize the input image (batch-size: 1)
     img = get_transforms(im)
     # propagate through the model
-    outputs = detector(img.to(device))
+    outputs = detector(img.to(DEVICE))
     # keep only predictions above set confidence
-    bboxes_keep = outputs[0, outputs[0, :, 4] > prob_threshold]
+    bboxes_keep = outputs[0, outputs[0, :, 4] > DET_THRESHOLD]
     probas = bboxes_keep[:, 4:]
     # convert boxes to image scales
     bboxes_scaled = rescale_bboxes(bboxes_keep[:, :4], im.size, tuple(img.size()[2:]))
     return probas, bboxes_scaled
 
 
-def classify_trash(im, clas_checkpoint, cls_th, probas, bboxes_scaled):
-    # classifier
-    classifier = load_learner(clas_checkpoint)
+def classify_trash(im, probas, bboxes_scaled):
+    # classifier, if checkpoint doesn't exist then download from hf
+    if not os.path.exists(CLAS_FILEPATH):
+        hf_hub_download(HF_CLAS_REPO_NAME, CLAS_FILENAME, local_dir=MODELS_PATH)
+    classifier = load_learner(CLAS_FILEPATH)
 
     bboxes_final = []
     cls_prob = []
@@ -32,26 +44,20 @@ def classify_trash(im, clas_checkpoint, cls_th, probas, bboxes_scaled):
         outputs = classifier.predict(img)
         p[1] = torch.topk(outputs[2], k=1).indices.squeeze(0).item()
         p[0] = torch.max(np.trunc(outputs[2] * 100))
-        if p[0] >= cls_th * 100:
+        if p[0] >= CLAS_THRESHOLD * 100:
             bboxes_final.append((xmin, ymin, xmax, ymax))
             cls_prob.append(p)
     return cls_prob, bboxes_final
 
 
-def detect_trash(
-    im, det_name, det_checkpoint, clas_checkpoint, device, prob_threshold, cls_th
-):
+def detect_trash(img):
     # prepare models for evaluation
     torch.set_grad_enabled(False)
 
     # 1) Localize
-    probas, bboxes_scaled = localize_trash(
-        im, det_name, det_checkpoint, device, prob_threshold
-    )
+    probas, bboxes_scaled = localize_trash(img)
 
     # 2) Classify
-    cls_prob, bboxes_final = classify_trash(
-        im, clas_checkpoint, cls_th, probas, bboxes_scaled
-    )
+    cls_prob, bboxes_final = classify_trash(img, probas, bboxes_scaled)
 
     return cls_prob, bboxes_final