Update app.py

app.py

@@ -1,86 +1,106 @@
-import streamlit as st
-import os
-import subprocess
-import cv2
-import matplotlib.pyplot as plt
-import glob
-
-# Function to display images side by side
-def display(img1, img2):
-    fig = plt.figure(figsize=(25, 10))
-    ax1 = fig.add_subplot(1, 2, 1) 
-    plt.title('Input image', fontsize=16)
-    ax1.axis('off')
-    ax2 = fig.add_subplot(1, 2, 2)
-    plt.title('Image-Blitz output', fontsize=16)
-    ax2.axis('off')
-    ax1.imshow(img1)
-    ax2.imshow(img2)
-    st.pyplot(fig)
-
-# Function to read an image
-def imread(img_path):
-    img = cv2.imread(img_path)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    return img
-
-# Function to run shell commands
-def run_shell_commands():
-    directories = [
-        "results/cropped_faces",
-        "results/restored_faces",
-        "results/restored_imgs",
-        "results/cmp"
-    ]
-    
-    # Create directories using os.makedirs()
-    for directory in directories:
-        os.makedirs(directory, exist_ok=True)
-    
-    # Run inference command
-    command = "python inference_gfpgan.py -i inputs/upload -o results -v 1.3 -s 2 --bg_upsampler realesrgan"
-    subprocess.run(command, shell=True, check=True)
-
-    # List the contents of the results/cmp directory
-    cmp_contents = subprocess.run("dir results\\cmp", shell=True, capture_output=True, text=True)
-    if cmp_contents.returncode != 0:
-        st.error(f"Error listing contents of results/cmp: {cmp_contents.stderr}")
-        return ""
-    return cmp_contents.stdout
-
-# Start Streamlit app
-st.title('Image Enhancer')
-st.write('This is a simple web app to enhance the quality of images')
-
-# Ask the user to input a file
-uploaded_file = st.file_uploader("Choose an image...", type="jpg")
-
-if uploaded_file is not None:
-    # Save the uploaded file to the inputs/upload directory
-    input_path = os.path.join('inputs', 'upload')
-    if not os.path.exists(input_path):
-        os.makedirs(input_path)
-    
-    file_path = os.path.join(input_path, uploaded_file.name)
-    with open(file_path, 'wb') as f:
-        f.write(uploaded_file.getbuffer())
-    
-    # Run shell commands
-    try:
-        output = run_shell_commands()
-        # Display the results
-        st.write('Contents of results/cmp:')
-        st.text(output)
-
-        # Display images
-        input_folder = 'results/cropped_faces'
-        result_folder = 'results/restored_faces'
-        input_list = sorted(glob.glob(os.path.join(input_folder, '*')))
-        output_list = sorted(glob.glob(os.path.join(result_folder, '*')))
-        
-        for input_path, output_path in zip(input_list, output_list):
-            img_input = imread(input_path)
-            img_output = imread(output_path)
-            display(img_input, img_output)
-    except Exception as e:
-        st.error(f"An error occurred: {str(e)}")
+import streamlit as st
+import os
+import subprocess
+import cv2
+import matplotlib.pyplot as plt
+import glob
+
+
+def modify_degradations_py():
+    file_path = '/usr/local/lib/python3.10/site-packages/basicsr/data/degradations.py'
+    with open(file_path, 'r') as f:
+        lines = f.readlines()
+
+    # Find the line containing 'from torchvision.transforms.functional_tensor import rgb_to_grayscale'
+    for i, line in enumerate(lines):
+        if 'from torchvision.transforms.functional_tensor import rgb_to_grayscale' in line:
+            # Replace it with 'from torchvision.transforms.functional import rgb_to_grayscale'
+            lines[i] = 'from torchvision.transforms.functional import rgb_to_grayscale\n'
+            break
+
+    with open(file_path, 'w') as f:
+        f.writelines(lines)
+
+# Call the function to modify the file
+modify_degradations_py()
+
+
+# Function to display images side by side
+def display(img1, img2):
+    fig = plt.figure(figsize=(25, 10))
+    ax1 = fig.add_subplot(1, 2, 1) 
+    plt.title('Input image', fontsize=16)
+    ax1.axis('off')
+    ax2 = fig.add_subplot(1, 2, 2)
+    plt.title('Image-Blitz output', fontsize=16)
+    ax2.axis('off')
+    ax1.imshow(img1)
+    ax2.imshow(img2)
+    st.pyplot(fig)
+
+# Function to read an image
+def imread(img_path):
+    img = cv2.imread(img_path)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    return img
+
+# Function to run shell commands
+def run_shell_commands():
+    directories = [
+        "results/cropped_faces",
+        "results/restored_faces",
+        "results/restored_imgs",
+        "results/cmp"
+    ]
+    
+    # Create directories using os.makedirs()
+    for directory in directories:
+        os.makedirs(directory, exist_ok=True)
+    
+    # Run inference command
+    command = "python inference_gfpgan.py -i inputs/upload -o results -v 1.3 -s 2 --bg_upsampler realesrgan"
+    subprocess.run(command, shell=True, check=True)
+
+    # List the contents of the results/cmp directory
+    cmp_contents = subprocess.run("dir results\\cmp", shell=True, capture_output=True, text=True)
+    if cmp_contents.returncode != 0:
+        st.error(f"Error listing contents of results/cmp: {cmp_contents.stderr}")
+        return ""
+    return cmp_contents.stdout
+
+# Start Streamlit app
+st.title('Image Enhancer')
+st.write('This is a simple web app to enhance the quality of images')
+
+# Ask the user to input a file
+uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+
+if uploaded_file is not None:
+    # Save the uploaded file to the inputs/upload directory
+    input_path = os.path.join('inputs', 'upload')
+    if not os.path.exists(input_path):
+        os.makedirs(input_path)
+    
+    file_path = os.path.join(input_path, uploaded_file.name)
+    with open(file_path, 'wb') as f:
+        f.write(uploaded_file.getbuffer())
+    
+    # Run shell commands
+    try:
+        output = run_shell_commands()
+        # Display the results
+        st.write('Contents of results/cmp:')
+        st.text(output)
+
+        # Display images
+        input_folder = 'results/cropped_faces'
+        result_folder = 'results/restored_faces'
+        input_list = sorted(glob.glob(os.path.join(input_folder, '*')))
+        output_list = sorted(glob.glob(os.path.join(result_folder, '*')))
+        
+        for input_path, output_path in zip(input_list, output_list):
+            img_input = imread(input_path)
+            img_output = imread(output_path)
+            display(img_input, img_output)
+    except Exception as e:
+        st.error(f"An error occurred: {str(e)}")