90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
import os
|
|
import requests
|
|
from urllib.parse import urlparse
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
|
|
def download_image(url: str, filename: str = None) -> str:
|
|
"""
|
|
Download an image from URL and save it in data/temp/ folder.
|
|
|
|
Args:
|
|
url (str): Image URL
|
|
filename (str, optional): Custom filename. If None, extracted from URL.
|
|
|
|
Returns:
|
|
str: Full path to the downloaded image
|
|
"""
|
|
try:
|
|
# Get project root directory (where your main script is)
|
|
root_dir = Path(os.path.dirname(os.path.abspath(__file__))).parent
|
|
|
|
# Create data/temp folder structure
|
|
temp_dir = root_dir / "data" / "temp"
|
|
temp_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Generate filename if not provided
|
|
if not filename:
|
|
parsed_url = urlparse(url)
|
|
filename = os.path.basename(parsed_url.path)
|
|
if not filename or "." not in filename:
|
|
# Fallback filename
|
|
ext = filename.split('.')[-1] if '.' in filename else 'jpg'
|
|
filename = f"image_{hash(url) % 100000}.{ext}"
|
|
|
|
# Ensure filename has extension
|
|
if '.' not in filename:
|
|
filename += ".jpg"
|
|
|
|
file_path = temp_dir / filename
|
|
|
|
# Download the image
|
|
response = requests.get(url, stream=True, timeout=30)
|
|
response.raise_for_status()
|
|
|
|
# Save image
|
|
with open(file_path, 'wb') as f:
|
|
for chunk in response.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
|
|
print(f"✅ Image downloaded: {file_path}")
|
|
return str(file_path)
|
|
|
|
except Exception as e:
|
|
print(f"❌ Failed to download image: {e}")
|
|
raise
|
|
|
|
def read_image(image_path: str) -> Image.Image:
|
|
"""
|
|
Read an image from the given path and return a PIL Image object.
|
|
|
|
Args:
|
|
image_path (str): Path to the image file
|
|
|
|
Returns:
|
|
PIL.Image.Image: Loaded image
|
|
|
|
Raises:
|
|
FileNotFoundError: If image doesn't exist
|
|
Exception: For other image loading errors
|
|
"""
|
|
try:
|
|
if not os.path.exists(image_path):
|
|
raise FileNotFoundError(f"Image not found at path: {image_path}")
|
|
|
|
# Open the image
|
|
image = Image.open(image_path)
|
|
|
|
# Convert to RGB (important for DINOv2 and most models)
|
|
if image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
|
|
print(f"✅ Image loaded successfully: {image_path} | Size: {image.size}")
|
|
return image
|
|
|
|
except FileNotFoundError as e:
|
|
print(f"❌ File not found: {e}")
|
|
raise
|
|
except Exception as e:
|
|
print(f"❌ Failed to read image: {e}")
|
|
raise |