Spaces:
Sleeping
Sleeping
File size: 2,440 Bytes
f548caf cdfb88f f548caf cdfb88f |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import os
import json
import shutil
import tempfile
import uvicorn
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"data/{guid}.jpg"
metadata = {
"file_name": path_in_repo,
"guid": guid,
"ip": ip,
"upload_timestamp": datetime.utcnow().isoformat(),
}
# Write metadata to metadata.jsonl
metadata_file = os.path.join(temp_dir, "metadata.jsonl")
with open(metadata_file, "a") as f:
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
def start_server(port: int = 7860):
uvicorn.run(
"app:app",
host="0.0.0.0",
port=port,
)
if __name__ == "__main__":
start_server() |