konerusudhir/mp_art_classification

1# -*- coding: utf-8 -*-2"""mp_art_classification.ipynb3 4Automatically generated by Colaboratory.5 6Original file is located at7    https://colab.research.google.com/drive/1mCMy50B9xHW2WdGNlxTq-wObAe-eMsQ58"""9 10import os11import shutil12import math 13import glob14import json15import pickle16import requests17import time18import re19import string20from datetime import datetime21 22import pandas as pd23import numpy as np24from PIL import Image25 26import matplotlib.pyplot as plt27 28import tensorflow as tf29 30if 'workspace/semantic_search' in os.getcwd():31    ROOT_FOLDER = os.path.join("./hf", "mp_art_classification")32else:33    ROOT_FOLDER = './'34 35PRE_TRAINED_MODELS_FOLDER = os.path.join(ROOT_FOLDER, "pre_trained_models")36TRAINED_WEIGHTS_FOLDER = os.path.join(ROOT_FOLDER, "trained_weights")37 38def clean_directories():39    shutil.rmtree(PRE_TRAINED_MODELS_FOLDER, ignore_errors=True)40 41# clean_directories()  42 43def create_directories():44    if not os.path.exists(PRE_TRAINED_MODELS_FOLDER):45        os.mkdir(PRE_TRAINED_MODELS_FOLDER)46 47create_directories()48 49from transformers import CLIPTokenizer, CLIPImageProcessor, TFCLIPTextModel, TFCLIPVisionModel50 51clip_model_id = "openai/clip-vit-large-patch14"52 53vision_model = TFCLIPVisionModel.from_pretrained(54    clip_model_id, 55    cache_dir=PRE_TRAINED_MODELS_FOLDER)56vision_processor = CLIPImageProcessor.from_pretrained(clip_model_id)57 58genre_classes_path = os.path.join(ROOT_FOLDER,'genre_class.txt')59# TSV headers [id, class]60genre_classes_df = pd.read_csv(genre_classes_path, sep = ' ',  header=None)61# print(genre_train_df.iloc[:,1])62genre_classes = []63for index, row in genre_classes_df.iterrows():64    genre_classes.append(row[1])65# print(genre_classes)66classes_count = len(genre_classes)67 68base_learning_rate = 0.000169steps_per_execution = 20070 71def create_classification_model():72    73    # Preprocess images74    inputs = tf.keras.Input(shape=(3, 224, 224))75    rescaling_layer = tf.keras.layers.Rescaling(1.0/255, offset=0.0)76    rescaled_input = rescaling_layer(inputs)77    # processed_inputs = vision_processor(images=[inputs], return_tensors="tf")78    # print(inputs)79    80    vision_model.trainable=False81    base_model_output = vision_model(rescaled_input)82    83    current_layer = base_model_output.pooler_output84    hidden_layers_nodes = [1024]85    for node_count in hidden_layers_nodes:86        hidden_layer = tf.keras.layers.Dense(node_count, activation='relu')87        dropout_layer = tf.keras.layers.Dropout(.2, input_shape=(2,))88        current_layer = hidden_layer(dropout_layer(current_layer))89    90    prediction_layer = tf.keras.layers.Dense(91        classes_count, activation='softmax')92    outputs = prediction_layer(current_layer)93    model = tf.keras.Model(inputs, outputs)94 95        96    model.compile(97        # Used leagcy optimizer due to tf 2.11 release issues with MACOS 98        # optimizer=tf.keras.optimizers.Adam(learning_rate=base_learning_rate),99        optimizer=tf.keras.optimizers.legacy.Adam(100            learning_rate=base_learning_rate),101        loss=tf.keras.losses.SparseCategoricalCrossentropy(),102        metrics=['accuracy']103        # steps_per_execution=steps_per_execution104    )105 106    return model107 108 109 110model = create_classification_model()111model.summary()112 113latest_weights = tf.train.latest_checkpoint(TRAINED_WEIGHTS_FOLDER)114model.load_weights(latest_weights)115 116# image_path = tf.constant('./hf/mp_image_search/examples/e1.jpeg')117# image = tf.io.read_file("/Users/skoneru/workspace/semantic_search/hf/mp_image_search/examples/e1.jpeg")118# print(image)119# decoded_image = tf.io.decode_image(120#     contents = image,121#     channels = 3,122#     expand_animations = False123# )124# print(decoded_image.shape)125# resized_image = tf.image.resize_with_pad(126#     image = decoded_image, 127#     target_height = 224,128#     target_width = 224,129# )130# print(resized_image.shape)131# # constant_new = tf.constant(132# #     resized_image, dtype=tf.float32, shape=(224,224,3), name='input_image'133# # )134# transposed_image = tf.transpose(135#     resized_image)136# print(transposed_image.shape)137# # constant = tf.constant(138# #     transposed_image, value_index=(3,224,224)139# # )140# # constant_new = tf.constant(141# #     transposed_image, dtype=tf.float32, shape=(3,224,224), name='input_image'142# # )143# ndarray = tf.make_ndarray(144#      tf.Variable(transposed_image, shape=(3,224,224))145# )146# # variable = tf.Variable(constant_new)147# # print(constant_new)     148# # print(inputs)149 150# image_path = './hf/mp_image_search/examples/e1.jpeg'151# img = Image.open(image_path).convert('RGB')152# desired_size =224153# old_size = img.size  # old_size[0] is in (width, height) format154# ratio = float(desired_size)/max(old_size)155# new_size = tuple([int(x*ratio) for x in old_size])156 157# img.thumbnail((desired_size, desired_size), Image.ANTIALIAS)158 159# new_im = Image.new("RGB", (desired_size, desired_size))160# new_im.paste(img, ((desired_size-new_size[0])//2,161#                     (desired_size-new_size[1])//2))162 163# # new_im.show()164# np_array = np.array(img)165# print(np_array.shape)166# transposed_np_array = np.transpose(np_array)167# print(transposed_np_array.shape)168# images_list = []169# images_list.append(transposed_np_array)170# np_input = np.asarray(images_list)171# print(np_input.shape)172# result = model.predict(np_input)173# print(result.flatten())174 175import gradio as gr176 177def process_image(input_image):178    desired_size =224179    old_size = input_image.size  # old_size[0] is in (width, height) format180    ratio = float(desired_size)/max(old_size)181    new_size = tuple([int(x*ratio) for x in old_size])182 183    input_image.thumbnail((desired_size, desired_size), Image.ANTIALIAS)184 185    new_im = Image.new("RGB", (desired_size, desired_size))186    new_im.paste(input_image, ((desired_size-new_size[0])//2,187                        (desired_size-new_size[1])//2))188 189    # new_im.show()190    np_array = np.array(input_image)191    # print(np_array.shape)192    transposed_np_array = np.transpose(np_array)193    # print(transposed_np_array.shape)194    images_list = []195    images_list.append(transposed_np_array)196    np_input = np.asarray(images_list)197    # print(np_input.shape)198    return model.predict(np_input).flatten()199 200def predict(input_image):201    # print(input_image)202    # img = Image.create(input_image)203    pil_image_object = Image.fromarray(input_image)204    probs = process_image(pil_image_object)205    return {genre_classes[i]: float(probs[i]) for i in range(len(genre_classes))}206 207image_path_prefx = os.path.join(ROOT_FOLDER,'examples')208examples = [f"{image_path_prefx}/e{n}.jpeg" for n in range(4)]209interpretation='shap'210title = "MP Art Classifier"211description = "<b>Classifies Art into 10 Genres</b>"212theme = 'grass'213 214gr.Interface(215    fn=predict,216    inputs=gr.inputs.Image(shape=((512,512))),217    outputs=gr.outputs.Label(num_top_classes=5),218    title = title,219    examples =  examples,220    theme = theme,221    interpretation = interpretation,222    description = description223).launch(debug=True)224 225clean_directories()