Update app.py

app.py

@@ -4,110 +4,132 @@ import subprocess
 import cv2
 import matplotlib.pyplot as plt
 import glob
+import logging
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 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)
+    try:
+        file_path = '/usr/local/lib/python3.10/site-packages/basicsr/data/degradations.py'
+        if not os.path.exists(file_path):
+            logger.warning(f"File not found: {file_path}")
+            return
+            
+        with open(file_path, 'r') as f:
+            lines = f.readlines()
 
-# Call the function to modify the file
-modify_degradations_py()
+        for i, line in enumerate(lines):
+            if 'from torchvision.transforms.functional_tensor import rgb_to_grayscale' in line:
+                lines[i] = 'from torchvision.transforms.functional import rgb_to_grayscale\n'
+                break
 
+        with open(file_path, 'w') as f:
+            f.writelines(lines)
+    except Exception as e:
+        logger.error(f"Error modifying degradations.py: {str(e)}")
 
-# 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)
+    try:
+        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)
+        plt.close(fig)  # Clean up
+    except Exception as e:
+        st.error(f"Error displaying images: {str(e)}")
 
-# Function to read an image
 def imread(img_path):
-    img = cv2.imread(img_path)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    return img
+    try:
+        if not os.path.exists(img_path):
+            raise FileNotFoundError(f"Image not found: {img_path}")
+        img = cv2.imread(img_path)
+        if img is None:
+            raise ValueError(f"Failed to load image: {img_path}")
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        return img
+    except Exception as e:
+        logger.error(f"Error reading image: {str(e)}")
+        return None
 
-# 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
+    try:
+        directories = [
+            "results/cropped_faces",
+            "results/restored_faces",
+            "results/restored_imgs",
+            "results/cmp"
+        ]
+        
+        for directory in directories:
+            os.makedirs(directory, exist_ok=True)
+        
+        command = "python inference_gfpgan.py -i inputs/upload -o results -v 1.3 -s 2 --bg_upsampler realesrgan"
+        process = subprocess.run(command, shell=True, capture_output=True, text=True)
+        
+        if process.returncode != 0:
+            raise subprocess.CalledProcessError(process.returncode, command, process.stdout, process.stderr)
+            
+        return process.stdout
+    except Exception as e:
+        logger.error(f"Error in shell commands: {str(e)}")
+        return None
 
-# Start Streamlit app
 st.title('Image Enhancer')
 st.write('This is a simple web app to enhance the quality of images')
 st.write('Kindly wait for 30-40 seconds after uploading a photo 🙂')
 
-# Ask the user to input a file
-uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
 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, '*')))
+        # Create input directory if it doesn't exist
+        input_path = os.path.join('inputs', 'upload')
+        os.makedirs(input_path, exist_ok=True)
+        
+        # Save uploaded file
+        file_path = os.path.join(input_path, uploaded_file.name)
+        with open(file_path, 'wb') as f:
+            f.write(uploaded_file.getbuffer())
         
-        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)
-            os.remove(file_path)
-            os.remove(input_path)
-            os.remove(output_path)
+        st.info("Processing image... Please wait.")
+        
+        output = run_shell_commands()
+        if output is not None:
+            st.success("Processing complete!")
+            
+            # 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, '*')))
+            
+            if not input_list or not output_list:
+                st.warning("No results found. The image might not contain any faces.")
+            else:
+                for input_path, output_path in zip(input_list, output_list):
+                    img_input = imread(input_path)
+                    img_output = imread(output_path)
+                    if img_input is not None and img_output is not None:
+                        display(img_input, img_output)
             
-    except subprocess.CalledProcessError as e:
-        st.error(f"Command failed with error: {e.stderr}")
+            # Cleanup
+            try:
+                os.remove(file_path)
+                for path in input_list + output_list:
+                    if os.path.exists(path):
+                        os.remove(path)
+            except Exception as e:
+                logger.error(f"Error during cleanup: {str(e)}")
+                
     except Exception as e:
-        st.error(f"An unexpected error occurred: {str(e)}")
+        st.error(f"An error occurred: {str(e)}")
+        logger.error(f"Error processing image: {str(e)}")