update callback

.ipynb_checkpoints/gradio_artist_classifier-checkpoint.py

@@ -5,11 +5,12 @@ prototype
 ---
 - 2022-01-18 jkang first created
 '''
-
+from PIL import Image
 import matplotlib.pyplot as plt
 import matplotlib.image as mpimg
 import seaborn as sns
 
+import io
 import json
 import skimage.io
 from loguru import logger
@@ -35,6 +36,11 @@ artist_model = from_pretrained_keras("jkang/drawing-artist-classifier")
 trend_model = from_pretrained_keras("jkang/drawing-artistic-trend-classifier")
 logger.info('both models loaded')
 
+def load_json_as_dict(json_file):
+    with open(json_file, 'r') as f:
+        out = json.load(f)
+    return dict(out)
+
 def load_image_as_array(image_file):
     img = skimage.io.imread(image_file, as_gray=False, plugin='matplotlib')
     if (img.shape[-1] > 3) & (remove_alpha_channel):  # if RGBA
@@ -48,17 +54,63 @@ def load_image_as_tensor(image_file):
 
 def predict(input_image):
     img_3d_array = load_image_as_array(input_image)
-    img_4d_tensor = load_image_as_tensor(input_image)
+    # img_4d_tensor = load_image_as_tensor(input_image)
+    img_4d_array = img_3d_array[np.newaxis,...]
     logger.info(f'--- {input_image} loaded')
     
-    artist_model(img_4d_tensor);
-    trend_model(img_4d_tensor);
+    artist2id = load_json_as_dict(ARTIST_META)
+    trend2id = load_json_as_dict(TREND_META)
+    id2artist = {artist2id[artist]:artist for artist in artist2id}
+    id2trend = {trend2id[trend]:trend for trend in trend2id}
+    
+    # Artist model
+    a_heatmap, a_pred_id, a_pred_out = make_gradcam_heatmap(artist_model,
+                                                            img_4d_array,
+                                                            pred_idx=None)
+    a_img_pil = align_image_with_heatmap(
+        img_4d_array, a_heatmap, alpha=alpha, cmap='jet')
+    a_img = np.asarray(a_img_pil).astype('float32')/255
+    a_label = id2artist[a_pred_id]
+    a_prob = a_pred_out[a_pred_id]
+
+    # Trend model
+    t_heatmap, t_pred_id, t_pred_out = make_gradcam_heatmap(trend_model,
+                                                            img_4d_array,
+                                                            pred_idx=None)
     
-    return img_3d_array
+    t_img_pil = align_image_with_heatmap(
+        img_4d_array, t_heatmap, alpha=alpha, cmap='jet')
+    t_img = np.asarray(t_img_pil).astype('float32')/255
+    t_label = id2trend[t_pred_id]
+    t_prob = t_pred_out[t_pred_id]
+    
+    with sns.plotting_context('poster', font_scale=0.7):
+        fig, (ax1, ax2, ax3) = plt.subplots(
+            1, 3, figsize=(12, 6), facecolor='white')
+        for ax in (ax1, ax2, ax3):
+            ax.set_xticks([])
+            ax.set_yticks([])
+
+        ax1.imshow(img_3d_array)
+        ax2.imshow(a_img)
+        ax3.imshow(t_img)
+
+        ax1.set_title(f'Artist: {artist}\nTrend: {trend}', ha='left', x=0, y=1.05)
+        ax2.set_title(f'Artist Prediction:\n =>{a_label} ({a_prob:.2f})', ha='left', x=0, y=1.05)
+        ax3.set_title(f'Trend Prediction:\n =>{t_label} ({t_prob:.2f})', ha='left', x=0, y=1.05)
+        fig.tight_layout()
+        
+    buf = io.BytesIO()
+    fig.save(buf, bbox_inces='tight', fotmat='jpg')
+    buf.seek(0)
+    pil_img = Image.open(buf)
+    plt.close()
+    logger.info('--- output generated')
+    return pil_img
 
 iface = gr.Interface(
     predict,
-    title='Predict Artist and Artistic Trend of Drawings 🎨👨🏻‍🎨 (prototype)',
+    title='Predict Artist and Artistic Style of Drawings 🎨👨🏻‍🎨 (prototype)',
     description='Upload a drawing and the model will predict how likely it seems given 10 artists and their trend/style',
     inputs=[
         gr.inputs.Image(label='Upload a drawing/image', type='file')

gradio_artist_classifier.py

@@ -5,11 +5,12 @@ prototype
 ---
 - 2022-01-18 jkang first created
 '''
-
+from PIL import Image
 import matplotlib.pyplot as plt
 import matplotlib.image as mpimg
 import seaborn as sns
 
+import io
 import json
 import skimage.io
 from loguru import logger
@@ -35,6 +36,11 @@ artist_model = from_pretrained_keras("jkang/drawing-artist-classifier")
 trend_model = from_pretrained_keras("jkang/drawing-artistic-trend-classifier")
 logger.info('both models loaded')
 
+def load_json_as_dict(json_file):
+    with open(json_file, 'r') as f:
+        out = json.load(f)
+    return dict(out)
+
 def load_image_as_array(image_file):
     img = skimage.io.imread(image_file, as_gray=False, plugin='matplotlib')
     if (img.shape[-1] > 3) & (remove_alpha_channel):  # if RGBA
@@ -48,17 +54,63 @@ def load_image_as_tensor(image_file):
 
 def predict(input_image):
     img_3d_array = load_image_as_array(input_image)
-    img_4d_tensor = load_image_as_tensor(input_image)
+    # img_4d_tensor = load_image_as_tensor(input_image)
+    img_4d_array = img_3d_array[np.newaxis,...]
     logger.info(f'--- {input_image} loaded')
     
-    artist_model(img_4d_tensor);
-    trend_model(img_4d_tensor);
+    artist2id = load_json_as_dict(ARTIST_META)
+    trend2id = load_json_as_dict(TREND_META)
+    id2artist = {artist2id[artist]:artist for artist in artist2id}
+    id2trend = {trend2id[trend]:trend for trend in trend2id}
+    
+    # Artist model
+    a_heatmap, a_pred_id, a_pred_out = make_gradcam_heatmap(artist_model,
+                                                            img_4d_array,
+                                                            pred_idx=None)
+    a_img_pil = align_image_with_heatmap(
+        img_4d_array, a_heatmap, alpha=alpha, cmap='jet')
+    a_img = np.asarray(a_img_pil).astype('float32')/255
+    a_label = id2artist[a_pred_id]
+    a_prob = a_pred_out[a_pred_id]
+
+    # Trend model
+    t_heatmap, t_pred_id, t_pred_out = make_gradcam_heatmap(trend_model,
+                                                            img_4d_array,
+                                                            pred_idx=None)
     
-    return img_3d_array
+    t_img_pil = align_image_with_heatmap(
+        img_4d_array, t_heatmap, alpha=alpha, cmap='jet')
+    t_img = np.asarray(t_img_pil).astype('float32')/255
+    t_label = id2trend[t_pred_id]
+    t_prob = t_pred_out[t_pred_id]
+    
+    with sns.plotting_context('poster', font_scale=0.7):
+        fig, (ax1, ax2, ax3) = plt.subplots(
+            1, 3, figsize=(12, 6), facecolor='white')
+        for ax in (ax1, ax2, ax3):
+            ax.set_xticks([])
+            ax.set_yticks([])
+
+        ax1.imshow(img_3d_array)
+        ax2.imshow(a_img)
+        ax3.imshow(t_img)
+
+        ax1.set_title(f'Artist: {artist}\nTrend: {trend}', ha='left', x=0, y=1.05)
+        ax2.set_title(f'Artist Prediction:\n =>{a_label} ({a_prob:.2f})', ha='left', x=0, y=1.05)
+        ax3.set_title(f'Trend Prediction:\n =>{t_label} ({t_prob:.2f})', ha='left', x=0, y=1.05)
+        fig.tight_layout()
+        
+    buf = io.BytesIO()
+    fig.save(buf, bbox_inces='tight', fotmat='jpg')
+    buf.seek(0)
+    pil_img = Image.open(buf)
+    plt.close()
+    logger.info('--- output generated')
+    return pil_img
 
 iface = gr.Interface(
     predict,
-    title='Predict Artist and Artistic Trend of Drawings 🎨👨🏻‍🎨 (prototype)',
+    title='Predict Artist and Artistic Style of Drawings 🎨👨🏻‍🎨 (prototype)',
     description='Upload a drawing and the model will predict how likely it seems given 10 artists and their trend/style',
     inputs=[
         gr.inputs.Image(label='Upload a drawing/image', type='file')

requirements-dev.txt

@@ -3,7 +3,8 @@ huggingface_hub==0.4.0
 loguru==0.5.3
 matplotlib==3.5.1
 numpy==1.22.0
+Pillow==9.0.0
 scikit_image==0.19.1
 seaborn==0.11.2
-scikit-image==0.19.1
+skimage==0.0
 tensorflow==2.7.0