Spaces:
Sleeping
Sleeping
| import io | |
| import uuid | |
| from contextlib import asynccontextmanager | |
| from pathlib import Path | |
| import numpy as np | |
| from fastapi import FastAPI, File, Form, HTTPException, UploadFile | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import StreamingResponse | |
| from PIL import Image, ImageFilter | |
| from rembg import new_session, remove | |
| BACKGROUNDS_DIR = Path("/app/backgrounds") | |
| _BG_WIDTH, _BG_HEIGHT = 1920, 1080 | |
| _BUILTIN_BACKGROUNDS = [ | |
| { | |
| "filename": "studio_white.jpg", | |
| "stops": [((210,210,215),0.0),((235,235,240),0.3),((248,248,252),0.6),((230,230,235),1.0)], | |
| "floor": (180, 180, 185), "vignette": 0.35, | |
| }, | |
| { | |
| "filename": "showroom_night.jpg", | |
| "stops": [((5,5,20),0.0),((10,12,40),0.3),((18,20,65),0.6),((8,10,30),1.0)], | |
| "floor": (12, 14, 45), "vignette": 0.5, | |
| }, | |
| { | |
| "filename": "canyon_sunset.jpg", | |
| "stops": [((180,60,20),0.0),((220,100,30),0.25),((240,150,40),0.5),((255,180,60),0.75),((200,80,25),1.0)], | |
| "floor": (120, 55, 20), "vignette": 0.4, | |
| }, | |
| { | |
| "filename": "mountain_road.jpg", | |
| "stops": [((80,160,220),0.0),((100,180,230),0.3),((50,120,80),0.6),((30,80,50),1.0)], | |
| "floor": (25, 70, 40), "vignette": 0.3, | |
| }, | |
| { | |
| "filename": "urban_garage.jpg", | |
| "stops": [((25,25,28),0.0),((38,38,42),0.3),((50,50,55),0.6),((30,30,34),1.0)], | |
| "floor": (20, 20, 22), "vignette": 0.45, | |
| }, | |
| ] | |
| def _linear_gradient(stops): | |
| arr = np.zeros((_BG_HEIGHT, _BG_WIDTH, 3), dtype=np.uint8) | |
| xs = np.linspace(0, 1, _BG_WIDTH) | |
| for xi, t in enumerate(xs): | |
| lc, lp = stops[0] | |
| rc, rp = stops[-1] | |
| for i in range(len(stops) - 1): | |
| c0, p0 = stops[i] | |
| c1, p1 = stops[i + 1] | |
| if p0 <= t <= p1: | |
| lc, lp, rc, rp = c0, p0, c1, p1 | |
| break | |
| blend = 0.0 if rp == lp else max(0.0, min(1.0, (t - lp) / (rp - lp))) | |
| arr[:, xi] = [int(lc[j] + (rc[j] - lc[j]) * blend) for j in range(3)] | |
| return Image.fromarray(arr, "RGB") | |
| def _add_floor(img, floor_color, ratio=0.78): | |
| arr = np.array(img).astype(np.float32) | |
| fy = int(_BG_HEIGHT * ratio) | |
| for y in range(fy, _BG_HEIGHT): | |
| a = (y - fy) / (_BG_HEIGHT - fy) | |
| arr[y] = arr[y] * (1 - a * 0.6) + np.array(floor_color) * (a * 0.6) | |
| return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), "RGB") | |
| def _vignette(img, strength=0.4): | |
| arr = np.array(img).astype(np.float32) | |
| cy, cx = _BG_HEIGHT / 2, _BG_WIDTH / 2 | |
| y, x = np.mgrid[0:_BG_HEIGHT, 0:_BG_WIDTH] | |
| dist = np.sqrt(((x - cx) / cx) ** 2 + ((y - cy) / cy) ** 2) | |
| arr *= (1.0 - np.clip(dist * strength, 0, strength))[:, :, np.newaxis] | |
| return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), "RGB") | |
| def generate_backgrounds(): | |
| BACKGROUNDS_DIR.mkdir(parents=True, exist_ok=True) | |
| for bg in _BUILTIN_BACKGROUNDS: | |
| path = BACKGROUNDS_DIR / bg["filename"] | |
| if path.exists(): | |
| print(f" [bg] exists: {bg['filename']}") | |
| continue | |
| img = _linear_gradient(bg["stops"]) | |
| img = _add_floor(img, bg["floor"]) | |
| img = _vignette(img, bg["vignette"]) | |
| img.save(str(path), "JPEG", quality=95, optimize=True) | |
| print(f" [bg] generated: {bg['filename']}") | |
| print(f" [bg] done") | |
| async def lifespan(app_: FastAPI): | |
| generate_backgrounds() | |
| yield | |
| app = FastAPI( | |
| title="Car Background Removal API", | |
| description="Remove car backgrounds and composite onto showroom backgrounds.", | |
| version="1.0.0", | |
| lifespan=lifespan, | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| rembg_session = new_session("u2net") | |
| def refresh_backgrounds(): | |
| return { | |
| bg.stem: bg | |
| for bg in BACKGROUNDS_DIR.glob("*") | |
| if bg.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"} | |
| } | |
| def remove_background(image): | |
| img_bytes = io.BytesIO() | |
| image.save(img_bytes, format="PNG") | |
| img_bytes.seek(0) | |
| result_bytes = remove(img_bytes.read(), session=rembg_session) | |
| result = Image.open(io.BytesIO(result_bytes)).convert("RGBA") | |
| return result | |
| def add_drop_shadow(car_rgba, shadow_offset=(20, 40), shadow_blur_radius=35, shadow_opacity=140): | |
| alpha = car_rgba.split()[3] | |
| shadow_alpha = Image.new("L", car_rgba.size, 0) | |
| shadow_alpha.paste(alpha, (0, 0)) | |
| dark_shadow = Image.new("RGBA", car_rgba.size, (0, 0, 0, shadow_opacity)) | |
| dark_shadow.putalpha(shadow_alpha) | |
| blurred = dark_shadow.filter(ImageFilter.GaussianBlur(radius=shadow_blur_radius)) | |
| canvas_w = car_rgba.width + abs(shadow_offset[0]) + shadow_blur_radius * 2 | |
| canvas_h = car_rgba.height + abs(shadow_offset[1]) + shadow_blur_radius * 2 | |
| canvas = Image.new("RGBA", (canvas_w, canvas_h), (0, 0, 0, 0)) | |
| pad_x = shadow_blur_radius | |
| pad_y = shadow_blur_radius | |
| canvas.paste(blurred, (pad_x + shadow_offset[0], pad_y + shadow_offset[1]), blurred) | |
| canvas.paste(car_rgba, (pad_x, pad_y), car_rgba) | |
| return canvas, (pad_x, pad_y) | |
| def fit_car_to_background(car_with_shadow, background, car_height_ratio=0.65): | |
| bg_w, bg_h = background.size | |
| shadow_w, shadow_h = car_with_shadow.size | |
| target_height = int(bg_h * car_height_ratio) | |
| scale = target_height / shadow_h | |
| target_width = int(shadow_w * scale) | |
| resized = car_with_shadow.resize((target_width, target_height), Image.LANCZOS) | |
| paste_x = (bg_w - target_width) // 2 | |
| paste_y = bg_h - target_height - int(bg_h * 0.05) | |
| return resized, paste_x, paste_y | |
| def composite_car_on_background(car_rgba, background): | |
| car_w, car_h = car_rgba.size | |
| shadow_blur = max(8, int(car_h * 0.025)) | |
| shadow_offset = (int(car_w * 0.008), int(car_h * 0.03)) | |
| car_with_shadow, _ = add_drop_shadow( | |
| car_rgba, | |
| shadow_offset=shadow_offset, | |
| shadow_blur_radius=shadow_blur, | |
| shadow_opacity=120, | |
| ) | |
| resized_car, paste_x, paste_y = fit_car_to_background(car_with_shadow, background) | |
| result = background.copy().convert("RGBA") | |
| result.paste(resized_car, (paste_x, paste_y), resized_car) | |
| return result.convert("RGB") | |
| def root(): | |
| return { | |
| "message": "Car Background Removal API is running.", | |
| "endpoints": { | |
| "POST /process": "Remove car background and composite onto a showroom background.", | |
| "GET /backgrounds": "List available background IDs.", | |
| "POST /upload-background": "Upload a custom background image.", | |
| }, | |
| } | |
| def list_backgrounds(): | |
| backgrounds = refresh_backgrounds() | |
| return {"available_background_ids": list(backgrounds.keys()), "count": len(backgrounds)} | |
| async def upload_background(background_image: UploadFile = File(...)): | |
| if background_image.content_type not in {"image/jpeg", "image/png", "image/webp"}: | |
| raise HTTPException(status_code=400, detail="Unsupported image type.") | |
| raw_bytes = await background_image.read() | |
| if len(raw_bytes) == 0: | |
| raise HTTPException(status_code=400, detail="Uploaded file is empty.") | |
| try: | |
| img = Image.open(io.BytesIO(raw_bytes)).convert("RGB") | |
| except Exception as exc: | |
| raise HTTPException(status_code=400, detail=f"Could not open image: {exc}") | |
| bg_id = f"custom_{uuid.uuid4().hex[:12]}" | |
| save_path = BACKGROUNDS_DIR / f"{bg_id}.jpg" | |
| try: | |
| img.save(save_path, "JPEG", quality=92, optimize=True) | |
| except Exception as exc: | |
| raise HTTPException(status_code=500, detail=f"Could not save background: {exc}") | |
| return {"background_id": bg_id, "filename": save_path.name, "width": img.width, "height": img.height} | |
| async def process_car_image( | |
| car_image: UploadFile = File(...), | |
| background_id: str = Form(...), | |
| ): | |
| if car_image.content_type not in {"image/jpeg", "image/png", "image/webp"}: | |
| raise HTTPException(status_code=400, detail="Unsupported image type.") | |
| backgrounds = refresh_backgrounds() | |
| if background_id not in backgrounds: | |
| raise HTTPException( | |
| status_code=404, | |
| detail=f"Background not found. Available: {list(backgrounds.keys())}", | |
| ) | |
| raw_bytes = await car_image.read() | |
| if len(raw_bytes) == 0: | |
| raise HTTPException(status_code=400, detail="Uploaded file is empty.") | |
| try: | |
| car_img = Image.open(io.BytesIO(raw_bytes)).convert("RGB") | |
| except Exception as exc: | |
| raise HTTPException(status_code=400, detail=f"Could not open image: {exc}") | |
| try: | |
| car_rgba = remove_background(car_img) | |
| except Exception as exc: | |
| raise HTTPException(status_code=500, detail=f"Background removal failed: {exc}") | |
| try: | |
| background_img = Image.open(backgrounds[background_id]).convert("RGB") | |
| except Exception as exc: | |
| raise HTTPException(status_code=500, detail=f"Could not load background: {exc}") | |
| try: | |
| final_image = composite_car_on_background(car_rgba, background_img) | |
| except Exception as exc: | |
| raise HTTPException(status_code=500, detail=f"Compositing failed: {exc}") | |
| output_buffer = io.BytesIO() | |
| final_image.save(output_buffer, format="JPEG", quality=92, optimize=True) | |
| output_buffer.seek(0) | |
| filename = Path(car_image.filename or "car").stem | |
| return StreamingResponse( | |
| output_buffer, | |
| media_type="image/jpeg", | |
| headers={"Content-Disposition": f"attachment; filename={filename}_composited.jpg"}, | |
| ) | |