import requests from bs4 import BeautifulSoup import os import json import gradio as gr 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") 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) if script: try: js_object_str = script.string.split('=', 1)[1].strip().rstrip(';') js_object_str = js_object_str.replace("'", '"') image_data = json.loads(js_object_str) return f"{image_data['domain']}{image_data['base_dir']}/{image_data['dir']}/{image_data['img']}" except json.JSONDecodeError as e: raise Exception(f"Failed to decode JSON: {str(e)}") img_tag = soup.find('img', alt=True) if img_tag and 'src' in img_tag.attrs: return img_tag['src'] return None def extract_tags(html_content): soup = BeautifulSoup(html_content, 'html.parser') tag_elements = soup.find_all('li', class_='tag-type-general') tags = [tag_element.find_all('a')[1].text for tag_element in tag_elements if len(tag_element.find_all('a')) > 1] return ','.join(tags) def download_image(url, cookies=None): try: response = make_request(url, cookies) response.raise_for_status() return Image.open(io.BytesIO(response.content)) except requests.RequestException as e: raise Exception(f"Failed to download image: {str(e)}") class DatasetBuilder: def __init__(self, dataset_name): self.dataset_name = dataset_name self.dataset = self.load_dataset() os.makedirs(IMAGES_DIR, exist_ok=True) def get_dataset_file(self): return os.path.join(DATA_DIR, f"{self.dataset_name}.json") def load_dataset(self): dataset_file = self.get_dataset_file() if os.path.exists(dataset_file): with open(dataset_file, 'r') as f: return json.load(f) return [] def save_dataset(self): dataset_file = self.get_dataset_file() with open(dataset_file, 'w') as f: json.dump(self.dataset, f) def add_image(self, url, cookies=None): try: response = make_request(url, cookies) response.raise_for_status() html_content = response.text image_url = extract_image_url(html_content) if not image_url: raise Exception("Failed to extract image URL") tags = extract_tags(html_content) image = download_image(image_url, cookies) filename = f"{uuid.uuid4()}.jpg" filepath = os.path.join(IMAGES_DIR, filename) image.save(filepath) self.dataset.append({ 'image': filename, 'tags': tags }) self.save_dataset() return f"Added image with tags: {tags}" except Exception as e: return f"Error: {str(e)}" def build_huggingface_dataset(self): if not self.dataset: return "Dataset is empty. Add some images first." try: hf_dataset = Dataset.from_dict({ 'image': [os.path.join(IMAGES_DIR, item['image']) for item in self.dataset], 'tags': [item['tags'] for item in self.dataset] }) return "HuggingFace Dataset created successfully!" except Exception as e: return f"Error creating HuggingFace Dataset: {str(e)}" def get_dataset_info(self): return f"Current dataset size ({self.dataset_name}): {len(self.dataset)} images" def get_dataset_preview(self, num_images=5): preview = [] for item in self.dataset[-num_images:]: image_path = os.path.join(IMAGES_DIR, item['image']) preview.append((image_path, item['tags'])) return preview def add_image_to_dataset(url, cookies, dataset_name): builder = DatasetBuilder(dataset_name) result = builder.add_image(url, cookies) return result, builder.get_dataset_info(), builder.get_dataset_preview() def create_huggingface_dataset(dataset_name): builder = DatasetBuilder(dataset_name) return builder.build_huggingface_dataset() def view_dataset(dataset_name): builder = DatasetBuilder(dataset_name) return builder.get_dataset_preview(num_images=20) # Create Gradio interface with gr.Blocks(theme="huggingface") as iface: gr.Markdown("# Image Dataset Builder") gr.Markdown("Enter a URL to add an image and its tags to the dataset. Progress is saved automatically.") with gr.Row(): dataset_name_input = gr.Textbox(lines=1, label="Dataset Name", placeholder="Enter dataset name...", value="default_dataset") url_input = gr.Textbox(lines=2, label="URL", placeholder="Enter image URL here...") cookies_input = gr.Textbox(lines=2, label="Cookies (optional)", placeholder="Enter cookies") add_button = gr.Button("Add Image") result_output = gr.Textbox(label="Result") dataset_info = gr.Textbox(label="Dataset Info") 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, cookies_input, dataset_name_input], outputs=[result_output, dataset_info, preview_gallery]) create_hf_button = gr.Button("Create HuggingFace Dataset") hf_result = gr.Textbox(label="Dataset Creation Result") create_hf_button.click(create_huggingface_dataset, inputs=[dataset_name_input], outputs=hf_result) view_dataset_button = gr.Button("View Dataset") dataset_gallery = gr.Gallery(label="Dataset Contents", show_label=False, elem_id="dataset_gallery", columns=5, rows=4, height="auto") view_dataset_button.click(view_dataset, inputs=[dataset_name_input], outputs=dataset_gallery) # Launch the interface iface.launch()