Update app.py

app.py

@@ -4,7 +4,8 @@ from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 import io
 import torch
-from transformers import AutoImageProcessor, AutoModelForObjectDetection
+from transformers import AutoImageProcessor, AutoBackbone
+import pytesseract  # OCR
 
 app = FastAPI()
 
@@ -16,41 +17,84 @@ app.add_middleware(
 )
 
 processor = AutoImageProcessor.from_pretrained("czczup/textnet-base")
-model = AutoModelForObjectDetection.from_pretrained("czczup/textnet-base")
+model = AutoBackbone.from_pretrained("czczup/textnet-base")
 model.eval()
 
 @app.post("/detect")
 async def detect_text(file: UploadFile = File(...)):
     try:
+        # Lire image
         image_bytes = await file.read()
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-        inputs = processor(images=image, return_tensors="pt")
-
+        # Entrée TextNet
+        inputs = processor(image, return_tensors="pt")
         with torch.no_grad():
             outputs = model(**inputs)
 
-        results = processor.post_process_object_detection(
-            outputs,
-            threshold=0.3,
-            target_sizes=[image.size[::-1]]
-        )[0]
+        # Feature map et heatmap
+        fm = outputs.feature_maps[-1][0]  # dernière layer
+        heatmap = fm.mean(dim=0).numpy()
+        H, W = heatmap.shape
+        threshold = heatmap.max() * 0.2
+
+        # Points chauds
+        points = [(x, y) for y in range(H) for x in range(W) if heatmap[y, x] > threshold]
+
+        if not points:
+            return JSONResponse([])
 
+        # Regrouper par lignes simples
+        lines = {}
+        for x, y in points:
+            key = int(y / 10)
+            lines.setdefault(key, []).append((x, y))
+
+        # Générer boxes et extraire texte OCR
+        scale_x = image.width / W
+        scale_y = image.height / H
         boxes = []
-        for poly, score in zip(results["polygons"], results["scores"]):
-            poly = poly.tolist()
-            xs = [p[0] for p in poly]
-            ys = [p[1] for p in poly]
+
+        for line in lines.values():
+            xs = [p[0] for p in line]
+            ys = [p[1] for p in line]
+            min_x, max_x = min(xs), max(xs)
+            min_y, max_y = min(ys), max(ys)
+
+            if (max_x - min_x) < 5 or (max_y - min_y) < 2:
+                continue
+
+            # Crop pour OCR
+            crop = image.crop((
+                int(min_x * scale_x),
+                int(min_y * scale_y),
+                int(max_x * scale_x),
+                int(max_y * scale_y)
+            ))
+
+            text = pytesseract.image_to_string(crop, lang='eng').strip()
+
+            if len(text) < 2:
+                continue
+
+            if len(boxes) == 0:
+                boxes.append({
+                    "x": 10,
+                    "y": 10,
+                    "w": 100,
+                    "h": 50,
+                    "text": "Aucun texte détecté"
+                })
 
             boxes.append({
-                "polygon": poly,
-                "x": int(min(xs)),
-                "y": int(min(ys)),
-                "w": int(max(xs) - min(xs)),
-                "h": int(max(ys) - min(ys)),
-                "score": float(score)
+                "x": int(min_x * scale_x),
+                "y": int(min_y * scale_y),
+                "w": int((max_x - min_x) * scale_x),
+                "h": int((max_y - min_y) * scale_y),
+                "text": text or "texte non reconnu"
             })
-
+            print("BOXES:", boxes)
+    
         return JSONResponse({
             "image_width": image.width,
             "image_height": image.height,
@@ -58,4 +102,4 @@ async def detect_text(file: UploadFile = File(...)):
         })
 
     except Exception as e:
-        return JSONResponse({"success": False, "error": str(e)}, status_code=500)
+        return JSONResponse({"success": False, "error": str(e)}, status_code=500)