added multiple files processing

app.py

@@ -1,4 +1,4 @@
-#importing the libraries
+# Import general purpose libraries
 import os, sys, re
 import streamlit as st
 import PIL
@@ -6,19 +6,8 @@ from PIL import Image
 import cv2
 import numpy as np
 import uuid
-
-# 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 create_directory
-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!
@@ -28,6 +17,10 @@ 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/'
@@ -66,42 +59,134 @@ 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_pil_img = resized_pil_img
+
+        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 = io.BytesIO()
+
+    with ZipFile(zip_buf, 'w', ZIP_DEFLATED) as zipObj:
+        for pil_img, img_name in zip(pil_images_list, img_names_list):
+            with io.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")
@@ -136,32 +221,21 @@ st_color_option = st.sidebar.selectbox('Select colorizer mode',
 try:
     colorizer = load_model(model_folder, st_color_option)
 except:
-    colorizer = None
-    print('Error while loading the model. Please refresh the page')
+    pass
+    colorizer = True
+    # print('Error while loading the model. Please refresh the page')
 
 if colorizer is not None:
     st_title_message.markdown("**To begin, please upload an image** 👇")
 
     #Choose your own image
-    uploaded_file = st_file_uploader.file_uploader("Upload a black and white photo", type=['png', 'jpg', 'jpeg'])
-
-    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)
-
-        st_title_message.markdown("**Processing your image, please wait** ⌛")
+    uploaded_files = st_file_uploader.file_uploader("Upload a black and white photo", 
+                                                    type=['png', 'jpg', 'jpeg'],
+                                                    accept_multiple_files=True)
 
-        output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
-        
-        st_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
-        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)
 
-        st_download_button.markdown(get_image_download_link(output_pil_img, img_name, 'Download Image'), unsafe_allow_html=True)