Update app.py

app.py

@@ -1,12 +1,10 @@
 import math
 import time
-import numpy as np
 import streamlit as st
 from PIL import Image
 import cv2
-
-
-st.set_page_config(layout="wide")
+import os
+import urllib.request
 
 hide_streamlit_style = """
 <style>
@@ -17,16 +15,36 @@ hide_streamlit_style = """
 
 """
 
-st.markdown(hide_streamlit_style, unsafe_allow_html=True)
+st.set_page_config(layout="wide")
+
+if os.path.isfile('models/ESPCN_x2.pb'):
+    print(True)
+else:
+    print("Downloading model x2 ...")
+    urllib.request.urlretrieve("https://github.com/fannymonori/TF-ESPCN/raw/master/export/ESPCN_x2.pb",
+                               "models/ESPCN_x2.pb")
+    print("Done!")
+
+    print("Downloading model x3 ...")
+    url2 = "https://github.com/fannymonori/TF-ESPCN/raw/master/export/ESPCN_x3.pb"
+    urllib.request.urlretrieve(url2, "models/ESPCN_x3.pb")
+    print("Done!")
 
+    print("Downloading model x4 ...")
+    url3 = "https://github.com/fannymonori/TF-ESPCN/raw/master/export/ESPCN_x4.pb"
+    urllib.request.urlretrieve(url3, "models/ESPCN_x4.pb")
+    print("Done!")
+
+
+@st.cache_resource
 def loadModel(n):
     super_res = cv2.dnn_superres.DnnSuperResImpl_create()
     super_res.readModel('models/ESPCN_x'+n+'.pb')
     return super_res
 
 # on removing (show_spinner=False), it will show that fuction is running on web app
-@st.experimental_memo(show_spinner=False)
-def upscale(file,task):
+@st.cache_data(show_spinner=False)
+def upscale(file,task,_progressBar = None):
     with open(file.name, "wb") as f:
         f.write(file.getbuffer())
         print('No file found, so added in list files')
@@ -38,7 +56,37 @@ def upscale(file,task):
             upscaled_image = super_res.upsample(img)
             print('I upscaled upto',task,'times')
             cv2.imwrite("processed_"+file.name,upscaled_image)
+            with st.sidebar:
+                st.success('Done!', icon="✅")
+            return True
+        elif file.type.split('/')[0] == 'video':
+            cap = cv2.VideoCapture(file.name)
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            fps = int(cap.get(cv2.CAP_PROP_FPS))
+            width = int(cap.get(3))
+            height = int(cap.get(4))
+            writer = cv2.VideoWriter("processed_"+file.name,fourcc,fps,(width*int(task),height*int(task)))
+            length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            step_size = 1.0/length
+            progress = 0
+
+            with st.sidebar:
+                st.info("Operation in progress. Please wait.",icon="ℹ️")
+                my_bar = st.progress(0,text="")
+
+            for percent_complete in range(length):
+                frame = cap.read()[1]
+                frame = super_res.upsample(frame)
+                writer.write(frame)
+                print(progress)
+                progress += step_size
+                my_bar.progress(progress-0.000000001)
+            with st.sidebar:
+                st.success('Done! Thankyou for your patience', icon="✅")
+            return True
         return True
+
+    # Second case where custom size is required
     else:
         req_width,req_height = int(task[0]),int(task[1])
         if file.type.split('/')[0] == 'image':
@@ -49,6 +97,8 @@ def upscale(file,task):
                 img = cv2.resize(img,(req_width,req_height))
                 print("I did resizing only!")
                 cv2.imwrite("processed_" + file.name, img)
+                with st.sidebar:
+                    st.success('Done!', icon="✅")
                 return True
             # rounding off the ratios
             w_ratio,h_ratio = math.ceil(w_ratio),math.ceil(h_ratio)
@@ -69,9 +119,37 @@ def upscale(file,task):
             print("I upscale upto", upscale_number , "times and then resize it.")
 
             cv2.imwrite("processed_" + file.name, upscaled_image)
-
+            with st.sidebar:
+                st.success('Done!', icon="✅")
             return True
 
+        # If file is video
+        elif file.type.split('/')[0] == 'video':
+            cap = cv2.VideoCapture(file.name)
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            fps = int(cap.get(cv2.CAP_PROP_FPS))
+            width = int(cap.get(3))
+            height = int(cap.get(4))
+            if height > 2160 or width > 3840:
+                with st.sidebar:
+                    st.success("Sorry, I can't processed Video with resolution above 4k. Please select custom size option.!", icon="ℹ️")
+            writer = cv2.VideoWriter("processed_" + file.name, fourcc, fps, (int(task[0]),int(task[1])))
+            length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            step_size = 1.0 / length
+            progress = 0
+            progress_text = "Operation in progress. Please wait."
+            with st.sidebar:
+                st.info("Operation in progress. Please wait.", icon="ℹ️")
+                my_bar = st.progress(0)
+            for percent_complete in range(length):
+                frame = cap.read()[1]
+                frame = cv2.resize(frame, (task[0], task[1]))
+                writer.write(frame)
+                progress += step_size
+                my_bar.progress(progress-0.000000001)
+            with st.sidebar:
+                st.success('Done! Thankyou for your patience', icon="✅")
+            return True
         return "It's second"
 
 
@@ -81,8 +159,12 @@ if 'disable_opt1' not in st.session_state:
     st.session_state.disable_opt1 = False
 if 'disable_download' not in st.session_state:
     st.session_state.disable_download = True
+if 'disable_proceed' not in st.session_state:
+    st.session_state.disable_proceed = False
 
 
+st.markdown(hide_streamlit_style, unsafe_allow_html=True)
+
 col1,_,col2 = st.columns([6,1,3],gap="small")
 
 def toggle_state_opt1():
@@ -106,19 +188,39 @@ def toggle_state_opt2():
 # Update the states based on user selection before drawing the widgets in the web page
 toggle_state_opt2()
 toggle_state_opt1()
+options = ["2", "3","4"]
+progressBar = None
 
 with col1:
-    file = st.file_uploader("",type=['png','jpeg','jpg','pgm','jpe'])
+    file = st.file_uploader(" ",type=['png','jpeg','jpg','pgm','jpe','mp4','mov'])
     if file is not None:
         # writing file and saving its details in dict for further processing
-
-
-
-        if file.type.split('/')[0] == "image":
+        bytes_data = file.getvalue()
+        file_size = len(bytes_data)
+        print("File size: ",file_size)
+
+        if file.type.split('/')[0] == "image" and file_size > 1550000:
+            st.session_state.disable_proceed = True
+            with st.sidebar:
+                st.info('Sorry, maximum size of image is 1.5MB', icon="ℹ️")
+        elif file.type.split('/')[0] == "image":
             image = Image.open(file)
+            st.session_state.disable_proceed = False
             st.image(image,caption="Upload Image", use_column_width=True)
+            st.session_state.disable_proceed = False
+        elif file.type.split('/')[0] == 'video' and file_size > 100000000:
+            with st.sidebar:
+                options = ["2", "3"]
+                st.info('Sorry, maximum size of video is 100MB', icon="ℹ️")
+                st.session_state.disable_proceed = True
         elif file.type.split('/')[0] == 'video':
-            st.video(file)
+            video = st.video(file)
+            print(type(video))
+            options = ["2", "3"]
+            st.session_state.disable_proceed = False
+            with st.sidebar:
+                st.info('For custom size, currently I can processed video without AI.', icon="ℹ️")
+
 
 
 with col2:
@@ -131,35 +233,30 @@ with col2:
     st.markdown("\n")
 
     opt1 = st.checkbox("MULTIPLES OF",key="opt1",value=True,on_change=toggle_state_opt1)
-    st.selectbox("SELECT", ["2", "3", "4"],key="opt1_selBox",disabled=st.session_state.disable_opt1)
+    st.selectbox("SELECT", options,key="opt1_selBox",disabled=st.session_state.disable_opt1)
 
     st.markdown("\n")
     st.markdown("\n")
     opt2 = st.checkbox("CUSTOM SIZE",key="opt2",on_change=toggle_state_opt2)
 
-    
-
-    
     st.number_input("Width", step=1, min_value=150,max_value=3840, value=900, key="width",disabled=st.session_state.disable_opt2)
 
-    
     st.number_input("Height", step=1, min_value=150,max_value=2160, value=900, key="height",disabled=st.session_state.disable_opt2)
 
     st.markdown("\n")
     st.markdown("\n")
 
-    #_, dcol, _ = st.columns([1,5,1],gap="small")
+    _, dcol, _ = st.columns([1,5,1],gap="small")
 
-    
 
-    if st.button(23*"&nbsp;"+"PROCEED"+"&nbsp;"*23) and file is not None:
+
+    if st.button(26*"&nbsp;"+"PROCEED"+"&nbsp;"*26,disabled=st.session_state.disable_proceed) and file is not None:
         if st.session_state.get('opt1') == True:
             task = st.session_state.opt1_selBox
         else:
             task = [st.session_state.width, st.session_state.height]
         print(task)
-        st.session_state.disable_download = not upscale(file,task)
-
+        st.session_state.disable_download = not upscale(file,task,progressBar)
 
         #print(resulted_file.shape)
 
@@ -170,15 +267,13 @@ with col2:
         st.session_state.disable_download = True
 
     if st.session_state.disable_download == True:
-        st.button(16*"&nbsp;"+"DOWNLOAD FILE"+"&nbsp;"*16,disabled=True)
+        st.button(19*"&nbsp;"+"DOWNLOAD FILE"+"&nbsp;"*19,disabled=True)
     else:
         with open('processed_'+file.name, "rb") as download_file:
-            st.download_button(label=16*"&nbsp;"+"DOWNLOAD FILE"+"&nbsp;"*16, data=download_file,
+            st.download_button(label=19*"&nbsp;"+"DOWNLOAD FILE"+"&nbsp;"*19, data=download_file,
                                 file_name= 'processed_'+file.name, mime= "image/png",
-                                disabled=st.session_state.disable_download)
+                                 disabled=st.session_state.disable_download)
 
-
-st.markdown("\n")
 st.markdown("\n")
 st.markdown("\n")
 st.info("DESCRIPTION :    This web app is a free tool designed to upscale or resize image resolution. While the app"+
@@ -186,3 +281,4 @@ st.info("DESCRIPTION :    This web app is a free tool designed to upscale or res
         " We welcome your feedback and suggestions, and encourage you to contact us at [email protected] "+
         "to share your thoughts. Thank you for your interest in our web application, and we look forward "+
         "to hearing from you as we continue to work towards making this project a resounding success.")
+