Spaces:
Sleeping
Sleeping
File size: 2,531 Bytes
21a0de8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | import os
import requests
from urllib.parse import quote
# Note: This is designed to run on a Cloud GPU instance with sufficient storage.
# The script will download the .ndjson files for a subset of the Quick, Draw! categories.
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA_DIR = os.path.join(base_dir, "data", "quickdraw_data")
BASE_URL = "https://storage.googleapis.com/quickdraw_dataset/full/simplified/"
def download_file(url, local_path):
# Only download if it doesn't already exist
if os.path.exists(local_path):
print(f"Already exists: {local_path}")
return
print(f"Downloading {url} ...")
response = requests.get(url, stream=True)
if response.status_code == 200:
with open(local_path, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
print("Success.")
else:
print(f"Failed with status code: {response.status_code}")
def main():
if not os.path.exists(DATA_DIR):
os.makedirs(DATA_DIR)
# Load from vocab.txt
vocab_path = os.path.join(base_dir, "data", "vocab.txt")
vocab_words = []
if os.path.exists(vocab_path):
with open(vocab_path, 'r') as f:
vocab_words = [line.strip().lower() for line in f if line.strip()]
# Fetch the official 345 QuickDraw categories from GitHub
categories_url = "https://raw.githubusercontent.com/googlecreativelab/quickdraw-dataset/master/categories.txt"
try:
r = requests.get(categories_url)
official_categories = [c.strip().lower() for c in r.text.split('\n') if c.strip()]
except:
official_categories = ["apple", "banana", "cat", "dog", "house"] # fallback
# Intersect
subset_categories = [cat for cat in official_categories if not vocab_words or cat in vocab_words]
if not subset_categories:
print("No valid QuickDraw categories found in vocab.txt! Downloading fallback categories...")
subset_categories = ["apple", "banana", "cat", "dog", "house"]
print(f"Found {len(subset_categories)} valid QuickDraw categories to download...")
for category in subset_categories:
# Encode URL because some categories might have spaces e.g., 'alarm clock'
filename = f"{category}.ndjson"
url = BASE_URL + quote(filename)
local_path = os.path.join(DATA_DIR, filename)
download_file(url, local_path)
if __name__ == "__main__":
main()
|