|
|
|
"""batch aesthetics predictor v2 - release.ipynb |
|
|
|
Automatically generated by Colaboratory. |
|
|
|
Original file is located at |
|
https://colab.research.google.com/drive/1zTrHop7pStcCwPAUP_nekK1rp6lcppYx |
|
""" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
checkpoint = "ViT-L/14" |
|
import numpy as np |
|
import torch |
|
import pytorch_lightning as pl |
|
import torch.nn as nn |
|
import clip |
|
import time |
|
global prev_time |
|
global isCpu |
|
|
|
|
|
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.Dropout(0.2), |
|
nn.Linear(1024, 128), |
|
|
|
nn.Dropout(0.2), |
|
nn.Linear(128, 64), |
|
|
|
nn.Dropout(0.1), |
|
nn.Linear(64, 16), |
|
|
|
nn.Linear(16, 1) |
|
) |
|
|
|
def forward(self, x): |
|
return self.layers(x) |
|
|
|
def normalized(a, axis=-1, order=2): |
|
l2 = np.atleast_1d(np.linalg.norm(a, order, axis)) |
|
l2[l2 == 0] = 1 |
|
return a / np.expand_dims(l2, axis) |
|
|
|
|
|
def load_models(): |
|
model = MLP(768) |
|
global device |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
global isCpu |
|
isCpu = device == "cpu" |
|
|
|
s = torch.load("sac+logos+ava1-l14-linearMSE.pth", map_location=device) |
|
|
|
model.load_state_dict(s) |
|
model.to(device) |
|
model.eval() |
|
|
|
model2, preprocess = clip.load(checkpoint, device=device) |
|
|
|
model_dict = {} |
|
model_dict['classifier'] = model |
|
model_dict['clip_model'] = model2 |
|
model_dict['clip_preprocess'] = preprocess |
|
model_dict['device'] = device |
|
|
|
return model_dict |
|
|
|
if __name__ == '__main__': |
|
print('\tinit models') |
|
|
|
global model_dict |
|
prev_time = time.time() |
|
|
|
model_dict = load_models() |
|
print('model load', time.time() - prev_time) |
|
description = f""" |
|
## Batch Image Aesthetic Predictor |
|
0. Based on https://huggingface.co/spaces/Geonmo/laion-aesthetic-predictor, I just expanded the GUI & added stats. |
|
1. This model is designed by adding five MLP layers on top of (frozen) CLIP <u>**{checkpoint}**</u> checkpoint and only the MLP layers are fine-tuned with a lot of images by a regression loss term such as MSE and MAE. |
|
2. Output is bounded from 0 to 10. The higher the better. |
|
3. The MLP being used currently is: **sac+logos+ava1-l14-linearMSE.pth** trained on 224x224 images. |
|
4. Running on **{device}**{', be patient. Progressive output & immediate stats are available.' if isCpu else '. Batch mode enabled, results after completion.'} |
|
5. Please don't click 'Submit' again during the processing, it'll mess things up. To stop processing, clear the file input. If the results are missing from the stats or export areas at the end, sort the table by any header & wait. |
|
{'6. The MLP was not retrained for this CLIP checkpoint, correct results are not guaranteed. It is also 2x slower.' if checkpoint != "ViT-L/14" else ''} |
|
""" |
|
|
|
|
|
global predict |
|
writeClip = False |
|
import os |
|
from PIL import Image |
|
|
|
if writeClip: |
|
import torchvision |
|
os.makedirs('CLIPped', exist_ok=True) |
|
|
|
|
|
def predict(image): |
|
img_input = model_dict['clip_preprocess'](Image.open(image)) |
|
|
|
clipped = None |
|
if writeClip: |
|
clipped = img_input |
|
|
|
image_input = img_input.unsqueeze(0).to(model_dict['device']) |
|
|
|
with torch.no_grad(): |
|
image_features = model_dict['clip_model'].encode_image(image_input) |
|
|
|
if model_dict['device'] == 'cuda': |
|
im_emb_arr = normalized(image_features.detach().cpu().numpy()) |
|
im_emb = torch.from_numpy(im_emb_arr).to(model_dict['device']).type(torch.cuda.FloatTensor) |
|
else: |
|
im_emb_arr = normalized(image_features.detach().numpy()) |
|
im_emb = torch.from_numpy(im_emb_arr).to(model_dict['device']).type(torch.FloatTensor) |
|
|
|
prediction = model_dict['classifier'](im_emb) |
|
score = prediction.item() |
|
|
|
return score, clipped |
|
|
|
|
|
DEBUG = False |
|
autoclearLog = True |
|
import csv |
|
import sys |
|
import gradio as gr |
|
if DEBUG: print(gr.__version__) |
|
|
|
def defStats(): |
|
return {'Max':{}, 'Max - min': {}} |
|
|
|
global Ready |
|
global avgScore |
|
global eta |
|
global speed |
|
global canPoll |
|
canPoll=Ready =False |
|
eta = avgScore = None |
|
speed = 0 |
|
|
|
Stats = defStats() |
|
global default_mode |
|
default_mode = list(Stats.keys())[1] |
|
|
|
def log(x = '', y = None): |
|
if not DEBUG: |
|
return x |
|
|
|
global prev_time |
|
print(f"<\033[97m{sys._getframe().f_back.f_code.co_name}\033[96m>:") |
|
if x: |
|
print(time.strftime('%M:%S '), x, round(time.time() - prev_time, 3), '\033[0m') |
|
if y: |
|
print(' extra: ', y, '\033[0m') |
|
prev_time = time.time() |
|
|
|
return x |
|
|
|
|
|
def pollStatus(table=[]): |
|
time.sleep(1) |
|
spd = speed and (f'{round(speed,1)} s/f' if speed >= 1 else f'{round(1/speed,1)} f/s') |
|
stext = f' Avg speed: {spd}.' if speed else '' |
|
etext = f' ETA: {eta} {"s." if type(eta) == int else ""}' if eta else '' |
|
atext = f'Running average: {avgScore}.' if avgScore else '' |
|
|
|
return f"[{time.strftime('%M:%S')}] {' Ready.' if not atext else ''} {atext} {etext} {stext}" if canPoll else 'idle' |
|
|
|
|
|
def switch_stats(mode): |
|
global default_mode |
|
default_mode = mode if mode else 'Max' |
|
|
|
return Stats[default_mode] |
|
|
|
|
|
def writeStats(labels): |
|
with open('stats.csv', 'w', newline='') as f: |
|
writer = csv.writer(f) |
|
log('actual write stats', labels and labels.values()) |
|
writer.writerow(gr.utils.sanitize_list_for_csv(labels.keys())) |
|
writer.writerow(gr.utils.sanitize_list_for_csv(labels.values())) |
|
|
|
|
|
def batch_predict(files=None, progress=gr.Progress()): |
|
run_time = time.time() |
|
if files and len(files) > 1: |
|
global eta |
|
eta = 'calculating...' |
|
|
|
results = list() |
|
log('has file(s)?', files and files[0]) |
|
global Stats |
|
global Ready |
|
Stats = defStats() |
|
|
|
if files is None: |
|
log('empty load') |
|
yield gr.update(), None, None, gr.update() |
|
log('ABORT') |
|
return |
|
else: |
|
maxSteps = min(len(files), 3) if isCpu else len(files) |
|
log('good2go') |
|
yield gr.update(visible=False), gr.update(visible=False), None, gr.update(variant="secondary") |
|
|
|
progress((1, maxSteps), unit='', desc='Importing...') |
|
clearStats() |
|
|
|
log('start the main loop') |
|
|
|
times=list() |
|
clips=list() |
|
|
|
for idx,file in enumerate(files, start=1): |
|
prev_time = time.time() |
|
score, clipped = predict(file) |
|
|
|
if not Ready: |
|
return |
|
|
|
results.append([file.orig_name, round(score, 5), None]) |
|
|
|
if writeClip: |
|
clips.append((clipped, file.orig_name)) |
|
|
|
times.append(time.time() - prev_time) |
|
asyncThreshold = 1 if isCpu else len(files)-1 |
|
|
|
if (idx <= asyncThreshold): |
|
progress((idx+1, maxSteps), unit='', desc='Starting...' if isCpu else 'Working...') |
|
|
|
if (idx > asyncThreshold) and (idx < len(files)): |
|
global avgScore |
|
global speed |
|
speed = np.mean(times) |
|
avgScore = statistics(results, False) |
|
eta = round(speed*(len(files)-idx+1)) |
|
log(idx) |
|
yield gr.update(), None, results, gr.update() |
|
|
|
table_data = results |
|
if DEBUG: print('RUN time', time.time() - run_time, 'avg', np.mean(times)) |
|
|
|
if len(results) > 1: |
|
eta = 'finishing...' |
|
log('finishing') |
|
stats = statistics(results) |
|
|
|
for i, row in enumerate(table_data): |
|
table_data[i][2] = round((row[1] - stats['AVG'])**2, 4) |
|
|
|
writeStats(stats) |
|
log('|2|', table_data) |
|
yield gr.update(visible=True), gr.update(value=switch_stats(default_mode), visible=True), table_data, gr.update(variant="primary") |
|
else: |
|
log('I', table_data) |
|
yield gr.update(visible=False), gr.update(value=None, visible=False), table_data, gr.update(variant="primary") |
|
|
|
avgScore = None |
|
if writeClip: |
|
log('beforeWrite') |
|
for c,f in clips: |
|
torchvision.utils.save_image(c, 'CLIPped/'+f+'.png', normalize=True) |
|
log('afterWrite') |
|
|
|
log('Exit main loop') |
|
speed = (time.time() - run_time)/len(files) |
|
|
|
|
|
def statistics(results, full=True): |
|
array = np.array(results).T[1].astype(float) |
|
|
|
max = np.max(array) |
|
avg = round(array.mean(), 3) |
|
if (not full): return avg |
|
med = round(np.median(array), 3) |
|
min = array.min() |
|
std = round(array.std(), 4) |
|
cov = round(std/avg*100, 2) |
|
rng = round(max-min, 3) |
|
range = max-min |
|
|
|
Stats['Max'][f'MAX: {round(max, 3)}'] = 1 |
|
Stats['Max'][f'min: {round(min, 3)}'] = min/max |
|
Stats['Max'][f"CoV: {cov}%"] = std/max |
|
Stats['Max'][f'AVG: {avg}'] = avg/max |
|
Stats['Max'][f'Med: {med}'] = med/max |
|
Stats['Max'][f'M-m: {rng}'] = range/max |
|
|
|
if (range == 0): |
|
range = 1 |
|
Stats['Max - min'][f'MAX: {round(max, 3)}'] = 1 |
|
Stats['Max - min'][f'min: {round(min, 3)}'] = 0 |
|
Stats['Max - min'][f"CoV: {cov}%"] = std/range |
|
Stats['Max - min'][f'AVG: {avg}'] = (avg-min)/range |
|
Stats['Max - min'][f'Med: {med}'] = (med-min)/range |
|
Stats['Max - min'][f'M-m: {rng}'] = rng/max |
|
|
|
return dict(zip(('AVG','CoV','M-m','min','Med','MAX'), (avg, cov, rng, round(min,3), med, round(max,3)))) |
|
|
|
|
|
def clearStats(): |
|
log('clst too many calls?') |
|
for root, dirs, files in os.walk('.'): |
|
for file in files: |
|
if (file.startswith(('scores','stats'))): |
|
os.remove(file) |
|
|
|
|
|
def scan(): |
|
r = ['scores.csv', 'stats.csv'] |
|
return [x for x in r if os.path.isfile(x)] |
|
|
|
|
|
def writeScores(table, files): |
|
statsVisible = False |
|
rows = table and table['data'] |
|
log('Entering the scores writer', 'from table change' if files and table else None) |
|
showStats = (gr.update(visible=statsVisible) for x in range(0,2)) |
|
|
|
if files is None: |
|
log('No files, exiting writer') |
|
resetStatus('from table') |
|
return [gr.update(value=scan()), *list(showStats)] |
|
|
|
def writes(tbl): |
|
with open('scores.csv', 'w', newline='') as f: |
|
writer = csv.writer(f) |
|
log('Actual saving scores', len(tbl['data'])) |
|
writer.writerow(gr.utils.sanitize_list_for_csv(tbl['headers'])) |
|
writer.writerows(gr.utils.sanitize_list_for_csv(tbl['data'])) |
|
|
|
if table and any([x for x in rows[0]]): |
|
if (len(rows) > 1): |
|
statsVisible = len(rows) >= len(files) |
|
|
|
if statsVisible: |
|
writes(table) |
|
log('Updating two', 'finished') |
|
global eta |
|
eta = 0 |
|
return [gr.update(value=scan()), *list(showStats)] |
|
|
|
else: |
|
statsVisible = False |
|
if (len(files) == 1): |
|
writes(table) |
|
log('updating 1') |
|
return [gr.update(value=scan()), *list(showStats)] |
|
|
|
log('Not ready for writing yet, exiting.', f'total files: {files and len(files)}, but ready rows: {rows and len(rows)}') |
|
return [gr.update(value=scan()), *list(showStats)] |
|
|
|
|
|
tableQueued_False = False |
|
queueConcurrency_2 = 10 |
|
queueUpdateInterval_0 = 0 |
|
prevent_thread_lock = False |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if DEBUG: |
|
import shutil |
|
import subprocess |
|
|
|
if writeClip: |
|
for root, dirs, files in os.walk('CLIPped'): |
|
for file in files: |
|
os.remove('CLIPped/'+file) |
|
|
|
if DEBUG: |
|
for root, dirs, files in os.walk('../tmp'): |
|
for dir in dirs: |
|
shutil.rmtree('../tmp/'+dir) |
|
for file in files: |
|
os.remove('../tmp/'+file) |
|
|
|
def resetStatus(msg = 'clear'): |
|
global avgScore |
|
global eta |
|
global speed |
|
avgScore = None |
|
eta = None |
|
speed = 0 |
|
log(msg) |
|
if msg != 'clear': |
|
clearStats() |
|
print('\n') |
|
|
|
|
|
Css = ''' |
|
#lbl .output-class { |
|
background-color: transparent; |
|
max-height: 0; |
|
color: transparent; |
|
padding: var(--size-3); |
|
} |
|
#add_img .file-preview .file td:first-child { |
|
overflow-wrap: anywhere; |
|
} |
|
#csv_out .file-preview { |
|
margin-bottom: var(--size-4); |
|
overflow-x: visible; |
|
} |
|
#tbl_out tbody .cell-wrap:first-child { |
|
overflow-wrap: anywhere; |
|
} |
|
button#sbmt:focus:not(:active) { |
|
opacity: 0.75; |
|
pointer-events: none; |
|
} |
|
#mid_col :not(#csv_out) .wrap.default { |
|
opacity: 0!important; |
|
} |
|
''' |
|
|
|
def toggleRun(files): |
|
global Ready |
|
Ready = files is not None |
|
log('Toggle', Ready) |
|
global canPoll |
|
canPoll = Ready |
|
|
|
if not Ready: |
|
if eta: |
|
log('INTERRUPTED at ss remaining (extra)', eta) |
|
resetStatus() |
|
if DEBUG and autoclearLog: |
|
subprocess.call('clear') |
|
print('\r') |
|
|
|
clearStats() |
|
return gr.Button.update(variant='primary' if Ready else 'secondary'), None, pollStatus() |
|
|
|
|
|
log('GUI start') |
|
|
|
blks = gr.Blocks(analytics_enabled=False, title="Batch Image Aesthetic Predictor", css=Css) |
|
with blks as demo: |
|
with gr.Accordion('README', open=False): |
|
gr.Markdown(description) |
|
|
|
|
|
with gr.Row().style(equal_height=False): |
|
with gr.Column(scale=2): |
|
imageinput = gr.Files(file_types=["image"], label="Add images", elem_id="addimg") |
|
submit_button = gr.Button('Submit', variant="secondary", elem_id='sbmt') |
|
with gr.Column(variant="compact", min_width=256, elem_id="mid_col"): |
|
stats_toggle = gr.Radio(list(Stats.keys()), show_label=True, label='Stats relative to:', value=default_mode, visible=False) |
|
stats_output = gr.Label(label='Stats', visible=False, elem_id="lbl") |
|
csv_output = gr.File( label="Export", elem_id='csv_out' ) |
|
with gr.Column(scale=2): |
|
table_output = gr.Dataframe(headers=['Image', 'Score', 'MSE'], max_rows=15, overflow_row_behaviour="paginate", interactive=False, wrap=True, elem_id="tbl_out") |
|
|
|
status = gr.Textbox(pollStatus(), max_lines=1, show_label=False, placeholder='Status bar').style(container=False) |
|
status.change(pollStatus, None, status, show_progress= False, queue=False) |
|
|
|
tch = table_output.change(writeScores, [table_output, imageinput], [csv_output, stats_output, stats_toggle], preprocess=False, queue= tableQueued_False, show_progress=not isCpu) |
|
stats_toggle.change(switch_stats, [stats_toggle], [stats_output], queue=False, show_progress=False) |
|
|
|
run = submit_button.click(batch_predict, imageinput, [stats_toggle, stats_output, table_output, submit_button], queue=True, scroll_to_output=True) |
|
|
|
imageinput.change(toggleRun, imageinput, [submit_button, table_output, status], queue= False, cancels=[run], show_progress=False) |
|
|
|
if DEBUG: |
|
demo.load(lambda: log('load'), queue=not True, show_progress=False) |
|
|
|
demo.queue(api_open= not DEBUG, status_update_rate='auto' if queueUpdateInterval_0 == 0 else queueUpdateInterval_0 , concurrency_count=max(queueConcurrency_2, 1)) |
|
log('Prelaunch') |
|
|
|
|
|
demo.launch(debug=DEBUG, quiet= not DEBUG, show_error= True, prevent_thread_lock=prevent_thread_lock, height=768) |
|
if (prevent_thread_lock and not DEBUG): demo.block_thread() |
|
|
|
|