Create app.py

app.py

@@ -0,0 +1,156 @@
+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 shutil
+
+DATA_DIR = "/data"
+IMAGES_DIR = os.path.join(DATA_DIR, "images")
+DATASET_FILE = os.path.join(DATA_DIR, "dataset.json")
+
+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:
+        image_data = json.loads(script.string.split('=', 1)[1].strip().rstrip(';'))
+        return f"{image_data['domain']}{image_data['base_dir']}/{image_data['dir']}/{image_data['img']}"
+    
+    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 = []
+    for tag_element in tag_elements:
+        tag_link = tag_element.find_all('a')[1]
+        if tag_link:
+            tags.append(tag_link.text)
+    
+    return ','.join(tags)
+
+def download_image(url):
+    try:
+        response = requests.get(url)
+        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):
+        self.dataset = self.load_dataset()
+        os.makedirs(IMAGES_DIR, exist_ok=True)
+
+    def load_dataset(self):
+        if os.path.exists(DATASET_FILE):
+            with open(DATASET_FILE, 'r') as f:
+                return json.load(f)
+        return []
+
+    def save_dataset(self):
+        with open(DATASET_FILE, 'w') as f:
+            json.dump(self.dataset, f)
+
+    def add_image(self, url):
+        try:
+            response = requests.get(url)
+            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)
+
+            # Generate a unique filename
+            filename = f"{uuid.uuid4()}.jpg"
+            filepath = os.path.join(IMAGES_DIR, filename)
+
+            # Save the image
+            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: {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
+
+dataset_builder = DatasetBuilder()
+
+def add_image_to_dataset(url):
+    result = dataset_builder.add_image(url)
+    return result, dataset_builder.get_dataset_info(), dataset_builder.get_dataset_preview()
+
+def create_huggingface_dataset():
+    return dataset_builder.build_huggingface_dataset()
+
+def view_dataset():
+    return dataset_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():
+        url_input = gr.Textbox(lines=2, placeholder="Enter image URL here...")
+        add_button = gr.Button("Add Image")
+
+    result_output = gr.Textbox(label="Result")
+    dataset_info = gr.Textbox(label="Dataset Info", value=dataset_builder.get_dataset_info())
+
+    gr.Markdown("## Dataset Preview")
+    preview_gallery = gr.Gallery(label="Recent Additions", show_label=False, elem_id="preview_gallery").style(grid=5, height="auto")
+
+    add_button.click(add_image_to_dataset, inputs=url_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=[], outputs=hf_result)
+
+    view_dataset_button = gr.Button("View Dataset")
+    dataset_gallery = gr.Gallery(label="Dataset Contents", show_label=False, elem_id="dataset_gallery").style(grid=5, height="auto")
+    view_dataset_button.click(view_dataset, inputs=[], outputs=dataset_gallery)
+
+# Launch the interface
+iface.launch()