from fastapi import FastAPI, File, UploadFile, HTTPException, Request
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import HfApi
import os
from dotenv import load_dotenv
import uvicorn
import requests
from io import BytesIO
import re
from urllib.parse import urlparse
from datetime import datetime
import os
import hashlib
import random
import string
load_dotenv()
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 环境变量配置
hf_token = os.getenv("HF_TOKEN")
hf_dataset_id = os.getenv("HF_DATASET_ID")
ACCESS_PASSWORD = os.getenv("ACCESS_PASSWORD", "your_default_password")
PROXY_DOMAIN = os.getenv("PROXY_DOMAIN", "huggingface.co")
# 初始化API并添加token
api = HfApi(token=hf_token)
# 设置通用请求头
headers = {
"Authorization": f"Bearer {hf_token}",
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
}
def is_valid_image_url(url):
try:
parsed = urlparse(url)
return bool(parsed.netloc) and bool(parsed.scheme)
except:
return False
def get_image_extension(content_type):
content_type = content_type.lower()
if 'jpeg' in content_type or 'jpg' in content_type:
return 'jpg'
elif 'png' in content_type:
return 'png'
elif 'gif' in content_type:
return 'gif'
elif 'webp' in content_type:
return 'webp'
return 'jpg'
def generate_random_string(length=6):
"""Generate a random string of fixed length"""
letters = string.ascii_lowercase + string.digits
return ''.join(random.choice(letters) for _ in range(length))
def generate_unique_filename(original_filename):
current_date = datetime.now().strftime('%Y-%m-%d')
ext = os.path.splitext(original_filename)[1]
if not ext:
ext = '.jpg'
timestamp = datetime.now().strftime('%H%M%S')
random_str = generate_random_string()
content = f"{timestamp}{random_str}{original_filename}".encode('utf-8')
hash_value = hashlib.md5(content).hexdigest()[:12]
unique_filename = f"{hash_value}{ext}"
return f"{current_date}/{unique_filename}"
@app.get("/", response_class=HTMLResponse)
async def root():
return """
Login
"""
@app.post("/verify")
async def verify_password(request: Request):
try:
data = await request.json()
password = data.get("password")
if password == ACCESS_PASSWORD:
return HTMLResponse("""
HuggingFace Dataset Images
""")
else:
raise HTTPException(status_code=401, detail="Invalid password")
except Exception as e:
raise HTTPException(status_code=400, detail=str(e))
@app.post("/upload/")
async def upload_image(file: UploadFile = File(...)):
if not file:
raise HTTPException(status_code=400, detail="No file uploaded")
if not file.content_type.startswith('image/'):
raise HTTPException(status_code=400, detail="File must be an image")
try:
contents = await file.read()
# Generate unique path
unique_path = generate_unique_filename(file.filename)
# Upload to HuggingFace
response = api.upload_file(
path_or_fileobj=contents,
path_in_repo=f"images/{unique_path}",
repo_id=hf_dataset_id,
repo_type="dataset",
token=hf_token
)
# 修改返回URL格式
image_url = f"https://{PROXY_DOMAIN}/datasets/{hf_dataset_id}/resolve/main/images/{unique_path}"
return {"url": image_url}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/fetch-url/")
async def fetch_url(request: Request):
try:
data = await request.json()
url = data.get("url")
if not url:
raise HTTPException(status_code=400, detail="No URL provided")
if not is_valid_image_url(url):
raise HTTPException(status_code=400, detail="Invalid image URL")
response = requests.get(url, headers=headers, timeout=10)
if not response.ok:
raise HTTPException(status_code=400, detail="Failed to fetch image")
content_type = response.headers.get('content-type', '')
if not content_type.startswith('image/'):
raise HTTPException(status_code=400, detail="URL does not point to an image")
original_filename = url.split('/')[-1]
if not original_filename or '.' not in original_filename:
ext = get_image_extension(content_type)
original_filename = f"downloaded_image.{ext}"
unique_path = generate_unique_filename(original_filename)
response = api.upload_file(
path_or_fileobj=response.content,
path_in_repo=f"images/{unique_path}",
repo_id=hf_dataset_id,
repo_type="dataset",
token=hf_token
)
# 修改返回URL格式
image_url = f"https://{PROXY_DOMAIN}/datasets/{hf_dataset_id}/resolve/main/images/{unique_path}"
return {"url": image_url, "filename": original_filename}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)