added images

main.py

@@ -3,13 +3,15 @@ import re
 import time
 
 import dotenv
+import fitz  # PyMuPDF
 import pandas as pd
 import requests
 import schedule
 import srsly
 from bs4 import BeautifulSoup
-from datasets import Dataset, load_dataset
+from datasets import Dataset, Image, load_dataset
 from huggingface_hub import create_repo, login, whoami
+from PIL import Image as PILImage
 from retry import retry
 from tqdm.auto import tqdm
 
@@ -17,7 +19,7 @@ dotenv.load_dotenv()
 login(token=os.environ.get("HF_TOKEN"))
 
 hf_user = whoami(os.environ.get("HF_TOKEN"))["name"]
-HF_REPO_ID = f"{hf_user}/zotero-answer-ai-articles"
+HF_REPO_ID = f"{hf_user}/zotero-articles"
 
 
 ########################################################
@@ -64,7 +66,7 @@ def get_zotero_items(debug=False):
         print(f"# items fetched {len(items)}")
 
         if debug:
-            if len(items) > 300:
+            if len(items) > 200:
                 break
 
     return items
@@ -309,6 +311,98 @@ def parse_markdown_content(md_content, arxiv_id):
     return parsed
 
 
+########################################################
+### Image Dataset
+########################################################
+
+
+def download_arxiv_pdf(arxiv_id):
+    arxiv_id = arxiv_id.split("v")[0]
+    url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
+    response = requests.get(url)
+    if response.status_code == 200:
+        return response.content
+    else:
+        raise Exception(f"Failed to download PDF. Status code: {response.status_code}")
+
+
+def pdf_to_jpegs(pdf_content, output_folder):
+    # Create output folder if it doesn't exist
+    os.makedirs(output_folder, exist_ok=True)
+
+    # Open the PDF
+    doc = fitz.open(stream=pdf_content, filetype="pdf")
+
+    # Iterate through pages
+    for page_num in range(len(doc)):
+        page = doc.load_page(page_num)
+
+        # Convert page to image
+        pix = page.get_pixmap()
+
+        # Save image as JPEG
+        image_path = os.path.join(output_folder, f"page_{page_num + 1}.jpg")
+        pix.save(image_path)
+        # print(f"Saved {image_path}")
+
+    doc.close()
+
+
+def save_arxiv_article_images(arxiv_id):
+    output_folder = os.path.join("data", "arxiv_images", arxiv_id)
+    try:
+        pdf_content = download_arxiv_pdf(arxiv_id)
+        pdf_to_jpegs(pdf_content, output_folder)
+    except Exception as e:
+        print(f"An error occurred: {str(e)}")
+
+
+def create_hf_image_dataset(base_dir):
+    data = []
+
+    # Walk through the directory
+    for root, dirs, files in os.walk(base_dir):
+        for file in files:
+            if file.endswith(".jpg"):
+                # Extract arxiv_id from the path
+                arxiv_id = os.path.basename(root)
+
+                # Extract page number from the filename
+                match = re.search(r"page_(\d+)", file)
+                if match:
+                    page_number = int(match.group(1))
+                else:
+                    continue  # Skip if page number can't be extracted
+
+                # Full path to the image
+                image_path = os.path.join(root, file)
+
+                # Open the image to get its size
+                with PILImage.open(image_path) as img:
+                    width, height = img.size
+
+                # Add the data
+                data.append(
+                    {"image": image_path, "arxiv_id": arxiv_id, "page_number": page_number, "width": width, "height": height}
+                )
+
+    # Create the dataset
+    dataset = Dataset.from_dict(
+        {
+            "image": [d["image"] for d in data],
+            "arxiv_id": [d["arxiv_id"] for d in data],
+            "page_number": [d["page_number"] for d in data],
+            "width": [d["width"] for d in data],
+            "height": [d["height"] for d in data],
+        }
+    )
+
+    # Cast the image column to Image
+    dataset = dataset.cast_column("image", Image())
+
+    return dataset
+
+
 ########################################################
 ### HF UPLOAD
 ########################################################
@@ -324,6 +418,10 @@ def upload_to_hf(abstract_df, contents_df, processed_arxiv_ids):
         exist_ok=True,
     )
 
+    # upload image dataset
+    img_ds = create_hf_image_dataset("data/arxiv_images")
+    img_ds.push_to_hub(repo_id, "images", token=os.environ.get("HF_TOKEN"))
+
     # push id_to_abstract
     abstract_ds = Dataset.from_pandas(abstract_df)
     abstract_ds.push_to_hub(repo_id, "abstracts", token=os.environ.get("HF_TOKEN"))
@@ -369,6 +467,10 @@ def main():
 
     processed_arxiv_ids = set()
     for item in arxiv_items:
+        # download images --
+        save_arxiv_article_images(item["arxiv_id"])
+
+        # parse html
         try:
             item["contents"] = parse_html_content(item["raw_html"])
             processed_arxiv_ids.add(item["arxiv_id"])

requirements.txt

@@ -10,4 +10,6 @@ python-dotenv
 beautifulsoup4
 retry
 pandas
-datasets
+datasets
+PyMuPDF
+pillow