Update opp.py

opp.py

@@ -1,11 +1,31 @@
-import os
+# -*- coding: utf-8 -*-
+"""batch aesthetics predictor v2 - release.ipynb
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/drive/1zTrHop7pStcCwPAUP_nekK1rp6lcppYx
+"""
+
+# Commented out IPython magic to ensure Python compatibility.
+# %%capture
+# #@title Install environment & dl MLP { form-width: "100%", display-mode: "form" }
+# !pip install git+https://github.com/openai/CLIP.git
+# !pip install gradio~=3.18.0
+# #!pip install torch==1.13.1#+cu116
+# !pip install pytorch-lightning~=2.0.1
+# !wget -nc  https://huggingface.co/spaces/Seedmanc/batch-laion-aesthetic-predictor/resolve/main/sac%2Blogos%2Bava1-l14-linearMSE.pth
+
+#@title CLIP dl & init { run: "auto", vertical-output: true, form-width: "25%", display-mode: "form" }
+checkpoint = "ViT-L/14" #@param ["ViT-L/14", "ViT-L/14@336px"]
 import numpy as np
 import torch
 import pytorch_lightning as pl
 import torch.nn as nn
 import clip
-from PIL import Image, ImageFile
-import gradio as gr
+import time
+global prev_time
+global isCpu
 
 # if you changed the MLP architecture during training, change it also here:
 class MLP(pl.LightningModule):
@@ -24,45 +44,26 @@ class MLP(pl.LightningModule):
             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)
 
+
 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)
 
@@ -70,7 +71,7 @@ def load_models():
     model.to(device)
     model.eval()
 
-    model2, preprocess = clip.load("ViT-L/14", device=device)
+    model2, preprocess = clip.load(checkpoint, device=device)
 
     model_dict = {}
     model_dict['classifier'] = model
@@ -80,54 +81,421 @@ def load_models():
 
     return model_dict
 
-def predict(image):
-    image_input = model_dict['clip_preprocess'](image).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 {'aesthetic score': score}
-
 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 ''}
+    """
 
-    inputs = [gr.inputs.Image(type='pil', label='Image')]
+#@title 👁️⃤ { run: "auto", form-width: "15%" }
+global predict#or
+writeClip = False #param {type:"boolean"}
+import os
+from PIL import Image
 
-    outputs = gr.outputs.JSON()
+if writeClip: #disabled in v1
+  import torchvision
+  os.makedirs('CLIPped', exist_ok=True)
 
-    title = 'image aesthetic predictor'
 
-    examples = ['example1.jpg', 'example2.jpg', 'example3.jpg']
+def predict(image):
+  img_input = model_dict['clip_preprocess'](Image.open(image))
 
-    description = """
-    # Image Aesthetic Predictor Demo
-    This model (Image Aesthetic Predictor) is trained by LAION Team. See [https://github.com/christophschuhmann/improved-aesthetic-predictor](https://github.com/christophschuhmann/improved-aesthetic-predictor)
-    1. This model is desgined by adding five MLP layers on top of (frozen) CLIP ViT-L/14 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.
-    """
+  clipped = None
+  if writeClip:
+    clipped = img_input
+      
+  image_input = img_input.unsqueeze(0).to(model_dict['device']) #try batch
+
+  with torch.no_grad():
+      image_features = model_dict['clip_model'].encode_image(image_input) 
+
+      if model_dict['device'] == 'cuda':    # add TPU support?
+          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() #optimize?
+
+  return score, clipped
+
+#@title Wrapper & stats { form-width: "10%" }
+DEBUG = True #@param {type:"boolean"}
+autoclearLog = True #@param {type:"boolean"}
+import csv
+import sys
+import gradio as gr 
+from google.colab import output
+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): # debug only
+  if not DEBUG:
+    return x
+
+  global prev_time
+  with output.use_tags('debug', append=True):
+    print(f"<\033[94m{sys._getframe( 1).f_back.f_code.co_name}\033[96m => \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=[]):  ### TODO idk what to do
+  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()))
+
+# MAIN ################################################################
+def batch_predict(files=None, progress=gr.Progress()): #=> stats_toggle, stats_output, table_output, submit_btn 
+  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: # the solution to the interruption bug, do not remove #
+      return 
+
+    results.append([file.orig_name, round(score, 5), None]) 
+
+    if writeClip: #disabled in v1
+      clips.append((clipped, file.orig_name))
+
+    times.append(time.time() - prev_time) #simplify
+    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)): # === False if not isCpu
+      global avgScore
+      global speed
+      speed = np.mean(times)
+      avgScore = statistics(results, False)
+      eta = round(speed*(len(files)-idx+1)) # +1 or [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)  # pylint: disable=report-general-type-issues
+    
+    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:     #supposedly runs async w/o delaying the results? disabled in v1 anyway
+    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)  
+# /main ##################################################################### 
+
+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
+  # TODO can this be shortened?
+  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'))): # TODO separate folder, names?
+        os.remove(file)
+
+
+def scan(): 
+  r = ['scores.csv', 'stats.csv']  
+  return [x for x in r if os.path.isfile(x)] 
+
+# buggy as fuck
+def writeScores(table, files): # => csv_output, stats_output, stats_toggle
+  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))  # add full return statement?
+
+  if files is None:
+    log('No files, exiting writer')#
+    resetStatus('from table') # refactor
+    return   [gr.update(value=scan()), *list(showStats)]
+######
+  def writes(tbl):
+    with open('scores.csv', 'w', newline='') as f: #try tsv, json
+      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)]
+
+#@title GUI { vertical-output: true, form-width: "50%", display-mode: "both" }
+tableQueued_False = False #@param {type:"boolean"}
+queueConcurrency_2 = 2 #@param {type:"integer", min:1}
+queueUpdateInterval_0 = 0 #@param {type:"slider", min:0, max:10, step:0.2}
+prevent_thread_lock = False #@param {type:"boolean"}
+#@markdown tableQueued == True + queueConcurrency == 1 guarantees stalling on CPU
+#@markdown 
+#@markdown tableQueued - unknown effect on speed or stability
+#@markdown 
+#@markdown queueConcurrency > 1 - technically should improve speed?
+#@markdown 
+#@markdown queueUpdateInterval - in (0, 1] slows down processing, otherwise seems useless. 
+
+#@markdown prevent_thread_lock - keep the "busy cell" behavior of debug mode without it to avoid multiple instances running in parallel;
+#@markdown effects on speed & stability unknown
+if DEBUG:
+  import shutil #i doshutilsya
+
+if writeClip: # disabled in v1
+  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'): #debug only
+    for dir in dirs:
+      shutil.rmtree('../tmp/'+dir)
+    for file in files:
+      os.remove('../tmp/'+file)               #/debug
+
+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): # => submit, dataframe, status
+  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):
+      output.clear(output_tags='debug')
+    print('\r')
+      
+  clearStats() 
+  return gr.Button.update(variant='primary' if Ready else 'secondary'), None, pollStatus()
+  # ''', interactive=True''')
+
+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)
+    if DEBUG and not autoclearLog:
+      gr.Button('Clear logs').click(lambda: output.clear(output_tags='debug')) #debug
+  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') #TODO interactive
+      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.clear(reset, [imageinput], [table_output], queue=False, show_progress=True, preprocess=False)
+  imageinput.change(toggleRun, imageinput, [submit_button, table_output, status], queue=  False, cancels=[run], show_progress=False) #
+  # try .then()
+  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.dev_mode = DEBUG
+  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()
 
-    article = "<p style='text-align: center'><a href='https://laion.ai/blog/laion-aesthetics/'>LAION aesthetics blog post</a></p>"
-
-    with gr.Blocks() as demo:
-        gr.Markdown(description)
-        with gr.Row():
-            with gr.Column():
-                image_input = gr.Image(type='pil', label='Input image')
-                submit_button = gr.Button('Submit')
-            json_output = gr.JSON(label='Output')
-        submit_button.click(predict, inputs=image_input, outputs=json_output)
-        gr.Examples(examples=examples, inputs=image_input)
-        gr.HTML(article)
-    demo.launch()
+#demo.close()