merged changes

.gitignore

@@ -0,0 +1 @@
+__pycache__/

app.py

@@ -1,4 +1,4 @@
-#importing the libraries
+# Import general purpose libraries
 import os, sys, re
 import streamlit as st
 import PIL
@@ -6,21 +6,8 @@ from PIL import Image
 import cv2
 import numpy as np
 import uuid
-
-import ssl
-ssl._create_default_https_context = ssl._create_unverified_context
-
-# Import torch libraries
-import fastai
-import torch
-
-# Import util functions from app_utils
-from app_utils import download
-from app_utils import generate_random_filename
-from app_utils import clean_me
-from app_utils import clean_all
-from app_utils import get_model_bin
-from app_utils import convertToJPG
+from zipfile import ZipFile, ZIP_DEFLATED
+from io import BytesIO
 
 # Import util functions from deoldify
 # NOTE:  This must be the first call in order to work properly!
@@ -30,13 +17,17 @@ from deoldify.device_id import DeviceId
 device.set(device=DeviceId.CPU)
 from deoldify.visualize import *
 
+# Import util functions from app_utils
+from app_utils import get_model_bin
+
+
 
 ####### INPUT PARAMS ###########
 model_folder = 'models/'
 max_img_size = 800
 ################################
 
-@st.cache(allow_output_mutation=True)
+@st.cache(allow_output_mutation=True, show_spinner=False)
 def load_model(model_dir, option):
     if option.lower() == 'artistic':
         model_url = 'https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth'
@@ -68,42 +59,132 @@ def resize_img(input_img, max_size):
 
     return img
 
-def get_image_download_link(img,filename,text):
+def get_image_download_link(img, filename, button_text):
     button_uuid = str(uuid.uuid4()).replace('-', '')
     button_id = re.sub('\d+', '', button_uuid)
-    
-    custom_css = f""" 
-        <style>
-            #{button_id} {{
-                background-color: rgb(255, 255, 255);
-                color: rgb(38, 39, 48);
-                padding: 0.25em 0.38em;
-                position: relative;
-                text-decoration: none;
-                border-radius: 4px;
-                border-width: 1px;
-                border-style: solid;
-                border-color: rgb(230, 234, 241);
-                border-image: initial;
-
-            }} 
-            #{button_id}:hover {{
-                border-color: rgb(246, 51, 102);
-                color: rgb(246, 51, 102);
-            }}
-            #{button_id}:active {{
-                box-shadow: none;
-                background-color: rgb(246, 51, 102);
-                color: white;
-                }}
-        </style> """
 
     buffered = BytesIO()
     img.save(buffered, format="JPEG")
     img_str = base64.b64encode(buffered.getvalue()).decode()
-    href =  custom_css + f'<a href="data:file/txt;base64,{img_str}" id="{button_id}" download="{filename}">{text}</a>'
+
+    return get_button_html_code(img_str, filename, 'txt', button_id, button_text)
+
+def get_button_html_code(data_str, filename, filetype, button_id, button_txt='Download file'):
+    custom_css = f""" 
+    <style>
+        #{button_id} {{
+            background-color: rgb(255, 255, 255);
+            color: rgb(38, 39, 48);
+            padding: 0.25em 0.38em;
+            position: relative;
+            text-decoration: none;
+            border-radius: 4px;
+            border-width: 1px;
+            border-style: solid;
+            border-color: rgb(230, 234, 241);
+            border-image: initial;
+
+        }} 
+        #{button_id}:hover {{
+            border-color: rgb(246, 51, 102);
+            color: rgb(246, 51, 102);
+        }}
+        #{button_id}:active {{
+            box-shadow: none;
+            background-color: rgb(246, 51, 102);
+            color: white;
+            }}
+    </style> """
+    
+    href =  custom_css + f'<a href="data:file/{filetype};base64,{data_str}" id="{button_id}" download="{filename}">{button_txt}</a>'
     return href
 
+def display_single_image(uploaded_file, img_size=800):
+    print('Type: ', type(uploaded_file))
+    st_title_message.markdown("**Processing your image, please wait** ⌛")
+    img_name = uploaded_file.name
+
+    # Open the image
+    pil_img = PIL.Image.open(uploaded_file)
+    img_rgb = np.array(pil_img)
+    resized_img_rgb = resize_img(img_rgb, img_size)
+    resized_pil_img = PIL.Image.fromarray(resized_img_rgb)
+
+    # Send the image to the model
+    output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
+
+    # Plot images
+    st_input_img.image(resized_pil_img, 'Input image', use_column_width=True)
+    st_output_img.image(output_pil_img, 'Output image', use_column_width=True)
+
+    # Show download button
+    st_download_button.markdown(get_image_download_link(output_pil_img, img_name, 'Download Image'), unsafe_allow_html=True)
+
+    # Reset the message
+    st_title_message.markdown("**To begin, please upload an image** 👇")
+
+def process_multiple_images(uploaded_files, img_size=800):
+    num_imgs = len(uploaded_files)
+
+    output_images_list = []
+    img_names_list = []
+    idx = 1
+    for idx, uploaded_file in enumerate(uploaded_files, start=1):
+        st_title_message.markdown("**Processing image {}/{}. Please wait** ⌛".format(idx,
+                                                                                    num_imgs))
+
+        img_name = uploaded_file.name
+        img_type = uploaded_file.type
+
+        # Open the image
+        pil_img = PIL.Image.open(uploaded_file)
+        img_rgb = np.array(pil_img)
+        resized_img_rgb = resize_img(img_rgb, img_size)
+        resized_pil_img = PIL.Image.fromarray(resized_img_rgb)
+
+        # Send the image to the model
+        output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
+
+        output_images_list.append(output_pil_img)
+        img_names_list.append(img_name.split('.')[0])
+
+    # Zip output files
+    zip_path = 'processed_images.zip'
+    zip_buf = zip_multiple_images(output_images_list, img_names_list, zip_path)
+
+    st_download_button.download_button(
+        label='Download ZIP file',
+        data=zip_buf.read(),
+        file_name=zip_path,
+        mime="application/zip"
+    )
+
+    # Show message
+    st_title_message.markdown("**Images are ready for download** 💾")
+
+def zip_multiple_images(pil_images_list, img_names_list, dest_path):
+    # Create zip file on memory
+    zip_buf = BytesIO()
+
+    with ZipFile(zip_buf, 'w', ZIP_DEFLATED) as zipObj:
+        for pil_img, img_name in zip(pil_images_list, img_names_list):
+            with BytesIO() as output:
+                # Save image in memory
+                pil_img.save(output, format="PNG")
+                
+                # Read data
+                contents = output.getvalue()
+
+                # Write it to zip file
+                zipObj.writestr(img_name+".png", contents)
+    zip_buf.seek(0)
+    return zip_buf
+
+
+
+###########################
+###### STREAMLIT CODE #####
+###########################
 
 # General configuration
 # st.set_page_config(layout="centered")
@@ -118,12 +199,16 @@ unsafe_allow_html=True)
 # Main window configuration
 st.title("Black and white colorizer")
 st.markdown("This app puts color into your black and white pictures")
-title_message = st.empty()
+st_title_message = st.empty()
+st_file_uploader = st.empty()
+st_input_img = st.empty()
+st_output_img = st.empty()
+st_download_button = st.empty()
 
-title_message.markdown("**Model loading, please wait** ⌛")
+st_title_message.markdown("**Model loading, please wait** ⌛")
 
 # # Sidebar
-color_option = st.sidebar.selectbox('Select colorizer mode',
+st_color_option = st.sidebar.selectbox('Select colorizer mode',
                                     ('Artistic', 'Stable'))
                                     
 # st.sidebar.title('Model parameters')
@@ -132,40 +217,27 @@ color_option = st.sidebar.selectbox('Select colorizer mode',
 
 # Load models
 try:
-    colorizer = load_model(model_folder, color_option)
+    print('before loading the model')
+    colorizer = load_model(model_folder, st_color_option)
+    print('after loading the model')
+
 except Exception as e: 
-    print(e)
     colorizer = None
     print('Error while loading the model. Please refresh the page')
+    print(e)
+    st_title_message.markdown("**Error while loading the model. Please refresh the page**")
 
 if colorizer is not None:
-    print('Running colorizer')
-    title_message.markdown("**To begin, please upload an image** 👇")
+    st_title_message.markdown("**To begin, please upload an image** 👇")
 
     #Choose your own image
-    uploaded_file = st.file_uploader("Upload a black and white photo", type=['png', 'jpg', 'jpeg'])
-
-    # show = st.image(use_column_width='auto')
-    input_img_pos = st.empty()
-    output_img_pos = st.empty()
+    uploaded_files = st_file_uploader.file_uploader("Upload a black and white photo", 
+                                                    type=['png', 'jpg', 'jpeg'],
+                                                    accept_multiple_files=True)
 
-    if uploaded_file is not None:
-        img_name = uploaded_file.name
-
-        pil_img = PIL.Image.open(uploaded_file)
-        img_rgb = np.array(pil_img)
-
-        resized_img_rgb = resize_img(img_rgb, max_img_size)
-        resized_pil_img = PIL.Image.fromarray(resized_img_rgb)
-
-        title_message.markdown("**Processing your image, please wait** ⌛")
-
-        output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
-        
-        title_message.markdown("**To begin, please upload an image** 👇")
+    if len(uploaded_files) == 1:
+        display_single_image(uploaded_files[0], max_img_size)
+    elif len(uploaded_files) > 1:
+        process_multiple_images(uploaded_files, max_img_size)
 
-        # Plot images
-        input_img_pos.image(resized_pil_img, 'Input image', use_column_width=True)
-        output_img_pos.image(output_pil_img, 'Output image', use_column_width=True)
 
-        st.markdown(get_image_download_link(output_pil_img, img_name, 'Download '+img_name), unsafe_allow_html=True)

app_utils.py

@@ -106,6 +106,7 @@ def clean_all(files):
         clean_me(me)
 
 
+<<<<<<< HEAD
 def get_model_bin(url, output_path):
     # print('Getting model dir: ', output_path)
     if not os.path.exists(output_path):
@@ -115,14 +116,27 @@ def get_model_bin(url, output_path):
         output_folder = output_path.replace('\\','/').split('/')[0]
         if not os.path.exists(output_folder):
             os.makedirs(output_folder, exist_ok=True)
+=======
+def create_directory(path):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+
+def get_model_bin(url, output_path):
+    # print('Getting model dir: ', output_path)
+    if not os.path.exists(output_path):
+        create_directory(output_path)
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
 
         urllib.request.urlretrieve(url, output_path)
 
         # cmd = "wget -O %s %s" % (output_path, url)
         # print(cmd)
         # os.system(cmd)
+<<<<<<< HEAD
     else:
         print('Model exists')
+=======
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
 
     return output_path
 

deoldify/generators.py

@@ -1,3 +1,4 @@
+<<<<<<< HEAD
 from fastai.basics import F, nn
 from fastai.basic_data import DataBunch
 from fastai.basic_train import Learner
@@ -5,12 +6,17 @@ from fastai.layers import NormType
 from fastai.torch_core import SplitFuncOrIdxList, to_device, apply_init
 from fastai.vision import *
 from fastai.vision.learner import cnn_config, create_body
+=======
+from fastai.vision import *
+from fastai.vision.learner import cnn_config
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
 from .unet import DynamicUnetWide, DynamicUnetDeep
 from .loss import FeatureLoss
 from .dataset import *
 
 # Weights are implicitly read from ./models/ folder
 def gen_inference_wide(
+<<<<<<< HEAD
     root_folder: Path, weights_name: str, nf_factor: int = 2,
                     arch=models.resnet101
                 ) -> Learner:
@@ -41,6 +47,16 @@ def get_inference(learn, root_folder, weights_name) -> Learner:
         print('Error while reading the model')
     learn.model.eval()
     
+=======
+    root_folder: Path, weights_name: str, nf_factor: int = 2, arch=models.resnet101) -> Learner:
+    data = get_dummy_databunch()
+    learn = gen_learner_wide(
+        data=data, gen_loss=F.l1_loss, nf_factor=nf_factor, arch=arch
+    )
+    learn.path = root_folder
+    learn.load(weights_name)
+    learn.model.eval()
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
     return learn
 
 
@@ -104,10 +120,29 @@ def unet_learner_wide(
 
 # ----------------------------------------------------------------------
 
+<<<<<<< HEAD
 def gen_learner_deep(data: ImageDataBunch, gen_loss, arch=models.resnet34, 
                         nf_factor: float = 1.5
                     ) -> Learner:
 
+=======
+# Weights are implicitly read from ./models/ folder
+def gen_inference_deep(
+    root_folder: Path, weights_name: str, arch=models.resnet34, nf_factor: float = 1.5) -> Learner:
+    data = get_dummy_databunch()
+    learn = gen_learner_deep(
+        data=data, gen_loss=F.l1_loss, arch=arch, nf_factor=nf_factor
+    )
+    learn.path = root_folder
+    learn.load(weights_name)
+    learn.model.eval()
+    return learn
+
+
+def gen_learner_deep(
+    data: ImageDataBunch, gen_loss, arch=models.resnet34, nf_factor: float = 1.5
+) -> Learner:
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
     return unet_learner_deep(
         data,
         arch,
@@ -123,6 +158,7 @@ def gen_learner_deep(data: ImageDataBunch, gen_loss, arch=models.resnet34,
 
 # The code below is meant to be merged into fastaiv1 ideally
 def unet_learner_deep(
+<<<<<<< HEAD
                         data: DataBunch,
                         arch: Callable,
                         pretrained: bool = True,
@@ -138,6 +174,22 @@ def unet_learner_deep(
                         **kwargs: Any
                     ) -> Learner:
                     
+=======
+    data: DataBunch,
+    arch: Callable,
+    pretrained: bool = True,
+    blur_final: bool = True,
+    norm_type: Optional[NormType] = NormType,
+    split_on: Optional[SplitFuncOrIdxList] = None,
+    blur: bool = False,
+    self_attention: bool = False,
+    y_range: Optional[Tuple[float, float]] = None,
+    last_cross: bool = True,
+    bottle: bool = False,
+    nf_factor: float = 1.5,
+    **kwargs: Any
+) -> Learner:
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
     "Build Unet learner from `data` and `arch`."
     meta = cnn_config(arch)
     body = create_body(arch, pretrained)

deoldify/visualize.py

@@ -19,13 +19,21 @@ from .generators import gen_inference_deep, gen_inference_wide
 
 
 
+<<<<<<< HEAD
+=======
+# class LoadedModel
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
 class ModelImageVisualizer:
     def __init__(self, filter: IFilter, results_dir: str = None):
         self.filter = filter
         self.results_dir = None if results_dir is None else Path(results_dir)
+<<<<<<< HEAD
 
         if self.results_dir is not None:
             self.results_dir.mkdir(parents=True, exist_ok=True)
+=======
+        self.results_dir.mkdir(parents=True, exist_ok=True)
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
 
     def _clean_mem(self):
         torch.cuda.empty_cache()
@@ -215,15 +223,22 @@ class ModelImageVisualizer:
         return rows, columns
 
 
+<<<<<<< HEAD
 def get_image_colorizer(root_folder: Path = Path('./'), render_factor: int = 35, artistic: bool = True
                         ) -> ModelImageVisualizer:
 
+=======
+def get_image_colorizer(
+    root_folder: Path = Path('./'), render_factor: int = 35, artistic: bool = True
+) -> ModelImageVisualizer:
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
     if artistic:
         return get_artistic_image_colorizer(root_folder=root_folder, render_factor=render_factor)
     else:
         return get_stable_image_colorizer(root_folder=root_folder, render_factor=render_factor)
 
 
+<<<<<<< HEAD
 def get_stable_image_colorizer(root_folder: Path = Path('./'), weights_name: str = 'ColorizeStable_gen',
                                 results_dir='result_images', render_factor: int = 35
                             ) -> ModelImageVisualizer:
@@ -243,4 +258,27 @@ def get_artistic_image_colorizer(root_folder: Path = Path('./'), weights_name: s
     filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
     # vis = ModelImageVisualizer(filtr, results_dir=results_dir)
     vis = ModelImageVisualizer(filtr)
+=======
+def get_stable_image_colorizer(
+    root_folder: Path = Path('./'),
+    weights_name: str = 'ColorizeStable_gen',
+    results_dir='result_images',
+    render_factor: int = 35
+) -> ModelImageVisualizer:
+    learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name)
+    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
+    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
+    return vis
+
+
+def get_artistic_image_colorizer(
+    root_folder: Path = Path('./'),
+    weights_name: str = 'ColorizeArtistic_gen',
+    results_dir='result_images',
+    render_factor: int = 35
+) -> ModelImageVisualizer:
+    learn = gen_inference_deep(root_folder=root_folder, weights_name=weights_name)
+    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
+    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
+>>>>>>> 878ecf212e9f3f2f6e923e3bfff6ec899dc40143
     return vis