Update app.py

app.py

@@ -4,93 +4,72 @@ 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():
-    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()
-
-        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)}")
+st.set_option('deprecation.showPyplotGlobalUse', False)
 
+# Function to display images side by side
 def display(img1, img2):
-    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)}")
+    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)
 
+# Function to read an image
 def imread(img_path):
-    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)}")
+    img = cv2.imread(img_path)
+    if img is None:
+        st.error(f"Failed to load image: {img_path}")
         return None
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    return img
 
+# Function to run shell commands
 def run_shell_commands():
+    st.info("Creating directories...")
+    directories = [
+        "results/cropped_faces",
+        "results/restored_faces",
+        "results/restored_imgs",
+        "results/cmp"
+    ]
+    
+    for directory in directories:
+        os.makedirs(directory, exist_ok=True)
+    
+    st.info("Running image enhancement...")
+    command = "python inference_gfpgan.py -i inputs/upload -o results -v 1.3 -s 2 --bg_upsampler realesrgan"
     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
+            st.error(f"Command failed with error: {process.stderr}")
+            return False
+        return True
     except Exception as e:
-        logger.error(f"Error in shell commands: {str(e)}")
-        return None
+        st.error(f"Failed to run enhancement: {str(e)}")
+        return False
 
+# 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 🙂')
 
+# File uploader with progress bar
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
     try:
-        # Create input directory if it doesn't exist
+        # Show processing status
+        status = st.empty()
+        status.info("Processing started...")
+        
+        # Create input directory
         input_path = os.path.join('inputs', 'upload')
         os.makedirs(input_path, exist_ok=True)
         
@@ -99,13 +78,13 @@ if uploaded_file is not None:
         with open(file_path, 'wb') as f:
             f.write(uploaded_file.getbuffer())
         
-        st.info("Processing image... Please wait.")
+        status.info("File uploaded successfully. Processing image...")
         
-        output = run_shell_commands()
-        if output is not None:
-            st.success("Processing complete!")
+        # Run enhancement
+        if run_shell_commands():
+            status.success("Processing complete!")
             
-            # Display images
+            # Display results
             input_folder = 'results/cropped_faces'
             result_folder = 'results/restored_faces'
             
@@ -113,7 +92,7 @@ if uploaded_file is not None:
             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.")
+                st.warning("No faces detected in the image.")
             else:
                 for input_path, output_path in zip(input_list, output_list):
                     img_input = imread(input_path)
@@ -122,14 +101,15 @@ if uploaded_file is not None:
                         display(img_input, img_output)
             
             # Cleanup
-            try:
+            if os.path.exists(file_path):
                 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)}")
-                
+            
+        else:
+            status.error("Failed to process image.")
+            
     except Exception as e:
-        st.error(f"An error occurred: {str(e)}")
-        logger.error(f"Error processing image: {str(e)}")
+        st.error(f"Error: {str(e)}")
+        
+    finally:
+        # Clear status
+        status.empty()