Update app.py

app.py

@@ -1,116 +1,33 @@
-import io
-import os
-from typing import Tuple
-
 import gradio as gr
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.responses import Response, JSONResponse
+from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
-from rembg import remove, new_session
-
-# ---- модель для удаления фона (CPU-дружелюбная) ----
-MODEL_NAME = os.getenv("REMBG_MODEL", "isnet-general-use")
-SESSION = new_session(MODEL_NAME)
-
-# ---------------- utils ----------------
-def _ensure_rgba(img: Image.Image) -> Image.Image:
-    return img.convert("RGBA") if img.mode != "RGBA" else img
-
-def _remove_bg_pil(img: Image.Image) -> Image.Image:
-    # rembg лучше кормить байтами; так стабильнее с профилями/EXIF
-    buf = io.BytesIO()
-    _ensure_rgba(img).save(buf, format="PNG")
-    out = remove(buf.getvalue(), session=SESSION)  # bytes PNG с альфой
-    return Image.open(io.BytesIO(out)).convert("RGBA")
-
-def _export(image: Image.Image, fmt: str = "PNG") -> Tuple[bytes, str]:
-    fmt = fmt.upper()
-    bio = io.BytesIO()
-    if fmt == "PNG":
-        image.save(bio, format="PNG", optimize=True)
-        return bio.getvalue(), "image/png"
-    elif fmt == "WEBP":
-        image.save(bio, format="WEBP", lossless=True)
-        return bio.getvalue(), "image/webp"
-    elif fmt in ("JPG", "JPEG"):
-        image.convert("RGB").save(bio, format="JPEG", quality=92, optimize=True, progressive=True)
-        return bio.getvalue(), "image/jpeg"
-    else:
-        raise ValueError("Unsupported format")
-
-# ---------------- FastAPI ----------------
-fast = FastAPI(title="Background Removal API")
-
-@fast.get("/")
-def root():
-    return {
-        "ok": True,
-        "service": "background-removal",
-        "model": MODEL_NAME,
-        "endpoints": {
-            "ui": "/",
-            "health": "/healthz",
-            "remove": "POST /remove-bg (multipart/form-data, field: file, optional ?fmt=png|webp|jpg)",
-        },
-    }
 
-@fast.get("/healthz")
-def healthz():
-    return {"status": "ok", "model": MODEL_NAME}
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-@fast.post("/remove-bg")
-async def remove_bg(file: UploadFile = File(...), fmt: str = "png"):
-    if not file.content_type or not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=415, detail="expected image/*")
-
-    raw = await file.read()
-    if not raw:
-        raise HTTPException(status_code=400, detail="empty file")
-
-    try:
-        src = Image.open(io.BytesIO(raw))
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=f"bad image: {e}")
+def generate_caption(image):
+    # Now directly using the PIL Image object
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model.generate(**inputs)
+    caption = processor.decode(outputs[0], skip_special_tokens=True)
+    return caption
 
+def caption_image(image):
+    """
+    Takes a PIL Image input and returns a caption.
+    """
     try:
-        cut = _remove_bg_pil(src)
-        fmt_norm = {"png": "PNG", "webp": "WEBP", "jpg": "JPEG", "jpeg": "JPEG"}[fmt.lower()]
-        body, media = _export(cut, fmt_norm)
-    except KeyError:
-        raise HTTPException(status_code=400, detail="fmt must be png|webp|jpg")
+        caption = generate_caption(image)
+        return caption
     except Exception as e:
-        return JSONResponse(status_code=400, content={"ok": False, "error": str(e)})
-
-    base = os.path.splitext(file.filename or "image")[0]
-    ext = {"image/png": "png", "image/webp": "webp", "image/jpeg": "jpg"}[media]
-    return Response(
-        content=body,
-        media_type=media,
-        headers={"Content-Disposition": f'inline; filename="{base}_no_bg.{ext}"'},
-    )
-
-# ---------------- Gradio UI ----------------
-def gr_remove(image: Image.Image, output_format: str):
-    """Возвращает картинку с прозрачным фоном (или перекодированную)."""
-    cut = _remove_bg_pil(image)
-    if output_format == "PNG (transparent)":
-        return cut
-    # если пользователь хочет JPG, просто вернём как RGB (фон будет чёрный у прозрачности)
-    if output_format == "JPG":
-        return cut.convert("RGB")
-    # WEBP отрендерится корректно в браузере
-    return cut
-
-demo = gr.Interface(
-    fn=gr_remove,
-    inputs=[
-        gr.Image(type="pil", label="Загрузите картинку"),
-        gr.Dropdown(choices=["PNG (transparent)", "WEBP", "JPG"], value="PNG (transparent)", label="Формат"),
-    ],
-    outputs=gr.Image(type="pil", label="Результат"),
-    title="Remove Background",
-    description="Простой вырезатель фона. Работает в брау��ере и по API одновременно.",
+        return f"An error occurred: {str(e)}"
+
+iface = gr.Interface(
+    fn=caption_image,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Image Captioning with BLIP",
+    description="Upload an image to generate a caption."
 )
 
-# Смонтируем Gradio внутрь FastAPI: UI на "/", API-роуты живут рядом.
-app = gr.mount_gradio_app(fast, demo, path="/")
+iface.launch().