DatasetUpdate / app.py
clone3's picture
Update app.py
06266a4 verified
Raw
History Blame Contribute Delete
3.38 kB
import os
import json
import shutil
import tempfile
from datetime import datetime
from flask import Flask, request, jsonify
from huggingface_hub import login, HfApi
from werkzeug.utils import secure_filename
app = Flask(__name__)
# Load environment variables
HF_TOKEN = os.getenv("HF_TOKEN")
REPO_ID = os.getenv("HF_REPO_ID")
# Initialize Hugging Face API
login(token=HF_TOKEN)
api = HfApi()
@app.route('/upload', methods=['POST'])
def upload_image():
try:
# Check if required data is present
if 'image' not in request.files or 'guid' not in request.form or 'ip' not in request.form:
return jsonify({"error": "Missing image, guid, or ip"}), 400
image = request.files['image']
guid = secure_filename(request.form['guid']) # Sanitize GUID
ip = request.form['ip']
# Validate image
if not image or not image.filename:
return jsonify({"error": "No valid image provided"}), 400
# Create temporary directory
with tempfile.TemporaryDirectory() as temp_dir:
# Create images folder
temp_image_dir = os.path.join(temp_dir, "images")
os.makedirs(temp_image_dir, exist_ok=True)
# Save image with guid.jpg extension
temp_image_path = os.path.join(temp_image_dir, f"{guid}.jpg")
image.save(temp_image_path)
# Create metadata
path_in_repo = f"images/{guid}.jpg"
metadata = {
"file_name": path_in_repo,
"guid": guid,
"ip": ip,
"upload_timestamp": datetime.utcnow().isoformat(),
}
# Path for temporary metadata file
metadata_file = os.path.join(temp_dir, "metadata.jsonl")
# Try to download existing metadata.jsonl from the repository
try:
api.hf_hub_download(
repo_id=REPO_ID,
filename="metadata.jsonl",
repo_type="dataset",
local_dir=temp_dir,
)
# Read existing metadata and append new entry
with open(metadata_file, "r") as f:
existing_metadata = f.readlines()
except Exception:
# If metadata.jsonl doesn't exist, start with an empty list
existing_metadata = []
# Append new metadata
with open(metadata_file, "w") as f:
# Write existing metadata back
for line in existing_metadata:
f.write(line)
# Append new metadata
f.write(json.dumps(metadata) + "\n")
# Upload to Hugging Face
api.upload_folder(
folder_path=temp_dir,
path_in_repo="",
repo_id=REPO_ID,
repo_type="dataset"
)
return jsonify({"message": "Image and metadata uploaded successfully"}), 200
except Exception as e:
return jsonify({"error": str(e)}), 500
if __name__ == '__main__':
print('''curl -X POST https://<your-space-name>.hf.space/upload \\
-F "image=@/path/to/image.jpg" \\
-F "guid=test-guid" \\
-F "ip=127.0.0.1"''')
port = int(os.environ.get("PORT", 7860))
app.run(host="0.0.0.0", port=port)