Update app.py

app.py

@@ -1,16 +1,11 @@
 import os
 import gradio as gr
 from zeroscratches import EraseScratches
-import torch
 import cv2
 import numpy as np
 from concurrent.futures import ThreadPoolExecutor
 from PIL import Image
 
-# ✅ Force GPU usage if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-restorer = EraseScratches().to(device)
-
 # ✅ Custom CSS for clean UI
 custom_css = """
 /* Dark theme styling */
@@ -64,16 +59,15 @@ img {
 }
 """
 
-# ✅ Optimized image processing function
+# ✅ Image processing function
 def predict(img_file):
-    """Run image processing with multi-threading"""
+    """Apply scratch removal with multi-threading"""
     with ThreadPoolExecutor(max_workers=4) as executor:
         future = executor.submit(process_image, img_file)
         restored_img = future.result()
-
     return restored_img
 
-# ✅ Image pre-processing function
+# ✅ Faster image processing with OpenCV
 def process_image(img_file):
     """Efficient image loading and restoration"""
     
@@ -81,18 +75,16 @@ def process_image(img_file):
     img = cv2.imdecode(np.frombuffer(img_file.read(), np.uint8), cv2.IMREAD_COLOR)
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
-    # Resize large images to speed up processing
+    # Resize large images for faster processing
     max_size = (1024, 1024)
     img = cv2.resize(img, max_size, interpolation=cv2.INTER_AREA)
 
-    # Convert image to tensor and send to GPU
-    img_tensor = torch.from_numpy(img).to(device)
-
-    # Perform restoration
-    restored_img = restorer.erase(img_tensor).cpu().numpy()
+    # Convert to PIL format
+    img_pil = Image.fromarray(img)
 
-    # Convert back to PIL format
-    restored_img = Image.fromarray(restored_img)
+    # Perform scratch removal
+    restorer = EraseScratches()
+    restored_img = restorer.erase(img_pil)
 
     return restored_img