Refactor and improve model, app, and training components

- Update dependencies in requirements.txt with pinned versions
- Enhance device handling in model and dataset classes
- Improve Streamlit app caching and error handling
- Optimize training and retraining procedures
- Add support for MPS device in model selection
- Update .gitignore to include output directory

Signed-off-by: Unai Garay <unaigaraymaestre@gmail.com>

.gitignore

@@ -2,4 +2,5 @@ feedback*
 new_model/
 __pycache__/
 data/
-events.out.*
+events.out.*
+output/

app.py

@@ -1,27 +1,30 @@
 import os
+import time
+
 import streamlit as st
 from PIL import Image
-import requests
-import io
-import time
+
 from model import ViTForImageClassification
 
 st.set_page_config(
-     page_title="Grocery Classifier",
-     page_icon="interface/shopping-cart.png",
-     initial_sidebar_state="expanded"
+    page_title="Grocery Classifier",
+    page_icon="interface/shopping-cart.png",
+    initial_sidebar_state="expanded",
 )
 
-@st.cache()
+
+@st.cache_resource()
 def load_model():
     with st.spinner("Loading model"):
-        model = ViTForImageClassification('google/vit-base-patch16-224')
-        model.load('model/')
+        model = ViTForImageClassification("google/vit-base-patch16-224")
+        model.load("model/")
     return model
-        
+
+
 model = load_model()
 feedback_path = "feedback"
 
+
 def predict(image):
     print("Predicting...")
     # Load using PIL
@@ -29,21 +32,24 @@ def predict(image):
 
     prediction, confidence = model.predict(image)
 
-    return {'prediction': prediction[0], 'confidence': round(confidence[0], 3)}, image
+    return {"prediction": prediction[0], "confidence": round(confidence[0], 3)}, image
+
 
 def submit_feedback(correct_label, image):
     folder_path = feedback_path + "/" + correct_label + "/"
     os.makedirs(folder_path, exist_ok=True)
     image.save(folder_path + correct_label + "_" + str(int(time.time())) + ".png")
-    
+
+
 def retrain_from_feedback():
     model.retrain_from_path(feedback_path, remove_path=True)
 
+
 def main():
     labels = set(list(model.label_encoder.classes_))
 
     st.title("🍇 Grocery Classifier 🥑")
-        
+
     if labels is None:
         st.warning("Received error from server, labels could not be retrieved")
     else:
@@ -54,37 +60,49 @@ def main():
         st.image(image_file)
 
         st.subheader("Classification")
-        
+
         if st.button("Predict"):
-            st.session_state['response_json'], st.session_state['image'] = predict(image_file)
+            st.session_state["response_json"], st.session_state["image"] = predict(
+                image_file
+            )
 
-        if 'response_json' in st.session_state and st.session_state['response_json'] is not None:
+        if (
+            "response_json" in st.session_state
+            and st.session_state["response_json"] is not None
+        ):
             # Show the result
-            st.markdown(f"**Prediction:** {st.session_state['response_json']['prediction']}")
-            st.markdown(f"**Confidence:** {st.session_state['response_json']['confidence']}")
-            
+            st.markdown(
+                f"**Prediction:** {st.session_state['response_json']['prediction']}"
+            )
+            st.markdown(
+                f"**Confidence:** {st.session_state['response_json']['confidence']}"
+            )
+
             # User feedback
             st.subheader("User Feedback")
-            st.markdown("If this prediction was incorrect, please select below the correct label")
+            st.markdown(
+                "If this prediction was incorrect, please select below the correct label"
+            )
             correct_labels = labels.copy()
-            correct_labels.remove(st.session_state['response_json']["prediction"])
+            correct_labels.remove(st.session_state["response_json"]["prediction"])
             correct_label = st.selectbox("Correct label", correct_labels)
             if st.button("Submit"):
                 # Save feedback
                 try:
-                    submit_feedback(correct_label, st.session_state['image'])
+                    submit_feedback(correct_label, st.session_state["image"])
                     st.success("Feedback submitted")
                 except Exception as e:
                     st.error("Feedback could not be submitted. Error: {}".format(e))
-                    
+
             # Retrain from feedback
             if st.button("Retrain from feedback"):
                 try:
-                    with st.spinner('Retraining...'):
+                    with st.spinner("Retraining..."):
                         retrain_from_feedback()
                     st.success("Model retrained")
                     st.balloons()
                 except Exception as e:
                     st.warning("Model could not be retrained. Error: {}".format(e))
-                    
-main()
+
+
+main()

dataset.py

@@ -1,27 +1,34 @@
 import torch
 
+
 class RetailDataset(torch.utils.data.Dataset):
-    def __init__(self, data, labels=None, transform=None):
+    def __init__(self, data, labels=None, transform=None, device=None):
         self.data = data
         self.labels = labels
         self.num_classes = len(set(labels))
         self.transform = transform
+        self.device = device if device is not None else torch.device("cpu")
 
     def __getitem__(self, idx):
-        item = {key: val[idx].detach().clone() for key, val in self.data.items()}
-        item['labels'] = self.labels[idx]
+        item = {
+            key: torch.tensor(val[idx].detach().clone(), device=self.device)
+            for key, val in self.data.items()
+        }
+        item["labels"] = torch.tensor(self.labels[idx], device=self.device)
         return item
 
     def __len__(self):
         return len(self.labels)
 
     def __repr__(self):
-        return 'RetailDataset'
+        return "RetailDataset"
 
     def __str__(self):
-        return str({
-            'data': self.data['pixel_values'].shape,
-            'labels': self.labels.shape,
-            'num_classes': self.num_classes,
-            'num_samples': len(self.labels)
-        })
+        return str(
+            {
+                "data": self.data["pixel_values"].shape,
+                "labels": self.labels.shape,
+                "num_classes": self.num_classes,
+                "num_samples": len(self.labels),
+            }
+        )

model.py

@@ -1,26 +1,39 @@
+import logging
+import os
 import shutil
 import time
+
 import numpy as np
-from tqdm import tqdm
-from transformers import ViTModel, ViTFeatureExtractor
-from transformers.modeling_outputs import SequenceClassifierOutput
-import torch.nn as nn
 import torch
+import torch.nn as nn
 from PIL import Image
-import logging
-import os
+from sklearn.metrics import classification_report
 from sklearn.preprocessing import LabelEncoder
+from tqdm import tqdm
+from transformers import ViTFeatureExtractor, ViTModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
 from train import (
-    re_training, metric, f1_score,
-    classification_report
+    f1_score,
+    metric,
+    re_training,
 )
 
-data_path = os.environ.get('DATA_PATH', "./data")
+data_path = os.environ.get("DATA_PATH", "./data")
 
 logging.basicConfig(level=os.getenv("LOGGER_LEVEL", logging.WARNING))
 logger = logging.getLogger(__name__)
 
+
 class ViTForImageClassification(nn.Module):
+    @staticmethod
+    def get_device():
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")
+
     def __init__(self, model_name, num_labels=24, dropout=0.25, image_size=224):
         logger.info("Loading model")
         super(ViTForImageClassification, self).__init__()
@@ -32,7 +45,8 @@ class ViTForImageClassification(nn.Module):
         self.classifier = nn.Linear(self.vit.config.hidden_size, num_labels)
         self.num_labels = num_labels
         self.label_encoder = LabelEncoder()
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = self.get_device()
+        logger.info(f"Using device: {self.device}")
         self.model_name = model_name
         # To device
         self.vit.to(self.device)
@@ -44,7 +58,7 @@ class ViTForImageClassification(nn.Module):
         logger.info("Forwarding")
         pixel_values = pixel_values.to(self.device)
         outputs = self.vit(pixel_values=pixel_values)
-        output = self.dropout(outputs.last_hidden_state[:,0])
+        output = self.dropout(outputs.last_hidden_state[:, 0])
         logits = self.classifier(output)
 
         loss = None
@@ -61,17 +75,21 @@ class ViTForImageClassification(nn.Module):
 
     def preprocess_image(self, images):
         logger.info("Preprocessing images")
-        return self.feature_extractor(images, return_tensors='pt')
+        return self.feature_extractor(images, return_tensors="pt")
 
-    def predict(self, images, batch_size=32, classes_names=True, return_probabilities=False):
+    def predict(
+        self, images, batch_size=32, classes_names=True, return_probabilities=False
+    ):
         logger.info("Predicting")
         if not isinstance(images, list):
             images = [images]
         classes_list = []
         confidence_list = []
-        for bs in tqdm(range(0, len(images), batch_size), desc="Preprocessing training images"):
-            images_batch = [image for image in images[bs:bs+batch_size]]
-            images_batch = self.preprocess_image(images_batch)['pixel_values']
+        for bs in tqdm(
+            range(0, len(images), batch_size), desc="Preprocessing training images"
+        ):
+            images_batch = [image for image in images[bs : bs + batch_size]]
+            images_batch = self.preprocess_image(images_batch)["pixel_values"]
             sequence_classifier_output = self.forward(images_batch, None)
             # Get max prob
             probs = sequence_classifier_output.logits.softmax(dim=-1).tolist()
@@ -96,19 +114,23 @@ class ViTForImageClassification(nn.Module):
         logger.info("Loading model")
         # Load label encoder
         # Check if label encoder and model exists
-        if not os.path.exists(path + "/label_encoder.npy") or not os.path.exists(path + "/model.pt"):
+        if not os.path.exists(path + "/label_encoder.npy") or not os.path.exists(
+            path + "/model.pt"
+        ):
             logger.warning("Label encoder or model not found")
             return
         self.label_encoder.classes_ = np.load(path + "/label_encoder.npy")
         # Reload classifier layer
-        self.classifier = nn.Linear(self.vit.config.hidden_size, len(self.label_encoder.classes_))
-        
+        self.classifier = nn.Linear(
+            self.vit.config.hidden_size, len(self.label_encoder.classes_)
+        )
+
         self.load_state_dict(torch.load(path + "/model.pt", map_location=self.device))
         self.vit.to(self.device)
         self.vit.eval()
         self.to(self.device)
         self.eval()
-        
+
     def evaluate(self, images, labels):
         logger.info("Evaluating")
         labels = self.label_encoder.transform(labels)
@@ -117,11 +139,18 @@ class ViTForImageClassification(nn.Module):
         # Evaluate
         metrics = metric.compute(predictions=y_pred, references=labels)
         f1 = f1_score.compute(predictions=y_pred, references=labels, average="macro")
-        print(classification_report(labels, y_pred, labels=[i for i in range(len(self.label_encoder.classes_))], target_names=self.label_encoder.classes_))
+        print(
+            classification_report(
+                labels,
+                y_pred,
+                labels=[i for i in range(len(self.label_encoder.classes_))],
+                target_names=self.label_encoder.classes_,
+            )
+        )
         print(f"Accuracy: {metrics['accuracy']}")
         print(f"F1: {f1}")
-        
-    def partial_fit(self, images, labels, save_model_path='new_model', num_epochs=10):
+
+    def partial_fit(self, images, labels, save_model_path="new_model", num_epochs=10):
         logger.info("Partial fitting")
         # Freeze ViT model but last layer
         # params = [param for param in self.vit.parameters()]
@@ -135,21 +164,27 @@ class ViTForImageClassification(nn.Module):
         self.vit.eval()
         self.eval()
         self.evaluate(images, labels)
-        
+
     def __load_from_path(self, path, num_per_label=None):
         images = []
         labels = []
         for label in os.listdir(path):
             count = 0
             label_folder_path = os.path.join(path, label)
-            for image_file in tqdm(os.listdir(label_folder_path), desc="Resizing images for label {}".format(label)):
+            for image_file in tqdm(
+                os.listdir(label_folder_path),
+                desc="Resizing images for label {}".format(label),
+            ):
                 file_path = os.path.join(label_folder_path, image_file)
                 try:
                     image = Image.open(file_path)
-                    image_shape = (self.feature_extractor.size, self.feature_extractor.size)
+                    image_shape = (
+                        self.feature_extractor.size,
+                        self.feature_extractor.size,
+                    )
                     if image.size != image_shape:
                         image = image.resize(image_shape)
-                    images.append(image.convert('RGB'))
+                    images.append(image.convert("RGB"))
                     labels.append(label)
                     count += 1
                 except Exception as e:
@@ -157,14 +192,16 @@ class ViTForImageClassification(nn.Module):
                 if num_per_label is not None and count >= num_per_label:
                     break
         return images, labels
-        
-    def retrain_from_path(self,
-                          path='./data/feedback',
-                          num_per_label=None,
-                          save_model_path='new_model',
-                          remove_path=False,
-                          num_epochs=10,
-                          save_new_data=data_path + '/new_data'):
+
+    def retrain_from_path(
+        self,
+        path="./data/feedback",
+        num_per_label=None,
+        save_model_path="new_model",
+        remove_path=False,
+        num_epochs=10,
+        save_new_data=data_path + "/new_data",
+    ):
         logger.info("Retraining from path")
         # Load path
         images, labels = self.__load_from_path(path, num_per_label)
@@ -173,19 +210,20 @@ class ViTForImageClassification(nn.Module):
         # Save new data
         if save_new_data is not None:
             logger.info("Saving new data")
-            for i ,(image, label) in enumerate(zip(images, labels)):
+            for i, (image, label) in enumerate(zip(images, labels)):
                 label_path = os.path.join(save_new_data, label)
                 os.makedirs(label_path, exist_ok=True)
-                image.save(os.path.join(label_path, str(int(time.time())) + f"_{i}.jpg"))
+                image.save(
+                    os.path.join(label_path, str(int(time.time())) + f"_{i}.jpg")
+                )
         # Remove path folder
         if remove_path:
             logger.info("Removing feedback path")
             shutil.rmtree(path)
-        
+
     def evaluate_from_path(self, path, num_per_label=None):
         logger.info("Evaluating from path")
         # Load images
         images, labels = self.__load_from_path(path, num_per_label)
         # Evaluate
         self.evaluate(images, labels)
-                    

requirements.txt

@@ -1,11 +1,14 @@
-Pillow
-requests
-numpy
-transformers
-scikit-learn
-datasets
-streamlit
-matplotlib
-scikit-image
-torch
-torchvision
+Pillow==10.4.0
+requests==2.32.3
+numpy==1.24.4
+transformers==4.46.3
+scikit-learn==1.3.2
+datasets==3.1.0
+streamlit==1.40.1
+matplotlib==3.7.5
+scikit-image==0.21.0
+torch==2.4.1
+torchvision==0.19.1
+altair==5.4.1
+evaluate==0.4.3
+accelerate==1.0.1

train.py

@@ -1,12 +1,10 @@
+import logging
 import os
+
 import numpy as np
-from sklearn.metrics import classification_report
-from tqdm import tqdm
-import logging
-from sklearn.model_selection import train_test_split
-from dataset import RetailDataset
+from evaluate import load
 from PIL import Image
-from datasets import load_metric
+from sklearn.model_selection import train_test_split
 from torchvision.transforms import (
     CenterCrop,
     Compose,
@@ -16,28 +14,38 @@ from torchvision.transforms import (
     Resize,
     ToTensor,
 )
-from transformers import Trainer, TrainingArguments, BatchFeature
-metric = load_metric("accuracy")
-f1_score = load_metric("f1")
+from tqdm import tqdm
+from transformers import BatchFeature, Trainer, TrainingArguments
+
+from dataset import RetailDataset
+
+metric = load("accuracy")
+f1_score = load("f1")
 np.random.seed(42)
 
 logging.basicConfig(level=os.getenv("LOGGER_LEVEL", logging.WARNING))
 logger = logging.getLogger(__name__)
-    
-def prepare_dataset(images,
-                    labels,
-                    model,
-                    test_size=.2,
-                    train_transform=None,
-                    val_transform=None,
-                    batch_size=512):
+
+
+def prepare_dataset(
+    images,
+    labels,
+    model,
+    test_size=0.2,
+    train_transform=None,
+    val_transform=None,
+    batch_size=512,
+):
     logger.info("Preparing dataset")
     # Split the dataset in train and test
     try:
-        images_train, images_test, labels_train, labels_test = \
-            train_test_split(images, labels, test_size=test_size)
+        images_train, images_test, labels_train, labels_test = train_test_split(
+            images, labels, test_size=test_size
+        )
     except ValueError:
-        logger.warning("Could not split dataset. Using all data for training and testing")
+        logger.warning(
+            "Could not split dataset. Using all data for training and testing"
+        )
         images_train = images
         labels_train = labels
         images_test = images
@@ -46,14 +54,24 @@ def prepare_dataset(images,
     # Preprocess images using model feature extractor
     images_train_prep = []
     images_test_prep = []
-    for bs in tqdm(range(0, len(images_train), batch_size), desc="Preprocessing training images"):
-        images_train_batch = [Image.fromarray(np.array(image)) for image in images_train[bs:bs+batch_size]]
+    for bs in tqdm(
+        range(0, len(images_train), batch_size), desc="Preprocessing training images"
+    ):
+        images_train_batch = [
+            Image.fromarray(np.array(image))
+            for image in images_train[bs : bs + batch_size]
+        ]
         images_train_batch = model.preprocess_image(images_train_batch)
-        images_train_prep.extend(images_train_batch['pixel_values'])
-    for bs in tqdm(range(0, len(images_test), batch_size), desc="Preprocessing test images"):
-        images_test_batch = [Image.fromarray(np.array(image)) for image in images_test[bs:bs+batch_size]]
+        images_train_prep.extend(images_train_batch["pixel_values"])
+    for bs in tqdm(
+        range(0, len(images_test), batch_size), desc="Preprocessing test images"
+    ):
+        images_test_batch = [
+            Image.fromarray(np.array(image))
+            for image in images_test[bs : bs + batch_size]
+        ]
         images_test_batch = model.preprocess_image(images_test_batch)
-        images_test_prep.extend(images_test_batch['pixel_values'])
+        images_test_prep.extend(images_test_batch["pixel_values"])
 
     # Create BatchFeatures
     images_train_prep = {"pixel_values": images_train_prep}
@@ -61,50 +79,67 @@ def prepare_dataset(images,
     images_test_prep = {"pixel_values": images_test_prep}
     test_batch_features = BatchFeature(data=images_test_prep)
 
-    # Create the datasets
-    train_dataset = RetailDataset(train_batch_features, labels_train, train_transform)
-    test_dataset = RetailDataset(test_batch_features, labels_test, val_transform)
+    # Create the datasets with proper device
+    train_dataset = RetailDataset(
+        train_batch_features, labels_train, train_transform, device=model.device
+    )
+    test_dataset = RetailDataset(
+        test_batch_features, labels_test, val_transform, device=model.device
+    )
     logger.info("Train dataset: %d images", len(labels_train))
     logger.info("Test dataset: %d images", len(labels_test))
     return train_dataset, test_dataset
 
-def re_training(images, labels, _model, save_model_path='new_model', num_epochs=10):
+
+def re_training(images, labels, _model, save_model_path="new_model", num_epochs=10):
     global model
     model = _model
     labels = model.label_encoder.transform(labels)
-    normalize = Normalize(mean=model.feature_extractor.image_mean, std=model.feature_extractor.image_std)
+    normalize = Normalize(
+        mean=model.feature_extractor.image_mean, std=model.feature_extractor.image_std
+    )
+
     def train_transforms(batch):
-        return Compose([
-            RandomResizedCrop(model.feature_extractor.size),
-            RandomHorizontalFlip(),
-            ToTensor(),
-            normalize,
-        ])(batch)
+        return Compose(
+            [
+                RandomResizedCrop(model.feature_extractor.size),
+                RandomHorizontalFlip(),
+                ToTensor(),
+                normalize,
+            ]
+        )(batch)
 
     def val_transforms(batch):
-        return Compose([
-            Resize(model.feature_extractor.size),
-            CenterCrop(model.feature_extractor.size),
-            ToTensor(),
-            normalize,
-        ])(batch)
+        return Compose(
+            [
+                Resize(model.feature_extractor.size),
+                CenterCrop(model.feature_extractor.size),
+                ToTensor(),
+                normalize,
+            ]
+        )(batch)
+
     train_dataset, test_dataset = prepare_dataset(
-        images, labels, model, .2, train_transforms, val_transforms)
+        images, labels, model, 0.2, train_transforms, val_transforms
+    )
+
     trainer = Trainer(
         model=model,
         args=TrainingArguments(
-            output_dir='output',
+            output_dir="output",
             overwrite_output_dir=True,
             num_train_epochs=num_epochs,
             per_device_train_batch_size=32,
             gradient_accumulation_steps=1,
             learning_rate=0.000001,
             weight_decay=0.01,
-            evaluation_strategy='steps',
+            eval_strategy="steps",
             eval_steps=1000,
-            save_steps=3000),
+            save_steps=3000,
+            use_cpu=model.device.type == "cpu",  # Only force CPU if that's our device
+        ),
         train_dataset=train_dataset,
-        eval_dataset=test_dataset
+        eval_dataset=test_dataset,
     )
     trainer.train()
-    model.save(save_model_path)
+    model.save(save_model_path)