Update app.py

app.py

@@ -193,32 +193,34 @@
 #         print("Model init warning:", e)
 #     app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), debug=False)
 
+
+
 import os
 import io
 import base64
 import threading
-import tempfile
 import traceback
 from typing import Optional
 
-from flask import Flask, request, jsonify, send_from_directory, send_file
-from PIL import Image, ImageDraw, ImageFont
+from flask import Flask, request, jsonify, send_from_directory
+from PIL import Image
 import numpy as np
 import requests
+import torch
 
-# Set writable cache dirs to avoid matplotlib/fontconfig warnings in containers
+# Set environment variables for CPU-only operation
 os.environ.setdefault("MPLCONFIGDIR", "/tmp/.matplotlib")
 os.environ.setdefault("FONTCONFIG_PATH", "/tmp/.fontconfig")
-# Ensure CPU-only (do not accidentally use GPU)
 os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+os.environ.setdefault("OMP_NUM_THREADS", "4")
+os.environ.setdefault("MKL_NUM_THREADS", "4")
+os.environ.setdefault("OPENBLAS_NUM_THREADS", "4")
+
+# Limit torch threads
+torch.set_num_threads(4)
 
-# --- Imports that may trigger the above warnings ---
-try:
-    import supervision as sv
-    from rfdetr import RFDETRSegPreview
-except Exception as e:
-    # Provide a clearer error at startup if imports fail
-    raise RuntimeError(f"Required library import failed: {e}")
+import supervision as sv
+from rfdetr import RFDETRSegPreview
 
 app = Flask(__name__, static_folder="static", static_url_path="/")
 
@@ -231,10 +233,12 @@ MODEL = None
 
 
 def download_file(url: str, dst: str, chunk_size: int = 8192):
+    """Download file if not exists"""
     if os.path.exists(dst) and os.path.getsize(dst) > 0:
+        print(f"[INFO] Checkpoint already exists at {dst}")
         return dst
     print(f"[INFO] Downloading weights from {url} -> {dst}")
-    r = requests.get(url, stream=True, timeout=60)
+    r = requests.get(url, stream=True, timeout=120)
     r.raise_for_status()
     with open(dst, "wb") as fh:
         for chunk in r.iter_content(chunk_size=chunk_size):
@@ -245,27 +249,29 @@ def download_file(url: str, dst: str, chunk_size: int = 8192):
 
 
 def init_model():
-    """
-    Lazily initialize the RF-DETR model and cache it in global MODEL.
-    Thread-safe.
-    """
+    """Lazily initialize the RF-DETR model and cache it in global MODEL."""
     global MODEL
     with MODEL_LOCK:
         if MODEL is not None:
             return MODEL
         try:
-            # ensure checkpoint present (best-effort)
+            # Ensure checkpoint present
             try:
                 download_file(CHECKPOINT_URL, CHECKPOINT_PATH)
             except Exception as e:
                 print("[WARN] Failed to download checkpoint:", e)
+                if not os.path.exists(CHECKPOINT_PATH):
+                    raise
 
             print("[INFO] Loading RF-DETR model (CPU mode)...")
-            MODEL = RFDETRSegPreview(pretrain_weights=CHECKPOINT_PATH if os.path.exists(CHECKPOINT_PATH) else None)
+            MODEL = RFDETRSegPreview(pretrain_weights=CHECKPOINT_PATH)
+            
+            # Try to optimize for inference
             try:
                 MODEL.optimize_for_inference()
             except Exception as e:
                 print("[WARN] optimize_for_inference() skipped/failed:", e)
+            
             print("[INFO] Model ready.")
             return MODEL
         except Exception:
@@ -274,14 +280,11 @@ def init_model():
 
 
 def decode_data_url(data_url: str) -> Image.Image:
-    """
-    Accepts a data URL (data:image/png;base64,...) or raw base64 and returns PIL.Image (RGB)
-    """
+    """Decode data URL to PIL Image"""
     if data_url.startswith("data:"):
         _, b64 = data_url.split(",", 1)
         data = base64.b64decode(b64)
     else:
-        # assume raw base64 or binary string
         try:
             data = base64.b64decode(data_url)
         except Exception:
@@ -290,117 +293,54 @@ def decode_data_url(data_url: str) -> Image.Image:
 
 
 def encode_pil_to_dataurl(pil_img: Image.Image, fmt="PNG") -> str:
+    """Encode PIL Image to data URL"""
     buf = io.BytesIO()
     pil_img.save(buf, format=fmt)
     return "data:image/{};base64,".format(fmt.lower()) + base64.b64encode(buf.getvalue()).decode("ascii")
 
 
-def overlay_mask_on_image(pil_img: Image.Image, detections, threshold: float = 0.25,
-                          mask_color=(255, 77, 166), alpha=0.45):
+def annotate_segmentation(image: Image.Image, detections: sv.Detections) -> Image.Image:
     """
-    Create annotated PIL image by overlaying per-instance masks (pink) and polygon borders,
-    and add confidence text (best confidence) on the image.
-    Uses supervision-like masks if available, otherwise attempts to use detections.masks.
-    Returns (annotated_pil_rgb, kept_confidences_list)
+    Annotate image with segmentation masks using supervision library.
+    This matches the visualization from rfdetr_seg_infer.py script.
     """
-    base = pil_img.convert("RGBA")
-    W, H = base.size
-
-    masks = getattr(detections, "masks", None)
-    confidences = []
-    try:
-        confidences = [float(x) for x in getattr(detections, "confidence", [])]
-    except Exception:
-        # fallback to 'scores' or empty
-        try:
-            confidences = [float(x) for x in getattr(detections, "scores", [])]
-        except Exception:
-            confidences = []
-
-    if masks is None:
-        # no masks -> return original image and empty list
-        return pil_img.convert("RGB"), []
-
-    # Normalize mask array to (N, H, W)
-    if isinstance(masks, list):
-        masks_arr = np.stack([np.asarray(m, dtype=bool) for m in masks], axis=0)
-    else:
-        masks_arr = np.asarray(masks)
-        # some outputs might be (H, W, N)
-        if masks_arr.ndim == 3 and masks_arr.shape[0] == H and masks_arr.shape[1] == W:
-            masks_arr = masks_arr.transpose(2, 0, 1)
-
-    # overlay image we will composite
-    overlay = Image.new("RGBA", (W, H), (0, 0, 0, 0))
-    kept_confidences = []
-
-    for i in range(masks_arr.shape[0]):
-        conf = confidences[i] if i < len(confidences) else 1.0
-        if conf < threshold:
-            continue
-        mask = masks_arr[i].astype(np.uint8) * 255
-        mask_img = Image.fromarray(mask).convert("L")
-        # if mask size doesn't match, resize
-        if mask_img.size != (W, H):
-            mask_img = mask_img.resize((W, H), resample=Image.NEAREST)
-
-        # color layer with alpha
-        color_layer = Image.new("RGBA", (W, H), mask_color + (0,))
-        # compute per-pixel alpha from mask (0..255) scaled by alpha
-        alpha_mask = mask_img.point(lambda p: int(p * alpha))
-        color_layer.putalpha(alpha_mask)
-        overlay = Image.alpha_composite(overlay, color_layer)
-        kept_confidences.append(float(conf))
-
-    # draw polygon outlines for visual crispness using supervision polygonifier if available
-    try:
-        # try to use supervision polygonizer if detections contains polygons
-        # fallback: create thin white outline by expanding mask boundaries
-        from skimage import measure
-        draw = ImageDraw.Draw(overlay)
-        for i in range(masks_arr.shape[0]):
-            conf = confidences[i] if i < len(confidences) else 1.0
-            if conf < threshold:
-                continue
-            mask = masks_arr[i].astype(np.uint8)
-            # resize mask for contour if needed
-            if mask.shape[1] != W or mask.shape[0] != H:
-                mask_pil = Image.fromarray((mask * 255).astype(np.uint8)).resize((W, H), resample=Image.NEAREST)
-                mask = np.asarray(mask_pil).astype(np.uint8) // 255
-            contours = measure.find_contours(mask, 0.5)
-            for contour in contours:
-                # contour is list of (row, col) -> convert to (x, y)
-                pts = [(float(c[1]), float(c[0])) for c in contour]
-                if len(pts) >= 3:
-                    # draw white outline
-                    draw.line(pts + [pts[0]], fill=(255, 255, 255, 255), width=2)
-    except Exception:
-        # ignore if skimage not available; outlines are optional
-        pass
-
-    annotated = Image.alpha_composite(base, overlay).convert("RGBA")
-
-    # annotate best confidence text (top-left)
-    if kept_confidences:
-        best = max(kept_confidences)
-        draw = ImageDraw.Draw(annotated)
-        try:
-            font = ImageFont.truetype("DejaVuSans-Bold.ttf", size=max(14, W // 32))
-        except Exception:
-            font = ImageFont.load_default()
-        text = f"Confidence: {best:.2f}"
-        tw, th = draw.textsize(text, font=font)
-        pad = 6
-        rect = [6, 6, 6 + tw + pad, 6 + th + pad]
-        draw.rectangle(rect, fill=(0, 0, 0, 180))
-        draw.text((6 + pad // 2, 6 + pad // 2), text, font=font, fill=(255, 255, 255, 255))
-
-    return annotated.convert("RGB"), kept_confidences
+    # Define color palette
+    palette = sv.ColorPalette.from_hex([
+        "#ffff00", "#ff9b00", "#ff8080", "#ff66b2", "#ff66ff", "#b266ff",
+        "#9999ff", "#3399ff", "#66ffff", "#33ff99", "#66ff66", "#99ff00",
+    ])
+
+    # Calculate optimal text scale based on image resolution
+    text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size)
+    
+    # Create annotators
+    mask_annotator = sv.MaskAnnotator(color=palette)
+    polygon_annotator = sv.PolygonAnnotator(color=sv.Color.WHITE)
+    label_annotator = sv.LabelAnnotator(
+        color=palette,
+        text_color=sv.Color.BLACK,
+        text_scale=text_scale,
+        text_position=sv.Position.CENTER_OF_MASS
+    )
+
+    # Create labels with class IDs and confidence scores
+    labels = [
+        f"Tulsi {float(conf):.2f}"
+        for conf in detections.confidence
+    ]
+
+    # Apply annotations
+    out = image.copy()
+    out = mask_annotator.annotate(out, detections)
+    out = polygon_annotator.annotate(out, detections)
+    out = label_annotator.annotate(out, detections, labels)
+    
+    return out
 
 
 @app.route("/", methods=["GET"])
 def index():
-    # serve the static UI file if present
+    """Serve the static UI"""
     index_path = os.path.join(app.static_folder or "static", "index.html")
     if os.path.exists(index_path):
         return send_from_directory(app.static_folder, "index.html")
@@ -421,11 +361,11 @@ def predict():
     except Exception as e:
         return jsonify({"error": f"Model initialization failed: {e}"}), 500
 
-    # parse input
+    # Parse input
     img: Optional[Image.Image] = None
     conf_threshold = 0.25
 
-    # If form file uploaded
+    # Check if file uploaded
     if "file" in request.files:
         file = request.files["file"]
         try:
@@ -434,7 +374,7 @@ def predict():
             return jsonify({"error": f"Invalid uploaded image: {e}"}), 400
         conf_threshold = float(request.form.get("conf", conf_threshold))
     else:
-        # try JSON payload
+        # Try JSON payload
         payload = request.get_json(silent=True)
         if not payload or "image" not in payload:
             return jsonify({"error": "No image provided. Upload 'file' or JSON with 'image' data-url."}), 400
@@ -444,28 +384,69 @@ def predict():
             return jsonify({"error": f"Invalid image data: {e}"}), 400
         conf_threshold = float(payload.get("conf", conf_threshold))
 
-    # run inference
+    # Optionally downscale large images to reduce memory usage
+    MAX_SIZE = 1024
+    if max(img.size) > MAX_SIZE:
+        w, h = img.size
+        scale = MAX_SIZE / float(max(w, h))
+        new_w, new_h = int(round(w * scale)), int(round(h * scale))
+        img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
+        print(f"[INFO] Resized image to {new_w}x{new_h}")
+
+    # Run inference with no_grad for memory efficiency
     try:
-        # set threshold=0.0 in model predict since we'll manually filter by conf_threshold
-        detections = model.predict(img, threshold=0.0)
+        with torch.no_grad():
+            detections = model.predict(img, threshold=conf_threshold)
+        
+        print(f"[INFO] Detected {len(detections)} objects")
+        
+        # Check if detections exist
+        if len(detections) == 0:
+            print("[INFO] No detections above threshold")
+            # Return original image with message
+            data_url = encode_pil_to_dataurl(img, fmt="PNG")
+            return jsonify({
+                "annotated": data_url,
+                "confidences": [],
+                "count": 0
+            })
+        
+        # Annotate image using supervision library
+        annotated_pil = annotate_segmentation(img, detections)
+        
+        # Extract confidence scores
+        confidences = [float(conf) for conf in detections.confidence]
+        
+        # Encode to data URL
+        data_url = encode_pil_to_dataurl(annotated_pil, fmt="PNG")
+        
+        return jsonify({
+            "annotated": data_url,
+            "confidences": confidences,
+            "count": len(confidences)
+        })
+        
     except Exception as e:
         traceback.print_exc()
         return jsonify({"error": f"Inference failed: {e}"}), 500
 
-    # overlay masks and extract confidences > threshold
-    annotated_pil, kept_conf = overlay_mask_on_image(img, detections, threshold=conf_threshold,
-                                                     mask_color=(255, 77, 166), alpha=0.45)
-
-    data_url = encode_pil_to_dataurl(annotated_pil, fmt="PNG")
-    return jsonify({"annotated": data_url, "confidences": kept_conf, "count": len(kept_conf)})
-
 
 if __name__ == "__main__":
-    # Warm model in a background thread to avoid blocking the container start logs too long
+    # Warm model in background thread
     def warm():
         try:
+            print("[INFO] Starting model warmup...")
             init_model()
+            print("[INFO] Model warmup complete")
         except Exception as e:
-            print("Model warmup failed:", e)
+            print(f"[ERROR] Model warmup failed: {e}")
+            traceback.print_exc()
+    
     threading.Thread(target=warm, daemon=True).start()
-    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), debug=False)
+    
+    # Run Flask app
+    app.run(
+        host="0.0.0.0",
+        port=int(os.environ.get("PORT", 7860)),
+        debug=False
+    )