Fix logits parsing

gradio-ocsai-d.py

@@ -1,28 +1,36 @@
 from pathlib import Path
-from transformers import pipeline
 from PIL import Image
 import gradio as gr
 import pandas as pd
 import numpy as np
 import altair as alt
+from transformers import AutoModelForImageClassification, AutoImageProcessor
 
+modelname = "POrg/ocsai-d-web"
+model = AutoModelForImageClassification.from_pretrained(modelname)
+model.eval()  # Set the model to evaluation mode
+image_processor = AutoImageProcessor.from_pretrained(modelname)
 
 prompt_images = {
-    "Images_11": "blanks/Images_11_blank.png",
-    "Images_4": "blanks/Images_4_blank.png",
-    "Images_17": "blanks/Images_17_blank.png",
-    "Images_9": "blanks/Images_9_blank.png",
-    "Images_3": "blanks/Images_3_blank.png",
-    "Images_8": "blanks/Images_8_blank.png",
-    "Images_13": "blanks/Images_13_blank.png",
-    "Images_15": "blanks/Images_15_blank.png",
-    "Images_12": "blanks/Images_12_blank.png",
-    "Images_7": "blanks/Images_7_blank.png",
-    "Images_56": "blanks/Images_56_blank.png",
-    "Images_19": "blanks/Images_19_blank.png",
+    "Images_11": "./blanks/Images_11_blank.png",
+    "Images_4": "./blanks/Images_4_blank.png",
+    "Images_17": "./blanks/Images_17_blank.png",
+    "Images_9": "./blanks/Images_9_blank.png",
+    "Images_3": "./blanks/Images_3_blank.png",
+    "Images_8": "./blanks/Images_8_blank.png",
+    "Images_13": "./blanks/Images_13_blank.png",
+    "Images_15": "./blanks/Images_15_blank.png",
+    "Images_12": "./blanks/Images_12_blank.png",
+    "Images_7": "./blanks/Images_7_blank.png",
+    "Images_56": "./blanks/Images_56_blank.png",
+    "Images_19": "./blanks/Images_19_blank.png",
 }
 
 dist = pd.read_csv('./score_norm_distribution.csv', dtype=float)
+base_chart = alt.Chart(dist).mark_line().encode(
+    x='percentile',
+    y='score_norm'
+)
 
 
 def get_percentile(score):
@@ -35,41 +43,43 @@ def inverse_scale(logits):
     return logits * (scaler_params['range']) + scaler_params['min']
 
 
-def get_predictions(img_dict):
-    # gradio passes a dictionary with background, composite, and layers
-    # the composite is what we want
-    img = img_dict['composite']
-    predictions = classifier(img)
+def get_predictions(img):
+    inputs = image_processor(img, return_tensors="pt")
+    prediction = model(**inputs)
+    score = prediction.logits[0].detach().numpy()[0]
+    score = min(max(score, 0), 1)
     return {
-        'originality' : np.round(predictions[0]['score'], 2),
-        'jrt' : np.round(inverse_scale(0), 2),
-        'percentile': get_percentile(predictions[0]['score'])
+        'originality': np.round(score, 2),
+        'jrt': np.round(inverse_scale(0), 2),
+        'percentile': get_percentile(score)
     }
 
 
-base_chart = alt.Chart(dist).mark_line().encode(
-    x='percentile',
-    y='score_norm'
-)
-
-def classify_image(img_dict):
+def classify_image(img_dict: dict):
     # gradio passes a dictionary with background, composite, and layers
     # the composite is what we want
-    img = img_dict['composite']
-    p = get_predictions(img_dict)
+    img = img_dict['composite'].convert('RGB')
+    p = get_predictions(img)
 
-    percentile_mark = alt.Chart(pd.DataFrame({'y': [p['originality']]})).mark_rule(color='red').encode(y='y')
+    score_round = str(np.round(p['originality'], 2))
+    label = f"Percentile: {p['percentile']}; Normalized Score: {score_round}"
+    label_df = pd.DataFrame({'y': [p['originality']], 'text': [label]})
 
-    # Text annotation for the percentile mark
-    text = alt.Chart(pd.DataFrame({'y': [p['originality']], 'text': [f"Percentile: {p['percentile']}; Normalized Score: {p['originality']}"]})).mark_text(
+    rule = alt.Chart(label_df).mark_rule(
+        color='red'
+    ).encode(
+        y='y'
+    )
+    
+    txt = alt.Chart(label_df).mark_text(
         align='left',
         baseline='middle',
-        dx=7, dy=-8  # Nudges text to right so it doesn't overlap with the line
+        dx=0, dy=-8
     ).encode(
         y='y',
         text='text'
     )
-    return base_chart + percentile_mark + text
+    return base_chart + rule + txt
 
 
 def update_editor(background, img_editor):
@@ -80,8 +90,6 @@ def update_editor(background, img_editor):
     return img_editor
 
 
-classifier = pipeline("image-classification", model='POrg/ocsai-d-large')
-
 editor = gr.ImageEditor(type='pil',
                         value=dict(
                             background=Image.open(prompt_images['Images_11']),
@@ -106,7 +114,7 @@ demo = gr.Interface(fn=classify_image,
                      inputs=[editor],
                      outputs=gr.Plot(),
                      title="Ocsai-D",
-                     description="Complete the drawing and classify the originality. Examples are from MTCI ([Barbot 2018](https://pubmed.ncbi.nlm.nih.gov/30618952/)). Choose the brush icon below the image to start editing.",
+                     description="Complete the drawing and classify the originality. Choose the brush icon below the image to start editing.\n\nModel from *A Comparison of Supervised and Unsupervised Learning Methods in Automated Scoring of Figural Tests of Creativity* ([preprint](http://dx.doi.org/10.13140/RG.2.2.26865.25444)).\n\nExamples are from MTCI ([Barbot 2018](https://pubmed.ncbi.nlm.nih.gov/30618952/)).",
                      examples=examples
                      )