Spaces:
Runtime error
Runtime error
| 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) | |