Spaces:
Running
Running
| import io | |
| import logging | |
| import traceback | |
| import numpy as np | |
| from fastapi import FastAPI, UploadFile, File | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import StreamingResponse | |
| from PIL import Image | |
| from rembg import new_session, remove | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| app = FastAPI() | |
| app.add_middleware(CORSMiddleware, allow_origins=['*'], allow_credentials=True, allow_methods=['*'], allow_headers=['*']) | |
| SESSION = None | |
| def get_session(): | |
| global SESSION | |
| if SESSION is None: | |
| SESSION = new_session('birefnet-general-lite') | |
| return SESSION | |
| def home(): | |
| return {'status': 'ok', 'engine': 'BiRefNet-Lite-Turbo'} | |
| async def remove_bg(file: UploadFile = File(...)): | |
| try: | |
| input_data = await file.read() | |
| img = Image.open(io.BytesIO(input_data)).convert('RGBA') | |
| W, H = img.size | |
| # 1. FAST MASK PASS (to find subject area) | |
| # We use a very small image to find where the person is | |
| low_res = img.resize((256, 256), Image.NEAREST) | |
| mask_low = remove(low_res, session=get_session(), only_mask=True) | |
| # Find Bounding Box coordinates | |
| mask_np = np.array(mask_low) | |
| coords = np.argwhere(mask_np > 0) | |
| if coords.size > 0: | |
| y_min, x_min = coords.min(axis=0) | |
| y_max, x_max = coords.max(axis=0) | |
| # Scale back to original size + padding | |
| pad = 20 | |
| l = max(0, int(x_min * W / 256) - pad) | |
| t = max(0, int(y_min * H / 256) - pad) | |
| r = min(W, int(x_max * W / 256) + pad) | |
| b = min(H, int(y_max * H / 256) + pad) | |
| # 2. CROP SUBJECT AREA | |
| subject_crop = img.crop((l, t, r, b)) | |
| # 3. HIGH QUALITY PASS ON CROP ONLY (Much faster because it's small) | |
| # Resize crop if it's still too huge for CPU | |
| if max(subject_crop.size) > 1024: | |
| subject_crop.thumbnail((1024, 1024), Image.LANCZOS) | |
| cutout_crop = remove(subject_crop, session=get_session()) | |
| # 4. PASTE BACK ON TRANSPARENT CANVAS | |
| final_img = Image.new('RGBA', (W, H), (0, 0, 0, 0)) | |
| # Resize cutout back to the crop rectangle size if thumbnail was used | |
| cutout_crop = cutout_crop.resize((r-l, b-t), Image.LANCZOS) | |
| final_img.paste(cutout_crop, (l, t), cutout_crop) | |
| else: | |
| # Fallback if no subject found | |
| final_img = remove(img.resize((800, 800)), session=get_session()).resize((W, H)) | |
| buf = io.BytesIO() | |
| final_img.save(buf, format='PNG') | |
| buf.seek(0) | |
| return StreamingResponse(buf, media_type='image/png') | |
| except Exception as e: | |
| return {'error': str(e)} | |
| if __name__ == '__main__': | |
| import uvicorn | |
| uvicorn.run(app, host='0.0.0.0', port=7860) |