Adds files from huggingface hub repo

.gitattributes

@@ -0,0 +1,31 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+feedback*
+new_model/
+__pycache__/
+data/
+events.out.*

README.md

@@ -0,0 +1,12 @@
+---
+title: Grocery Classifier Demo
+emoji: 🛒
+colorFrom: red
+colorTo: green
+sdk: streamlit
+sdk_version: 1.10.0
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,88 @@
+import os
+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"
+)
+
+@st.cache()
+def load_model():
+    with st.spinner("Loading 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
+    image = Image.open(image)
+
+    prediction, confidence = model.predict(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:
+        st.write("Labels:", labels)
+
+    image_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    if image_file is not None:
+        st.image(image_file)
+
+        st.subheader("Classification")
+        
+        if st.button("Predict"):
+            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:
+            # Show the result
+            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")
+            correct_labels = labels.copy()
+            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'])
+                    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:
+                    retrain_from_feedback()
+                    st.success("Model retrained")
+                except Exception as e:
+                    st.warning("Model could not be retrained. Error: {}".format(e))
+                    
+main()

dataset.py

@@ -0,0 +1,27 @@
+import torch
+
+class RetailDataset(torch.utils.data.Dataset):
+    def __init__(self, data, labels=None, transform=None):
+        self.data = data
+        self.labels = labels
+        self.num_classes = len(set(labels))
+        self.transform = transform
+
+    def __getitem__(self, idx):
+        item = {key: val[idx].detach().clone() for key, val in self.data.items()}
+        item['labels'] = self.labels[idx]
+        return item
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __repr__(self):
+        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)
+        })

model.py

@@ -0,0 +1,191 @@
+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
+from PIL import Image
+import logging
+import os
+from sklearn.preprocessing import LabelEncoder
+from train import (
+    re_training, metric, f1_score,
+    classification_report
+)
+
+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):
+    def __init__(self, model_name, num_labels=24, dropout=0.25, image_size=224):
+        logger.info("Loading model")
+        super(ViTForImageClassification, self).__init__()
+        self.vit = ViTModel.from_pretrained(model_name)
+        self.feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+        self.feature_extractor.do_resize = True
+        self.feature_extractor.size = image_size
+        self.dropout = nn.Dropout(dropout)
+        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.model_name = model_name
+        # To device
+        self.vit.to(self.device)
+        self.to(self.device)
+        self.classifier.to(self.device)
+        logger.info("Model loaded")
+
+    def forward(self, pixel_values, labels):
+        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])
+        logits = self.classifier(output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def preprocess_image(self, images):
+        logger.info("Preprocessing images")
+        return self.feature_extractor(images, return_tensors='pt')
+
+    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']
+            sequence_classifier_output = self.forward(images_batch, None)
+            # Get max prob
+            probs = sequence_classifier_output.logits.softmax(dim=-1).tolist()
+            classes = np.argmax(probs, axis=1)
+            confidences = np.max(probs, axis=1)
+            classes_list.extend(classes)
+            confidence_list.extend(confidences)
+        if classes_names:
+            classes_list = self.label_encoder.inverse_transform(classes_list)
+        if return_probabilities:
+            return classes_list, confidence_list, probs
+        return classes_list, confidence_list
+
+    def save(self, path):
+        logger.info("Saving model")
+        os.makedirs(path, exist_ok=True)
+        torch.save(self.state_dict(), path + "/model.pt")
+        # Save label encoder
+        np.save(path + "/label_encoder.npy", self.label_encoder.classes_)
+
+    def load(self, path):
+        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"):
+            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.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)
+        # Predict
+        y_pred, _ = self.predict(images, classes_names=False)
+        # 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(f"Accuracy: {metrics['accuracy']}")
+        print(f"F1: {f1}")
+        
+    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()]
+        # for param in params[:-1]:
+        #     param.requires_grad = False
+        # Model in training mode
+        self.vit.train()
+        self.train()
+        re_training(images, labels, self, save_model_path, num_epochs)
+        self.load(save_model_path)
+        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)):
+                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)
+                    if image.size != image_shape:
+                        image = image.resize(image_shape)
+                    images.append(image.convert('RGB'))
+                    labels.append(label)
+                    count += 1
+                except Exception as e:
+                    print(f"ERROR - Could not resize image {file_path} - {e}")
+                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'):
+        logger.info("Retraining from path")
+        # Load path
+        images, labels = self.__load_from_path(path, num_per_label)
+        # Retrain
+        self.partial_fit(images, labels, save_model_path, num_epochs)
+        # Save new data
+        if save_new_data is not None:
+            logger.info("Saving new data")
+            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"))
+        # 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

@@ -0,0 +1,11 @@
+Pillow
+requests
+numpy
+transformers
+scikit-learn
+datasets
+streamlit
+matplotlib
+scikit-image
+torch
+torchvision

train.py

@@ -0,0 +1,110 @@
+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 PIL import Image
+from datasets import load_metric
+from torchvision.transforms import (
+    CenterCrop,
+    Compose,
+    Normalize,
+    RandomHorizontalFlip,
+    RandomResizedCrop,
+    Resize,
+    ToTensor,
+)
+from transformers import Trainer, TrainingArguments, BatchFeature
+metric = load_metric("accuracy")
+f1_score = load_metric("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):
+    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)
+    except ValueError:
+        logger.warning("Could not split dataset. Using all data for training and testing")
+        images_train = images
+        labels_train = labels
+        images_test = images
+        labels_test = labels
+
+    # 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]]
+        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_test_batch = model.preprocess_image(images_test_batch)
+        images_test_prep.extend(images_test_batch['pixel_values'])
+
+    # Create BatchFeatures
+    images_train_prep = {"pixel_values": images_train_prep}
+    train_batch_features = BatchFeature(data=images_train_prep)
+    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)
+    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):
+    global model
+    model = _model
+    labels = model.label_encoder.transform(labels)
+    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)
+
+    def val_transforms(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)
+    trainer = Trainer(
+        model=model,
+        args=TrainingArguments(
+            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_steps=1000,
+            save_steps=3000),
+        train_dataset=train_dataset,
+        eval_dataset=test_dataset
+    )
+    trainer.train()
+    model.save(save_model_path)