{"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"},{"sourceId":1102,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":949,"modelId":117}],"dockerImageVersionId":30762,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Fast Neural Style Transfer (google/arbitrary-image-stylization-v1) \nref:https://arxiv.org/abs/1705.06830","metadata":{}},{"cell_type":"code","source":"import tensorflow as tf","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:05.453360Z","iopub.execute_input":"2024-09-04T16:34:05.453797Z","iopub.status.idle":"2024-09-04T16:34:05.458728Z","shell.execute_reply.started":"2024-09-04T16:34:05.453759Z","shell.execute_reply":"2024-09-04T16:34:05.457599Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow_hub as hub\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport os\nimport glob\nfrom pathlib import Path\nimport shutil\nimport numpy as np\nimport operator\nfrom dataclasses import dataclass\nimport time\nfrom datetime import timedelta","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:05.660763Z","iopub.execute_input":"2024-09-04T16:34:05.661144Z","iopub.status.idle":"2024-09-04T16:34:05.666915Z","shell.execute_reply.started":"2024-09-04T16:34:05.661105Z","shell.execute_reply":"2024-09-04T16:34:05.665816Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Import Model from Model Hub","metadata":{}},{"cell_type":"code","source":"@dataclass\nclass G:\n    env=\"kaggle\"\n    if env == \"kaggle\":\n        style_img_dir_path=\"/kaggle/input/gan-getting-started/monet_jpg\"\n        content_img_dir_path=\"/kaggle/input/gan-getting-started/photo_jpg\"\n        model = hub.load('/kaggle/input/arbitrary-image-stylization-v1/tensorflow1/256/2')\n    else:\n        model  = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')\n        style_img_dir_path = os.getcwd()+\"/monet_jpg\"\n        content_img_dir_path=os.getcwd()+\"/photo_jpg\"","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:06.028492Z","iopub.execute_input":"2024-09-04T16:34:06.029216Z","iopub.status.idle":"2024-09-04T16:34:08.285642Z","shell.execute_reply.started":"2024-09-04T16:34:06.029173Z","shell.execute_reply":"2024-09-04T16:34:08.284802Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Image","metadata":{}},{"cell_type":"code","source":"def load_img(path_to_img):\n    image = tf.io.read_file(path_to_img)\n    image = tf.image.decode_jpeg(image)\n    image = tf.image.convert_image_dtype(image, tf.uint8)\n\n    return image\n\ndef load_images(content_path, style_path):\n    content_image = load_img(\"{}\".format(content_path))\n    style_image = load_img(\"{}\".format(style_path))\n\n    return content_image, style_image","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:08.287555Z","iopub.execute_input":"2024-09-04T16:34:08.287968Z","iopub.status.idle":"2024-09-04T16:34:08.294251Z","shell.execute_reply.started":"2024-09-04T16:34:08.287923Z","shell.execute_reply":"2024-09-04T16:34:08.293155Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Transfer Style","metadata":{}},{"cell_type":"code","source":"def load_style_ds(style_img_dir):\n    ds=tf.keras.preprocessing.image_dataset_from_directory(style_img_dir,labels=None, color_mode='grayscale',\n                                                          batch_size=100)\n    ds=ds.map(lambda image: tf.image.convert_image_dtype(image, tf.uint8))\n    return ds","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:08.295593Z","iopub.execute_input":"2024-09-04T16:34:08.295954Z","iopub.status.idle":"2024-09-04T16:34:08.305349Z","shell.execute_reply.started":"2024-09-04T16:34:08.295908Z","shell.execute_reply":"2024-09-04T16:34:08.304369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def calculate_ssmin(content_image, style_image):\n    ssim_score= tf.image.ssim(content_image, style_image,256)\n    return ssim_score","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:08.307670Z","iopub.execute_input":"2024-09-04T16:34:08.308376Z","iopub.status.idle":"2024-09-04T16:34:08.315251Z","shell.execute_reply.started":"2024-09-04T16:34:08.308300Z","shell.execute_reply":"2024-09-04T16:34:08.314292Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trial=load_style_ds(G.style_img_dir_path)","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:08.316515Z","iopub.execute_input":"2024-09-04T16:34:08.316811Z","iopub.status.idle":"2024-09-04T16:34:08.500989Z","shell.execute_reply.started":"2024-09-04T16:34:08.316780Z","shell.execute_reply":"2024-09-04T16:34:08.500155Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def find_matches(content_img_path,style_img_ds, verbose):\n    scores={}\n    starttime = time.perf_counter()\n    \n    content_image=load_img(\"{}\".format(content_img_path))\n    content_image=tf.image.rgb_to_grayscale(content_image)\n    style_image_paths=list(os.walk(G.style_img_dir_path))[0][2]\n    for images in style_img_ds:\n        for style_image_path,style_img in zip(style_image_paths,images):\n            ssim_score=calculate_ssmin(content_image, style_img)\n            ssim_score=ssim_score.numpy() \n            style_image_path_full=G.style_img_dir_path+\"/\"+style_image_path\n            scores[style_image_path_full]=abs(ssim_score)\n    \n    match=max(scores.items(), key=operator.itemgetter(1))[0]\n    duration = timedelta(seconds=time.perf_counter()-starttime)\n    if verbose:\n        print(\"computation time for finding match: \", duration)\n    return match","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:08.502094Z","iopub.execute_input":"2024-09-04T16:34:08.502405Z","iopub.status.idle":"2024-09-04T16:34:08.510524Z","shell.execute_reply.started":"2024-09-04T16:34:08.502373Z","shell.execute_reply":"2024-09-04T16:34:08.509639Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def imshow(image, title=None):\n    if len(image.shape) > 3:\n        image = tf.squeeze(image, axis=0)\n\n    plt.imshow(image)\n    if title:\n        plt.title(title)","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:08.935482Z","iopub.execute_input":"2024-09-04T16:34:08.935854Z","iopub.status.idle":"2024-09-04T16:34:08.941008Z","shell.execute_reply.started":"2024-09-04T16:34:08.935820Z","shell.execute_reply":"2024-09-04T16:34:08.940051Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def predict(content_dir, style_img_dir, path=\"/kaggle/working/tmp/\"):\n    if os.path.isdir(path):\n        shutil.rmtree(path)\n    Path(path).mkdir(parents=True, exist_ok=True)\n    style_image_ds=load_style_ds(style_img_dir).take(1)\n    print_sample=5\n    for i, content_img_path in enumerate(content_dir):\n        if i<print_sample:\n            style_img_path=find_matches(content_img_path,style_image_ds, verbose=True)\n        else:\n            style_img_path=find_matches(content_img_path,style_image_ds, verbose=False)\n        \n        if i<print_sample:\n            starttime = time.perf_counter()\n        content_image, style_image = load_images(content_img_path, style_img_path)\n        if i<print_sample:\n            duration = timedelta(seconds=time.perf_counter()-starttime)\n            print(\"computation time for style and content images loading: \", duration)\n        if i<print_sample:\n            plt.subplot(1, 3, 1)\n            imshow(content_image, 'Content Image')\n\n            plt.subplot(1, 3, 2)\n            imshow(style_image, 'Style Image')\n        if i<print_sample:\n            starttime = time.perf_counter()\n        stylized_image = G.model(tf.image.convert_image_dtype(content_image[tf.newaxis, :], tf.float32),\n                            tf.image.convert_image_dtype(style_image[tf.newaxis, :], tf.float32))[0]\n        if i<print_sample:\n            duration = timedelta(seconds=time.perf_counter()-starttime)\n            print(\"computation time for style transfer: \", duration)\n        \n        tensor_shape=tf.shape(stylized_image)\n        number_elem_shape = tf.shape(tensor_shape)\n        if number_elem_shape > 3:\n            assert tensor_shape[0] == 1\n            stylized_image = stylized_image[0]\n            \n        if i<print_sample:\n            plt.subplot(1, 3, 3)\n            imshow(stylized_image, 'Stylized Image')\n            plt.show()\n            plt.close()\n        if i<print_sample:\n            starttime = time.perf_counter()\n        im=tf.keras.preprocessing.image.array_to_img(stylized_image)\n        im.save(path + str(i) + \".jpg\")\n        if i<print_sample:\n            duration = timedelta(seconds=time.perf_counter()-starttime)\n            print(\"time it took to save image: \", duration)\n            print(\"*\"*66)\n\n\nif G.env==\"kaggle\":\n    content_dir = glob.glob(os.path.join(\"/kaggle/input/gan-getting-started/photo_jpg\", \"*.jpg\"))\n    predict(content_dir, G.style_img_dir_path)\nelse:\n    content_dir = glob.glob(os.path.join(os.getcwd()+\"/photo_jpg\", \"*.jpg\"))\n    predict(content_dir, G.style_img_dir_path, path=os.getcwd()+\"/tmp/\")","metadata":{"execution":{"iopub.status.busy":"2024-09-04T16:34:09.244615Z","iopub.execute_input":"2024-09-04T16:34:09.245011Z","iopub.status.idle":"2024-09-04T16:34:51.306790Z","shell.execute_reply.started":"2024-09-04T16:34:09.244972Z","shell.execute_reply":"2024-09-04T16:34:51.305430Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nif G.env==\"kaggle\":\n    shutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/working/tmp/\")\nelse:\n    shutil.make_archive(os.getcwd()+\"/images\", 'zip', os.getcwd()+\"/tmp/\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}