File size: 1,989 Bytes
7b2449b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
# This script prepares the training images and ratings for the training.
# It assumes that all images are stored as files that PIL can read.
# It also assumes that the paths to the images files and the average ratings are in a .parquet files that can be read into a dataframe ( df ).
from datasets import load_dataset
import pandas as pd
import statistics
from torch.utils.data import Dataset, DataLoader
import clip
import torch
from PIL import Image, ImageFile
import numpy as np
import time
def normalized(a, axis=-1, order=2):
import numpy as np # pylint: disable=import-outside-toplevel
l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
l2[l2 == 0] = 1
return a / np.expand_dims(l2, axis)
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-L/14", device=device)
f = "trainingdata.parquet"
df = pd.read_parquet(f) #assumes that the df has the columns IMAGEPATH & AVERAGE_RATING
x = []
y = []
c= 0
for idx, row in df.iterrows():
start = time.time()
average_rating = float(row.AVERAGE_RATING)
print(average_rating)
if average_rating <1:
continue
img= row.IMAGEPATH #assumes that the df has the column IMAGEPATH
print(img)
try:
image = preprocess(Image.open(img)).unsqueeze(0).to(device)
except:
continue
with torch.no_grad():
image_features = model.encode_image(image)
im_emb_arr = image_features.cpu().detach().numpy()
x.append(normalized ( im_emb_arr) ) # all CLIP embeddings are getting normalized. This also has to be done when inputting an embedding later for inference
y_ = np.zeros((1, 1))
y_[0][0] = average_rating
#y_[0][1] = stdev # I initially considered also predicting the standard deviation, but then didn't do it
y.append(y_)
print(c)
c+=1
x = np.vstack(x)
y = np.vstack(y)
print(x.shape)
print(y.shape)
np.save('x_OpenAI_CLIP_L14_embeddings.npy', x)
np.save('y_ratings.npy', y)
|