Take screenshots

app.py

@@ -8,18 +8,25 @@ import datetime
 
 app = FastAPI()
 
+
 @app.get("/")
 def index():
-    return HTMLResponse("""
+    return HTMLResponse(
+        """
     <p>Backend for gradio theme gallery.
     <a href="https://huggingface.co/spaces/freddyaboulton/theme-gallery-static">https://huggingface.co/spaces/freddyaboulton/theme-gallery-stati</a>
-    </p>""")
+    </p>"""
+    )
+
 
 scheduler = BackgroundScheduler()
-scheduler.add_job(func=process_spaces, trigger="interval", seconds=360,
-                  next_run_time=datetime.datetime.now())
+scheduler.add_job(
+    func=process_spaces,
+    trigger="interval",
+    seconds=1200,
+    next_run_time=datetime.datetime.now(),
+)
 scheduler.start()
 
 if __name__ == "__main__":
     uvicorn.run(app, port=7860, host="0.0.0.0")
-

processing.py

@@ -10,6 +10,10 @@ import json
 import datetime
 import tqdm
 import requests
+from pathlib import Path
+from screenshot import get_screen_shot
+import boto3
+
 
 class SpaceData(TypedDict):
     id: str
@@ -18,20 +22,32 @@ class SpaceData(TypedDict):
     lastModified: str
     status: str
 
+
 repo = huggingface_hub.Repository(
     local_dir="data",
     repo_type="dataset",
     clone_from="freddyaboulton/gradio-theme-subdomains",
-    token=os.getenv("HF_TOKEN")
+    token=os.getenv("HF_TOKEN"),
 )
 repo.git_pull()
 
+screen_shot_dir = Path("data") / "images"
+screen_shot_dir.mkdir(exist_ok=True, parents=True)
+
+s3_client = boto3.client(
+    "s3",
+    aws_access_key_id=os.getenv("AWS_ACCESS_KEY"),
+    aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
+)
+
 
 api = huggingface_hub.HfApi(token=os.getenv("HF_TOKEN"))
 
+
 def get_theme_preview_spaces() -> List[SpaceInfo]:
     return list(iter(api.list_spaces(filter="gradio-theme")))
 
+
 def get_info(space_name: SpaceInfo) -> SpaceData | None:
     if not space_name.id:
         print(f"no space_name for {space_name}")
@@ -44,17 +60,42 @@ def get_info(space_name: SpaceInfo) -> SpaceData | None:
     if subdomain is None:
         print(f"no subdomain for {space_info.id}")
         return None
-    
-    status = space_info.runtime['stage']
+
+    status = space_info.runtime["stage"]
+    img_id = space_info.id.replace("/", "_")
+    light_file = str(screen_shot_dir / Path(img_id + "_light.jpg"))
+    dark_file = str(screen_shot_dir / Path(img_id + "_dark.jpg"))
+    if status == "RUNNING":
+        get_screen_shot(
+            f"https://{space_info.subdomain}.hf.space?__theme=light", 10, light_file
+        )
+        get_screen_shot(
+            f"https://{space_info.subdomain}.hf.space?__theme=dark", 10, dark_file
+        )
+        s3_client.upload_file(
+            light_file,
+            "gradio-theme-screenshots",
+            img_id + "_light.jpg",
+            ExtraArgs={"ContentType": "image/jpg"},
+        )
+        s3_client.upload_file(
+            dark_file,
+            "gradio-theme-screenshots",
+            img_id + "_dark.jpg",
+            ExtraArgs={"ContentType": "image/jpg"},
+        )
     if status not in ["SLEEPING", "RUNNING", "RUNNING_BUILDING", "BUILDING"]:
         print(f"Space not running, building, or sleeping {space_info.id}")
     elif status == "SLEEPING":
         requests.get(f"https://huggingface.co/spaces/{space_info.id}")
-
-    return {"id": space_info.id, "likes": space_info.likes,
-            "subdomain": f"https://{space_info.subdomain}.hf.space",
-            "lastModified": space_info.lastModified,
-            "status": status} # type: ignore
+    return {
+        "id": space_info.id,
+        "likes": space_info.likes,
+        "subdomain": f"https://{space_info.subdomain}.hf.space",
+        "lastModified": space_info.lastModified,
+        "screenshot_id": img_id,
+        "status": status,
+    }  # type: ignore
 
 
 def get_all_info(spaces: List[SpaceInfo]) -> List[SpaceData]:
@@ -68,7 +109,7 @@ def process_spaces():
 
     all_info = get_all_info(theme_spaces)
 
-    json.dump(all_info, open("data/subdomains.json", "w"))
-    repo.push_to_hub(blocking=False, commit_message=f"Updating data at {datetime.datetime.now()}")
-
-
+    json.dump(all_info, open("data/val_subdomains.json", "w"))
+    repo.push_to_hub(
+        blocking=False, commit_message=f"Updating data at {datetime.datetime.now()}"
+    )

requirements.txt

@@ -1 +1,4 @@
-apscheduler
+apscheduler
+selenium
+webdriver-manager
+boto3

screenshot.py

@@ -0,0 +1,15 @@
+from selenium import webdriver
+from selenium.webdriver import ChromeOptions
+from webdriver_manager.chrome import ChromeDriverManager
+import time
+
+options = ChromeOptions()
+options.add_argument("--headless=new")
+ChromeDriverManager().install()
+
+
+def get_screen_shot(url: str, load_time: int, path: str):
+    driver = webdriver.Chrome(options=options)
+    driver.get(url)
+    time.sleep(load_time)
+    driver.save_screenshot(path)