91 lines
3.4 KiB
Python
91 lines
3.4 KiB
Python
"""CLIP-based indoor room classifier. Runs ViT-B/32 on CPU, ~100ms per image."""
|
|
|
|
import io
|
|
import logging
|
|
import os
|
|
|
|
import open_clip
|
|
import torch
|
|
from PIL import Image
|
|
from flask import Flask, jsonify, request
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
|
log = logging.getLogger("clip-server")
|
|
app = Flask(__name__)
|
|
app.config["MAX_CONTENT_LENGTH"] = 50 * 1024 * 1024 # 50 MB max upload
|
|
|
|
# Load once at startup
|
|
MODEL_NAME = os.getenv("CLIP_MODEL", "ViT-B-32")
|
|
PRETRAINED = os.getenv("CLIP_PRETRAINED", "laion2b_s34b_b79k")
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
log.info(f"Loading CLIP model: {MODEL_NAME} pretrained={PRETRAINED} device={DEVICE}")
|
|
model, _, preprocess = open_clip.create_model_and_transforms(MODEL_NAME, pretrained=PRETRAINED)
|
|
model = model.to(DEVICE).eval()
|
|
tokenizer = open_clip.get_tokenizer(MODEL_NAME)
|
|
|
|
# Text prompts for indoor/outdoor classification
|
|
PROMPTS = [
|
|
"a photo of an indoor room interior with floor visible",
|
|
"a photo of an outdoor scene, landscape, street, or non-room image",
|
|
]
|
|
|
|
|
|
@app.route("/health", methods=["GET"])
|
|
def health():
|
|
return jsonify({"status": "ok", "model": MODEL_NAME, "device": DEVICE})
|
|
|
|
|
|
@app.route("/check", methods=["POST"])
|
|
def check():
|
|
"""Classify whether an image is an indoor room photo suitable for floor replacement.
|
|
Expects raw image bytes as the request body. Returns {score: 0-10, passed: bool}."""
|
|
if not request.data:
|
|
return jsonify({"error": "no image data"}), 400
|
|
|
|
try:
|
|
img = Image.open(io.BytesIO(request.data)).convert("RGB")
|
|
except Exception as e:
|
|
return jsonify({"error": f"bad image: {e}"}), 400
|
|
|
|
# Multiple prompt pairs for robust ensemble (indoor vs outdoor)
|
|
prompt_pairs = [
|
|
("a photo of an indoor room interior with floor", "a photo of an outdoor scene"),
|
|
("interior room photograph, real estate photo", "exterior photograph, outdoor photo"),
|
|
("an indoor space with visible flooring", "an outdoor space, landscape, street view"),
|
|
]
|
|
|
|
image_input = preprocess(img).unsqueeze(0).to(DEVICE)
|
|
wins = 0
|
|
gap_sum = 0.0
|
|
for indoor_text, outdoor_text in prompt_pairs:
|
|
text_input = tokenizer([indoor_text, outdoor_text]).to(DEVICE)
|
|
with torch.no_grad():
|
|
image_features = model.encode_image(image_input)
|
|
text_features = model.encode_text(text_input)
|
|
image_features /= image_features.norm(dim=-1, keepdim=True)
|
|
text_features /= text_features.norm(dim=-1, keepdim=True)
|
|
raw = (image_features @ text_features.T)[0] # cosine similarities
|
|
indoor_sim = float(raw[0].item())
|
|
outdoor_sim = float(raw[1].item())
|
|
if indoor_sim > outdoor_sim:
|
|
wins += 1
|
|
gap_sum += indoor_sim - outdoor_sim # positive = indoor wins
|
|
|
|
indoor_wins = wins >= 2 # majority vote (2 out of 3)
|
|
avg_gap = gap_sum / len(prompt_pairs)
|
|
# Map avg_gap to 1-10: gap ~0.05 = score 6, gap ~0.15+ = score 10
|
|
score_1_10 = max(1, min(10, int(5 + avg_gap * 40)))
|
|
passed = indoor_wins
|
|
|
|
log.info(
|
|
f"wins={wins}/3 avg_gap={avg_gap:.3f} score_1_10={score_1_10} passed={passed}"
|
|
)
|
|
return jsonify({"score": score_1_10, "passed": passed})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
port = int(os.getenv("CLIP_PORT", "5100"))
|
|
log.info(f"CLIP server listening on :{port}")
|
|
app.run(host="0.0.0.0", port=port, debug=False)
|