ffff / server1.py
afdx2's picture
Update server1.py
882b1f1 verified
# comentarios sin tildes / sin enye
import io
import numpy as np
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from PIL import Image
import torch
from transformers import SamModel, SamProcessor
app = FastAPI(title="Accudoctor Strip Analyzer")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model = SamModel.from_pretrained("facebook/sam-vit-base").to(DEVICE)
processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
def detect_strip_mask(img):
inputs = processor(img, return_tensors="pt").to(DEVICE)
with torch.no_grad():
out = model(**inputs)
masks = processor.post_process_masks(
out.pred_masks.cpu(),
inputs["original_sizes"].cpu(),
inputs["reshaped_input_sizes"].cpu()
)[0].numpy()
# cojo la mascara mas grande (la tira)
best_mask = max(masks, key=lambda m: np.sum(m))
best_mask = np.squeeze(best_mask)
ys, xs = np.where(best_mask > 0.5)
x1, x2 = xs.min(), xs.max()
y1, y2 = ys.min(), ys.max()
return x1, y1, x2, y2
def split_into_11(img_strip):
w, h = img_strip.size
block_h = h // 11
blocks = []
for i in range(11):
y1 = i * block_h
y2 = (i + 1) * block_h
crop = img_strip.crop((0, y1, w, y2))
blocks.append({
"index": i + 1,
"bbox": [0, y1, w, y2]
})
return blocks
def detect_blocks(image_bytes):
img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# 1) SAM detecta la tira completa
x1, y1, x2, y2 = detect_strip_mask(img)
strip = img.crop((x1, y1, x2, y2))
# 2) se divide en 11 bloques
blocks = split_into_11(strip)
return blocks
@app.post("/strip/")
async def strip(front: UploadFile = File(...)):
try:
img_bytes = await front.read()
blocks = detect_blocks(img_bytes)
return JSONResponse(content={
"code": 200,
"num_blocks": len(blocks),
"blocks": blocks
})
except Exception as e:
return JSONResponse(content={"code": 500, "error": str(e)})