Update app.py

app.py

@@ -1,7 +1,8 @@
 import cv2
 import gradio as gr
 import os
-from huggingface_hub import hf_hub_download  # Added for Hugging Face download
+import requests  # Added for GitHub downloads
+from huggingface_hub import hf_hub_download  # For Hugging Face download
 from PIL import Image
 import numpy as np
 import torch
@@ -9,18 +10,24 @@ from torch.autograd import Variable
 from torchvision import transforms
 import torch.nn.functional as F
 
-# Project imports (assumes data_loader_cache.py and models.py are uploaded)
-from data_loader_cache import normalize, im_reader, im_preprocess
-from models import *
+# Automatically download required files
+# 1. data_loader_cache.py from GitHub
+if not os.path.exists("data_loader_cache.py"):
+    print("Downloading data_loader_cache.py...")
+    response = requests.get("https://raw.githubusercontent.com/xuebinqin/DIS/main/IS-Net/data_loader_cache.py")
+    with open("data_loader_cache.py", "wb") as f:
+        f.write(response.content)
 
-# Helpers
-device = 'cpu'  # Free Hugging Face Space uses CPU
+# 2. models.py from GitHub
+if not os.path.exists("models.py"):
+    print("Downloading models.py...")
+    response = requests.get("https://raw.githubusercontent.com/xuebinqin/DIS/main/IS-Net/models.py")
+    with open("models.py", "wb") as f:
+        f.write(response.content)
 
-# Create directory for model weights
+# 3. isnet.pth from ECCV2022/dis-background-removal
 if not os.path.exists("saved_models"):
     os.makedirs("saved_models")
-
-# Automatically download isnet.pth from ECCV2022/dis-background-removal if not present
 isnet_path = "saved_models/isnet.pth"
 if not os.path.exists(isnet_path):
     print("Downloading isnet.pth from ECCV2022/dis-background-removal...")
@@ -31,6 +38,13 @@ if not os.path.exists(isnet_path):
         local_dir_use_symlinks=False
     )
 
+# Project imports
+from data_loader_cache import normalize, im_reader, im_preprocess
+from models import *
+
+# Helpers
+device = 'cpu'  # Free Hugging Face Space uses CPU
+
 class GOSNormalize(object):
     def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
         self.mean = mean
@@ -60,21 +74,21 @@ def build_model(hypar, device):
 def predict(net, inputs_val, shapes_val, hypar, device):
     net.eval()
     inputs_val = inputs_val.type(torch.FloatTensor).to(device)
-    with torch.no_grad():  # Reduce memory usage
+    with torch.no_grad():
         inputs_val_v = Variable(inputs_val)
         ds_val = net(inputs_val_v)[0]
-        pred_val = ds_val[0][0, :, :, :]  # B x 1 x H x W
+        pred_val = ds_val[0][0, :, :, :]
         pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val, 0), (shapes_val[0][0], shapes_val[0][1]), mode='bilinear'))
         ma = torch.max(pred_val)
         mi = torch.min(pred_val)
-        pred_val = (pred_val - mi) / (ma - mi)  # Normalize to [0, 1]
+        pred_val = (pred_val - mi) / (ma - mi)
         return (pred_val.cpu().numpy() * 255).astype(np.uint8)
 
 # Set Parameters
 hypar = {
     "model_path": "saved_models",
     "restore_model": "isnet.pth",
-    "cache_size": [512, 512],  # Optimized for CPU
+    "cache_size": [512, 512],
     "input_size": [512, 512],
     "crop_size": [512, 512],
     "model": ISNetDIS()