visulaize_100k_from_LAION400M.py

import webdataset as wds
from PIL import Image
import io
import matplotlib.pyplot as plt
import os
import json

from warnings import filterwarnings


# os.environ["CUDA_VISIBLE_DEVICES"] = "0" # choose GPU if you are on a multi GPU server
import numpy as np
import torch
import pytorch_lightning as pl
import torch.nn as nn
from torchvision import datasets, transforms
import tqdm

from os.path import join
from datasets import load_dataset
import pandas as pd
from torch.utils.data import Dataset, DataLoader
import json

import clip
#import open_clip

from PIL import Image, ImageFile


# if you changed the MLP architecture during training, change it also here:

class MLP(pl.LightningModule):
 def __init__(self, input_size, xcol='emb', ycol='avg_rating'):
 super().__init__()
 self.input_size = input_size
 self.xcol = xcol
 self.ycol = ycol
 self.layers = nn.Sequential(
 nn.Linear(self.input_size, 1024),
 #nn.ReLU(),
 nn.Dropout(0.2),
 nn.Linear(1024, 128),
 #nn.ReLU(),
 nn.Dropout(0.2),
 nn.Linear(128, 64),
 #nn.ReLU(),
 nn.Dropout(0.1),

 nn.Linear(64, 16),
 #nn.ReLU(),

 nn.Linear(16, 1)
 )

 def forward(self, x):
 return self.layers(x)

 def training_step(self, batch, batch_idx):
 x = batch[self.xcol]
 y = batch[self.ycol].reshape(-1, 1)
 x_hat = self.layers(x)
 loss = F.mse_loss(x_hat, y)
 return loss
 
 def validation_step(self, batch, batch_idx):
 x = batch[self.xcol]
 y = batch[self.ycol].reshape(-1, 1)
 x_hat = self.layers(x)
 loss = F.mse_loss(x_hat, y)
 return loss

 def configure_optimizers(self):
 optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
 return optimizer

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)


model = MLP(768) # CLIP embedding dim is 768 for CLIP ViT L 14

s = torch.load("ava+logos-l14-linearMSE.pth") # load the model you trained previously or the model available in this repo

model.load_state_dict(s)


model.to("cuda")
model.eval()





device = "cuda" if torch.cuda.is_available() else "cpu"
model2, preprocess = clip.load("ViT-L/14", device=device) #RN50x64 



c=0
urls= []
predictions=[]

# this will run inference over 10 webdataset tar files from LAION 400M and sort them into 20 categories
# you can DL LAION 400M and convert it to wds tar files with img2dataset ( https://github.com/rom1504/img2dataset ) 


for j in range(10):
 if j<10:
 # change the path to the tar files accordingly
 dataset = wds.WebDataset("pipe:aws s3 cp s3://s-datasets/laion400m/laion400m-dat-release/0000"+str(j)+".tar -") #"pipe:aws s3 cp s3://s-datasets/laion400m/laion400m-dat-release/00625.tar -")
 else:
 dataset = wds.WebDataset("pipe:aws s3 cp s3://s-datasets/laion400m/laion400m-dat-release/000"+str(j)+".tar -") #"pipe:aws s3 cp s3://s-datasets/laion400m/laion400m-dat-release/00625.tar -")


 for i, d in enumerate(dataset):
 print(c)

 metadata= json.loads(d['json']) 

 pil_image = Image.open(io.BytesIO(d['jpg']))
 c=c+1
 try:
 image = preprocess(pil_image).unsqueeze(0).to(device)

 except:
 continue

 with torch.no_grad():
 image_features = model2.encode_image(image)


 im_emb_arr = normalized(image_features.cpu().detach().numpy() )

 prediction = model(torch.from_numpy(im_emb_arr).to(device).type(torch.cuda.FloatTensor))
 urls.append(metadata["url"])
 predictions.append(prediction)


df = pd.DataFrame(list(zip(urls, predictions)),
 columns =['filepath', 'prediction'])


buckets = [(i, i+1) for i in range(20)]


html= "<h1>Aesthetic subsets in LAION 100k samples</h1>"

i =0
for [a,b] in buckets:
 a = a/2
 b = b/2
 total_part = df[( (df["prediction"] ) *1>= a) & ( (df["prediction"] ) *1 <= b)]
 print(a,b)
 print(len(total_part) )
 count_part = len(total_part) / len(df) * 100
 estimated =int ( len(total_part) )
 part = total_part[:50]

 html+=f"<h2>In bucket {a} - {b} there is {count_part:.2f}% samples:{estimated:.2f} </h2> <div>"
 for filepath in part["filepath"]:
 html+='<img src="'+filepath +'" height="200" />'


 html+="</div>"
 i+=1
 print(i)
with open("./aesthetic_viz_laion_ava+logos_L14_100k-linearMSE.html", "w") as f:
 f.write(html)