ffff / server1.py
addgbf's picture
Update server1.py
b8df022 verified
raw
history blame
2.08 kB
# comentarios sin tildes / sin enye
import io, os
import numpy as np
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from PIL import Image
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
import torch
app = FastAPI(title="Accudoctor Strip Analyzer")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# cargar modelo SAM2
sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
sam.to(device=DEVICE)
mask_generator = SamAutomaticMaskGenerator(sam)
def dominant_color(pil_img):
img = pil_img.resize((60, 60))
arr = np.array(img)
arr = arr.reshape((-1, 3))
pixels, counts = np.unique(arr, axis=0, return_counts=True)
dom = pixels[counts.argmax()]
return "#{:02x}{:02x}{:02x}".format(dom[0], dom[1], dom[2])
def analyze_strip(image_bytes):
img = Image.open(io.BytesIO(image_bytes))
if img.mode != "RGB":
img = img.convert("RGB")
np_img = np.array(img)
masks = mask_generator.generate(np_img)
blocks = []
H = np_img.shape[0]
for m in masks:
x, y, w, h = m["bbox"]
aspect = h / (w + 1e-6)
if aspect < 3:
continue
if h < H * 0.04:
continue
crop = img.crop((x, y, x+w, y+h))
color = dominant_color(crop)
blocks.append({
"bbox": [int(x), int(y), int(x+w), int(y+h)],
"color_hex": color,
"y_center": y + h/2
})
blocks = sorted(blocks, key=lambda b: b["y_center"])
for i,b in enumerate(blocks):
b["index"] = i+1
del b["y_center"]
return blocks[:11]
@app.post("/strip/")
async def strip(front: UploadFile = File(...)):
try:
bytes_img = await front.read()
result = analyze_strip(bytes_img)
return JSONResponse(
status_code=200,
content={"code": 200, "blocks": result}
)
except Exception as e:
return JSONResponse(
status_code=200,
content={"code": 500, "error": str(e)}
)