Update app.py

app.py

@@ -7,11 +7,38 @@ from datasets import Dataset
 from PIL import Image
 import io
 import uuid
+import time
+import random
 
 DATA_DIR = "/data"
 IMAGES_DIR = os.path.join(DATA_DIR, "images")
 DATASET_FILE = os.path.join(DATA_DIR, "dataset.json")
 
+# Add a user agent rotation list
+USER_AGENTS = [
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
+    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.0.3 Safari/605.1.15",
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0"
+]
+
+def get_headers(cookies=None):
+    headers = {
+        "User-Agent": random.choice(USER_AGENTS),
+        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
+        "Accept-Language": "en-US,en;q=0.5",
+        "Referer": "https://www.google.com/",
+        "DNT": "1",
+        "Connection": "keep-alive",
+        "Upgrade-Insecure-Requests": "1"
+    }
+    if cookies:
+        headers["Cookie"] = cookies
+    return headers
+
+def make_request(url, cookies=None):
+    time.sleep(random.uniform(1, 3))  # Add a random delay between requests
+    return requests.get(url, headers=get_headers(cookies), timeout=10)
+
 def extract_image_url(html_content):
     soup = BeautifulSoup(html_content, 'html.parser')
     script = soup.find('script', type='text/javascript', string=lambda text: 'image =' in text if text else False)
@@ -38,9 +65,9 @@ def extract_tags(html_content):
     
     return ','.join(tags)
 
-def download_image(url):
+def download_image(url, cookies=None):
     try:
-        response = requests.get(url)
+        response = make_request(url, cookies)
         response.raise_for_status()
         return Image.open(io.BytesIO(response.content))
     except requests.RequestException as e:
@@ -61,9 +88,9 @@ class DatasetBuilder:
         with open(DATASET_FILE, 'w') as f:
             json.dump(self.dataset, f)
 
-    def add_image(self, url):
+    def add_image(self, url, cookies=None):
         try:
-            response = requests.get(url)
+            response = make_request(url, cookies)
             response.raise_for_status()
             html_content = response.text
 
@@ -72,7 +99,7 @@ class DatasetBuilder:
                 raise Exception("Failed to extract image URL")
 
             tags = extract_tags(html_content)
-            image = download_image(image_url)
+            image = download_image(image_url, cookies)
 
             # Generate a unique filename
             filename = f"{uuid.uuid4()}.jpg"
@@ -116,8 +143,8 @@ class DatasetBuilder:
 
 dataset_builder = DatasetBuilder()
 
-def add_image_to_dataset(url):
-    result = dataset_builder.add_image(url)
+def add_image_to_dataset(url, cookies):
+    result = dataset_builder.add_image(url, cookies)
     return result, dataset_builder.get_dataset_info(), dataset_builder.get_dataset_preview()
 
 def create_huggingface_dataset():
@@ -133,6 +160,7 @@ with gr.Blocks(theme="huggingface") as iface:
 
     with gr.Row():
         url_input = gr.Textbox(lines=2, placeholder="Enter image URL here...")
+        cookies_input = gr.Textbox(lines=2, placeholder="Enter cookies (optional)")
         add_button = gr.Button("Add Image")
 
     result_output = gr.Textbox(label="Result")
@@ -141,7 +169,7 @@ with gr.Blocks(theme="huggingface") as iface:
     gr.Markdown("## Dataset Preview")
     preview_gallery = gr.Gallery(label="Recent Additions", show_label=False, elem_id="preview_gallery", columns=5, rows=1, height="auto")
 
-    add_button.click(add_image_to_dataset, inputs=url_input, outputs=[result_output, dataset_info, preview_gallery])
+    add_button.click(add_image_to_dataset, inputs=[url_input, cookies_input], outputs=[result_output, dataset_info, preview_gallery])
 
     create_hf_button = gr.Button("Create HuggingFace Dataset")
     hf_result = gr.Textbox(label="Dataset Creation Result")