Update app.py

app.py

@@ -1,35 +1,160 @@
-import gradio as gr
+import math
+import streamlit as st
+from PIL import Image
 import cv2
-import numpy as np
-
-# 이미지 업스케일 함수
-def upscale_image(input_image, radio_input):
-    upscale_factor = int(radio_input)  # 업스케일 비율을 정수형으로 변환
-    # 입력 이미지를 주어진 업스케일 비율로 크기 변경 (보간법: INTER_CUBIC)
-    output_image = cv2.resize(input_image, None, fx=upscale_factor, fy=upscale_factor, interpolation=cv2.INTER_CUBIC)
-    # 업스케일된 이미지를 PNG 파일로 저장
-    output_file_path = "/mnt/data/upscaled_image.png"
-    cv2.imwrite(output_file_path, output_image)
-    return output_file_path
-
-# 인터페이스 설명
-DESCRIPTION = """
-이미지의 크기와 품질을 향상시켜 보세요! (you can increase the size and quality of your images)
+import os
+import threading
+
+hide_streamlit_style = """
+<style>
+    #root > div:nth-child(1) > div > div > div > div > section > div {padding-top: 0rem;
+    padding-left: 1%;
+    }
+</style>
 """
 
-# 업스케일 레벨 선택 라디오 버튼
-radio_input = gr.Radio(label="업스케일 레벨 선택 (Select Upscaling level)", choices=[2, 4, 6, 8, 10], value=2)
+st.set_page_config(layout="wide")
+
+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 function is running on web app
+@st.experimental_memo(show_spinner=False)
+def upscale(file, task):
+    with open(file.name, "wb") as f:
+        f.write(file.getbuffer())
+        print('No file found, so added in list files')
+    if isinstance(task, str):
+        super_res = loadModel(task)
+        super_res.setModel('espcn', int(task))
+        if file.type.split('/')[0] == 'image':
+            img = cv2.imread(file.name)
+            upscaled_image = super_res.upsample(img)
+            print('I upscaled up to', task, 'times')
+            output_file_path = "processed_" + file.name.split('.')[0] + ".png"
+            cv2.imwrite(output_file_path, upscaled_image)
+            with st.sidebar:
+                st.success('Done!', icon="✅")
+            return output_file_path
+        elif file.type.split('/')[0] == 'video':
+            # 영상 업스케일링 코드는 생략 (필요시 추가 가능)
+            pass
+        return True
+    else:
+        # 커스텀 사이즈 코드 (필요시 추가)
+        pass
+
+    return True
+
+if 'disable_opt2' not in st.session_state:
+    st.session_state.disable_opt2 = True
+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():
+    if st.session_state.get("opt1") == True:
+        st.session_state.opt2 = False
+        st.session_state.disable_opt2 = True
+    else:
+        st.session_state.opt2 = True
+        st.session_state.disable_opt2 = False
+
+def toggle_state_opt2():
+    if st.session_state.get("opt2") == True:
+        st.session_state.opt1 = False
+        st.session_state.disable_opt1 = True
+    else:
+        st.session_state.opt1 = True
+        st.session_state.disable_opt1 = False
+
+toggle_state_opt2()
+toggle_state_opt1()
+options = ["2", "3", "4"]
+
+with col1:
+    file = st.file_uploader(" ", type=['png', 'jpeg', 'jpg', 'pgm', 'jpe', 'mp4', 'mov'])
+    if file is not None:
+        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 > 200000000:
+            with st.sidebar:
+                options = ["2", "3"]
+                st.info('Sorry, maximum size of video is 200MB', icon="ℹ️")
+                st.session_state.disable_proceed = True
+        elif file.type.split('/')[0] == 'video':
+            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 process video without AI.', icon="ℹ️")
+
+with col2:
+    st.markdown("\n")
+    st.markdown("\n")
+    st.markdown("\n")
+
+    st.subheader(" UPSCALE RESOLUTION UP TO")
+    st.markdown("\n")
+    st.markdown("\n")
+
+    opt1 = st.checkbox("MULTIPLES OF", key="opt1", value=True, on_change=toggle_state_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")
+
+    if st.button(15*"&nbsp;"+"PROCEED"+"&nbsp;"*15, 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)
+        result_file_path = upscale(file, task)
+
+        if result_file_path:
+            st.session_state.disable_download = False
+            st.session_state.result_file_path = result_file_path
 
-# Gradio 인터페이스 정의
-demo = gr.Interface(
-    fn=upscale_image,
-    inputs=[gr.Image(label="입력 이미지 (Input Image)", type="numpy"), radio_input],
-    outputs=gr.File(label="PNG 파일 다운로드 (Download PNG File)"),
-    title="Image Upscaler",
-    description=DESCRIPTION
-)
+    st.markdown("\n")
+    st.markdown("\n")
 
-# 인터페이스 실행
-demo.launch(show_api=False)
+    if file is None:
+        st.session_state.disable_download = True
 
+    if st.session_state.disable_download:
+        st.button(13*"&nbsp;"+"DOWNLOAD"+"&nbsp;"*13, disabled=True)
+    else:
+        with open(st.session_state.result_file_path, "rb") as download_file:
+            st.download_button(label=13*"&nbsp;"+"DOWNLOAD"+"&nbsp;"*13, data=download_file, file_name='processed_' + file.name.split('.')[0] + ".png", mime="image/png")
 
+st.markdown("\n")
+st.markdown("\n")