Spaces:
Sleeping
Sleeping
| from __future__ import print_function, division | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from torch.optim import lr_scheduler | |
| import torch.backends.cudnn as cudnn | |
| import numpy as np | |
| import torchvision | |
| from torchvision import datasets, models, transforms | |
| from torch.utils.data import TensorDataset, DataLoader | |
| from PIL import Image | |
| import matplotlib.pyplot as plt | |
| from dataloader import imgDataset | |
| import time | |
| import os | |
| import copy | |
| from transformers import BlipProcessor, BlipForConditionalGeneration | |
| from transformers import AutoImageProcessor, ResNetModel | |
| from translate import Translator | |
| PATH = './images/' | |
| class CUPredictor_v2(nn.Module): | |
| def __init__(self, num_class=2): | |
| super(CUPredictor_v2, self).__init__() | |
| self.base = ResNetModel.from_pretrained("microsoft/resnet-50") | |
| num_ftrs = 2048 | |
| #self.base.fc = nn.Linear(num_ftrs, num_ftrs//2) | |
| self.classifier = nn.Linear(num_ftrs, num_class) | |
| self.height_regressor = nn.Linear(num_ftrs, 1) | |
| self.relu = nn.ReLU() | |
| def forward(self, input_img): | |
| output = self.base(input_img['pixel_values'].squeeze(1)).pooler_output.squeeze() | |
| predict_cls = self.classifier(output) | |
| predict_height = self.relu(self.height_regressor(output)) | |
| return predict_cls, predict_height | |
| class CUPredictor(nn.Module): | |
| def __init__(self, num_class=2): | |
| super(CUPredictor, self).__init__() | |
| self.base = torchvision.models.resnet50(pretrained=True) | |
| for param in self.base.parameters(): | |
| param.requires_grad = False | |
| num_ftrs = self.base.fc.in_features | |
| self.base.fc = nn.Sequential( | |
| nn.Linear(num_ftrs, num_ftrs//4), | |
| nn.ReLU(), | |
| nn.Linear(num_ftrs//4, num_ftrs//8), | |
| nn.ReLU() | |
| ) | |
| self.classifier = nn.Linear(num_ftrs//8, num_class) | |
| self.regressor_h = nn.Linear(num_ftrs//8, 1) | |
| self.regressor_b = nn.Linear(num_ftrs//8, 1) | |
| self.regressor_w = nn.Linear(num_ftrs//8, 1) | |
| self.regressor_hi = nn.Linear(num_ftrs//8, 1) | |
| self.relu = nn.ReLU() | |
| def forward(self, input_img): | |
| output = self.base(input_img) | |
| predict_cls = self.classifier(output) | |
| predict_h = self.relu(self.regressor_h(output)) | |
| predict_b = self.relu(self.regressor_b(output)) | |
| predict_w = self.relu(self.regressor_w(output)) | |
| predict_hi = self.relu(self.regressor_hi(output)) | |
| return predict_cls, predict_h, predict_b, predict_w, predict_hi | |
| def imshow(inp, title=None): | |
| """Imshow for Tensor.""" | |
| inp = inp.numpy().transpose((1, 2, 0)) | |
| mean = np.array([0.485, 0.456, 0.406]) | |
| std = np.array([0.229, 0.224, 0.225]) | |
| inp = std * inp + mean | |
| inp = np.clip(inp, 0, 1) | |
| plt.imshow(inp) | |
| if title is not None: | |
| plt.title(title) | |
| plt.pause(0.001) # pause a bit so that plots are updated | |
| plt.savefig(f'images/preds/prediction.png') | |
| def train_model(model, device, dataloaders, dataset_sizes, num_epochs=25): | |
| since = time.time() | |
| ce = nn.CrossEntropyLoss() | |
| mse = nn.MSELoss() | |
| optimizer = optim.AdamW(model.parameters(), lr=0.0008) | |
| best_model_wts = copy.deepcopy(model.state_dict()) | |
| best_acc = 0.0 | |
| for epoch in range(num_epochs): | |
| print(f'Epoch {epoch+1}/{num_epochs}') | |
| print('-' * 10) | |
| # Each epoch has a training and validation phase | |
| for phase in ['train', 'val']: | |
| if phase == 'train': | |
| model.train() # Set model to training mode | |
| else: | |
| model.eval() # Set model to evaluate mode | |
| running_ce_loss = 0.0 | |
| running_rmse_loss = 0.0 | |
| running_corrects = 0 | |
| # Iterate over data. | |
| for inputs, labels, heights, bust, waist, hips in dataloaders[phase]: | |
| inputs = inputs.to(device) | |
| labels = labels.to(device) | |
| heights = heights.to(device) | |
| bust = bust.to(device) | |
| waist, hips = waist.to(device), hips.to(device) | |
| # zero the parameter gradients | |
| optimizer.zero_grad() | |
| # forward | |
| # track history if only in train | |
| with torch.set_grad_enabled(phase == 'train'): | |
| outputs_c, outputs_h, outputs_b, outputs_w, outputs_hi = model(inputs) | |
| _, preds = torch.max(outputs_c, 1) | |
| ce_loss = ce(outputs_c, labels) | |
| rmse_loss_h = torch.sqrt(mse(outputs_h, heights.unsqueeze(-1))) | |
| rmse_loss_b = torch.sqrt(mse(outputs_b, bust.unsqueeze(-1))) | |
| rmse_loss_w = torch.sqrt(mse(outputs_w, waist.unsqueeze(-1))) | |
| rmse_loss_hi = torch.sqrt(mse(outputs_hi, hips.unsqueeze(-1))) | |
| rmse_loss = rmse_loss_h*4 + rmse_loss_b*2 + rmse_loss_w + rmse_loss_hi | |
| loss = ce_loss + (rmse_loss)*1 | |
| # backward + optimize only if in training phase | |
| if phase == 'train': | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) | |
| optimizer.step() | |
| # statistics | |
| running_ce_loss += ce_loss.item() * inputs.size(0) | |
| running_rmse_loss += rmse_loss.item() * inputs.size(0) | |
| running_corrects += torch.sum(preds == labels.data) | |
| epoch_ce_loss = running_ce_loss / dataset_sizes[phase] | |
| epoch_rmse_loss = running_rmse_loss / dataset_sizes[phase] | |
| epoch_acc = running_corrects.double() / dataset_sizes[phase] | |
| print(f'{phase} CE_Loss: {epoch_ce_loss:.4f} RMSE_Loss: {epoch_rmse_loss:.4f} Acc: {epoch_acc:.4f}') | |
| # deep copy the model | |
| if phase == 'val' and epoch_acc > best_acc: | |
| best_acc = epoch_acc | |
| best_model_wts = copy.deepcopy(model.state_dict()) | |
| #if epoch %2 == 0 and phase == 'val':print(outputs_c, outputs_h) | |
| print() | |
| time_elapsed = time.time() - since | |
| print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s') | |
| print(f'Best val Acc: {best_acc:4f}') | |
| # load best model weights | |
| model.load_state_dict(best_model_wts) | |
| return model | |
| def visualize_model(model, device, dataloaders, class_names, num_images=6): | |
| was_training = model.training | |
| model.eval() | |
| images_so_far = 0 | |
| fig = plt.figure() | |
| with torch.no_grad(): | |
| for i, (inputs, labels) in enumerate(dataloaders['val']): | |
| inputs = inputs.to(device) | |
| labels = labels.to(device) | |
| outputs = model(inputs) | |
| _, preds = torch.max(outputs, 1) | |
| for j in range(inputs.size()[0]): | |
| images_so_far += 1 | |
| ax = plt.subplot(num_images//2, 2, images_so_far) | |
| ax.axis('off') | |
| ax.set_title(f'pred: {class_names[preds[j]]}|tar: {class_names[labels[j]]}') | |
| imshow(inputs.cpu().data[j]) | |
| if images_so_far == num_images: | |
| model.train(mode=was_training) | |
| return | |
| model.train(mode=was_training) | |
| def evaluation(model, epoch, device, dataloaders): | |
| model.load_state_dict(torch.load(f'models/model_{epoch}.pt')) | |
| model.eval() | |
| with torch.no_grad(): | |
| for i, (inputs, labels) in enumerate(dataloaders['val']): | |
| inputs = inputs.to(device) | |
| labels = labels.to(device) | |
| outputs = model(inputs) | |
| _, preds = torch.max(outputs, 1) | |
| print(preds) | |
| def inference(inp_img, classes = ['big', 'small'], epoch = 6): | |
| device = torch.device("cpu") | |
| translator= Translator(to_lang="zh-TW") | |
| model = CUPredictor() | |
| model.load_state_dict(torch.load(f'models/model_{epoch}.pt', map_location=torch.device('cpu'))) | |
| # load image-to-text model | |
| processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") | |
| model_blip = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base") | |
| model.eval() | |
| trans = transforms.Compose([ | |
| transforms.Resize(256), | |
| transforms.CenterCrop(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| image_tensor = trans(inp_img) | |
| image_tensor = image_tensor.unsqueeze(0) | |
| with torch.no_grad(): | |
| inputs = image_tensor.to(device) | |
| outputs_c, outputs_h, outputs_b, outputs_w, outputs_hi = model(inputs) | |
| _, preds = torch.max(outputs_c, 1) | |
| idx = preds.numpy()[0] | |
| # unconditional image captioning | |
| inputs = processor(inp_img, return_tensors="pt") | |
| out = model_blip.generate(**inputs) | |
| description = processor.decode(out[0], skip_special_tokens=True) | |
| description_tw = translator.translate(description) | |
| return outputs_c, classes[idx], f"{outputs_h.numpy()[0][0]:.2f}", f"{outputs_b.numpy()[0][0]:.2f}", f"{outputs_w.numpy()[0][0]:.2f}", f"{outputs_hi.numpy()[0][0]:.2f}", [description, description_tw] | |
| def main(epoch = 15, mode = 'val'): | |
| cudnn.benchmark = True | |
| plt.ion() # interactive mode | |
| model = CUPredictor() | |
| train_dataset = imgDataset('labels.txt', mode='train', use_processor=False) | |
| test_dataset = imgDataset('labels.txt', mode='val', use_processor=False) | |
| dataloaders = { | |
| "train": DataLoader(train_dataset, batch_size=64, shuffle=True), | |
| "val": DataLoader(test_dataset, batch_size=64, shuffle=False) | |
| } | |
| dataset_sizes = { | |
| "train": len(train_dataset), | |
| "val": len(test_dataset) | |
| } | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| #device = torch.device("cpu") | |
| model = model.to(device) | |
| model_conv = train_model(model, device, dataloaders, dataset_sizes, num_epochs=epoch) | |
| torch.save(model_conv.state_dict(), f'models/model_{epoch}.pt') | |
| def divide_class_dir(path): | |
| file_list = os.listdir(path) | |
| for img_name in file_list: | |
| dest_path = os.path.join(path, img_name.split('-')[3]) | |
| if not os.path.exists(dest_path): | |
| os.mkdir(dest_path) # 建立資料夾 | |
| os.replace(os.path.join(path, img_name), os.path.join(dest_path, img_name)) | |
| def get_label(types): | |
| with open('labels.txt', 'w', encoding='utf-8') as f: | |
| for f_type in types: | |
| for img_type in CLASS: | |
| path = os.path.join('images', f_type, img_type) | |
| file_list = os.listdir(path) | |
| for file_name in file_list: | |
| file_name_list = file_name.split('-') | |
| f.write(" ".join([f_type, file_name, img_type, file_name_list[4].split('_')[0], '\n'])) | |
| if __name__ == "__main__": | |
| CLASS = ['big', 'small'] | |
| mode = 'train' | |
| get_label(['train', 'val']) | |
| epoch = 7 | |
| #main(epoch, mode = mode) | |
| outputs, preds, heights, bust, waist, hips, description = inference('images/test/lin.png', CLASS, epoch=epoch) | |
| print(outputs, preds, heights, bust, waist, hips) | |
| #print(CUPredictor()) | |
| #divide_class_dir('./images/train_all') | |
| #divide_class_dir('./images/val_all') | |
| '''''' |