app.py

import torch
import pickle
import cv2
import os
import numpy as np
from PIL import Image
from transformers import ViTForImageClassification, AutoImageProcessor, AdamW, ViTImageProcessor, VisionEncoderDecoderModel, AutoTokenizer
from torch.utils.data import DataLoader, TensorDataset

model_path = 'model'
train_pickle_path = 'train_data.pickle'
valid_pickle_path = 'valid_data.pickle'
image_directory = 'images' 
test_image_path = 'test.jpg'
num_epochs = 5    # Fine-tune the model
label_list = ["小白", "巧巧", "冏媽", "乖狗", "花捲", "超人", "黑胖", "橘子"]
label_dictionary = {"小白": 0, "巧巧": 1, "冏媽": 2, "乖狗": 3, "花捲": 4, "超人": 5, "黑胖": 6, "橘子": 7}
num_classes = len(label_dictionary)  # Adjust according to your classification task
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# device = torch.device("mps")

def data_generate(dataset):
    images = []
    labels = []
    image_processor = AutoImageProcessor.from_pretrained('google/vit-large-patch16-224-in21k')
    for folder_name in os.listdir(image_directory):
        folder_path = os.path.join(image_directory, folder_name)
        if os.path.isdir(folder_path):
            for image_file in os.listdir(folder_path):
                if image_file.startswith(dataset):
                    image_path = os.path.join(folder_path, image_file)
                    # print(image_path)

                    img = cv2.imread(image_path)
                    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                    img = Image.fromarray(img)
                    img = img.resize((224, 224))
                    inputs = image_processor(images=img, return_tensors="pt")
                    images.append(inputs['pixel_values'].squeeze(0).numpy())
                    labels.append(int(folder_name.split('_')[0]))

    images = np.array(images)
    labels = np.array(labels)

    # Now you can pickle this data
    train_data = {'img': images, 'label': labels}
    with open(f'{dataset}_data.pickle', 'wb') as f:
        pickle.dump(train_data, f)

def train_model():
    
    if not os.path.exists(valid_pickle_path):
        data_generate('valid')
    if not os.path.exists(train_pickle_path):
        data_generate('train')

    # Load the train and vaild
    with open("train_data.pickle", "rb") as f:
        train_data = pickle.load(f)

    with open("valid_data.pickle", "rb") as f:
        valid_data = pickle.load(f)

    # Convert the dataset into torch tensors
    train_inputs = torch.tensor(train_data["img"])
    train_labels = torch.tensor(train_data["label"])
    valid_inputs = torch.tensor(valid_data["img"])
    valid_labels = torch.tensor(valid_data["label"])

    # Create the TensorDataset
    train_dataset = TensorDataset(train_inputs, train_labels)
    valid_dataset = TensorDataset(valid_inputs, valid_labels)

    # Create the DataLoader
    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
    valid_loader = DataLoader(valid_dataset, batch_size=16, shuffle=True)

    # Define the model and move it to the GPU
    
    model = ViTForImageClassification.from_pretrained('google/vit-large-patch16-224-in21k', num_labels=num_classes)
    model.to(device)

    # Define the optimizer
    optimizer = AdamW(model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):

        model.train()
        total_loss = 0

        for i, batch in enumerate(train_loader):

            # Move batch to the GPU
            batch = [r.to(device) for r in batch]

            # Unpack the inputs from our dataloader
            inputs, labels = batch

            # Clear out the gradients (by default they accumulate)
            optimizer.zero_grad()

            # Forward pass
            outputs = model(inputs, labels=labels)

            # Compute loss
            loss = outputs.loss

            # Backward pass
            loss.backward()


            # Update parameters and take a step using the computed gradient
            optimizer.step()

            # Update the loss
            total_loss += loss.item()

            # print(f'{i}/{len(train_loader)} ')
            
        # Get the average loss for the entire epoch
        avg_loss = total_loss / len(train_loader)
        
        # Print the loss
        print('Epoch:', epoch + 1, 'Training Loss:', avg_loss)
        
        
    # Evaluate the model on the validation set
    model.eval()
    total_correct = 0

    for batch in valid_loader:
        # Move batch to the GPU
        batch = [t.to(device) for t in batch]

        # Unpack the inputs from our dataloader
        inputs, labels = batch

        # Forward pass
        with torch.no_grad():
            outputs = model(inputs)

        # Get the predictions
        predictions = torch.argmax(outputs.logits, dim=1)

        # Update the total correct
        total_correct += torch.sum(predictions == labels)
        
    # Calculate the accuracy
    accuracy = total_correct / len(valid_dataset)
    print('Validation accuracy:', accuracy.item())

    model.save_pretrained("model")

def predict():
    # Load the model
    model = ViTForImageClassification.from_pretrained(model_path, num_labels=num_classes)

    image_processor = AutoImageProcessor.from_pretrained('google/vit-large-patch16-224-in21k')


    # Load the test data
    # Load the image

    img = cv2.imread(test_image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # Resize the image to 224x224 pixels
    img = Image.fromarray(img)
    img = img.resize((224, 224))

    # img to tensor
    # Preprocess the image and generate features
    inputs = image_processor(images=img, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits

    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    predicted_class_idx = logits.argmax(-1).item()

    return label_list[predicted_class_idx] if probabilities.max().item() > 0.90 else '不是校狗'

def captioning():
    
    model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
    feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
    tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    max_length = 16
    num_beams = 4
    gen_kwargs = {"max_length": max_length, "num_beams": num_beams}

    images = []
    for image_path in [test_image_path]:
        i_image = Image.open(image_path)
        if i_image.mode != "RGB":
            i_image = i_image.convert(mode="RGB")

        images.append(i_image)

    pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(device)

    output_ids = model.generate(pixel_values, **gen_kwargs)

    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    preds = [pred.strip() for pred in preds]
    return preds[-1]

def output(predict_class, caption):
    conj = ['are', 'is', 'dog']
    if predict_class == '不是校狗' or caption.find('dog') == -1:
        print(f'{caption} ({predict_class})')
    else:
        for c in conj:
            if caption.find(c) != -1:
                print(f'{predict_class} is{caption[caption.find(c) + len(c):]}')
                return
        print(f'{caption} ({predict_class})')


if __name__ == '__main__':

    if not os.path.exists(model_path):
        train_model()
    output(predict(), captioning())