Update server1.py

server1.py

@@ -2,54 +2,67 @@
 
 import io
 import numpy as np
-import cv2
 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")
 
-def contar_cuadrados(image_bytes):
+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")
+
 
-    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    np_img = np.array(img)
+def detect_blocks(image_bytes):
 
-    gray = cv2.cvtColor(np_img, cv2.COLOR_RGB2GRAY)
-    blur = cv2.GaussianBlur(gray, (5,5), 0)
+    img = Image.open(io.BytesIO(image_bytes))
+    if img.mode != "RGB":
+        img = img.convert("RGB")
 
-    # bordes
-    edges = cv2.Canny(blur, 50, 150)
+    # preparar imagen
+    inputs = processor(img, return_tensors="pt").to(DEVICE)
 
-    # contornos
-    contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+    # obtener mascaras
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-    cuadrados = []
+    masks = processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu()
+    )[0].numpy()
 
-    for cnt in contours:
-        approx = cv2.approxPolyDP(cnt, 0.04 * cv2.arcLength(cnt, True), True)
+    blocks = []
 
-        # queremos formas cuadradas / rectangulares
-        if len(approx) == 4 and cv2.contourArea(cnt) > 200:
-            x, y, w, h = cv2.boundingRect(cnt)
+    # convertir mascaras en bounding boxes
+    for mask in masks:
+        ys, xs = np.where(mask > 0.5)
+        if len(xs) == 0:
+            continue
 
-            if w > 10 and h > 10:
-                cuadrados.append({
-                    "bbox": [int(x), int(y), int(x+w), int(y+h)]
-                })
+        x1, x2 = int(xs.min()), int(xs.max())
+        y1, y2 = int(ys.min()), int(ys.max())
 
-    return cuadrados
+        blocks.append([x1, y1, x2, y2])
+
+    return blocks
 
 
 @app.post("/strip/")
 async def strip(front: UploadFile = File(...)):
     try:
-        bytes_img = await front.read()
-        cuadrados = contar_cuadrados(bytes_img)
+        img_bytes = await front.read()
+        blocks = detect_blocks(img_bytes)
 
         return JSONResponse(content={
             "code": 200,
-            "count": len(cuadrados),
-            "cuadrados": cuadrados
+            "num_blocks": len(blocks),
+            "blocks": blocks
         })
+
     except Exception as e:
         return JSONResponse(content={"code": 500, "error": str(e)})