Update server1.py

server1.py

@@ -6,7 +6,6 @@ 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")
@@ -17,44 +16,58 @@ model = SamModel.from_pretrained("facebook/sam-vit-base").to(DEVICE)
 processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
 
-def detect_blocks(image_bytes):
-
-    img = Image.open(io.BytesIO(image_bytes))
-    if img.mode != "RGB":
-        img = img.convert("RGB")
-
+def detect_strip_mask(img):
     inputs = processor(img, return_tensors="pt").to(DEVICE)
-
     with torch.no_grad():
-        outputs = model(**inputs)
+        out = model(**inputs)
 
     masks = processor.post_process_masks(
-        outputs.pred_masks.cpu(),
+        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 mask in masks:
+    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
 
-        mask = np.array(mask)
 
-        # forzar a 2D
-        while mask.ndim > 2:
-            mask = mask[0]
+def detect_blocks(image_bytes):
 
-        if mask.ndim != 2:
-            continue
+    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-        ys, xs = np.where(mask > 0.5)
-        if len(xs) == 0:
-            continue
+    # 1) SAM detecta la tira completa
+    x1, y1, x2, y2 = detect_strip_mask(img)
 
-        x1, x2 = int(xs.min()), int(xs.max())
-        y1, y2 = int(ys.min()), int(ys.max())
+    strip = img.crop((x1, y1, x2, y2))
 
-        blocks.append([x1, y1, x2, y2])
+    # 2) se divide en 11 bloques
+    blocks = split_into_11(strip)
 
     return blocks