Update app.py

app.py

@@ -1,205 +1,9 @@
-# import os
-# import io
-# import base64
-# import tempfile
-# import threading
-# from PIL import Image, ImageDraw, ImageFont
-# import numpy as np
-# from flask import Flask, request, jsonify, send_from_directory
-# import requests
-
-# # Force CPU-only (prevents accidental GPU usage); works by hiding CUDA devices
-# os.environ["CUDA_VISIBLE_DEVICES"] = ""
-
-# # --- model import (ensure rfdetr package is available in requirements) ---
-# try:
-#     from rfdetr import RFDETRSegPreview
-# except Exception as e:
-#     raise RuntimeError("rfdetr package import failed. Make sure `rfdetr` is in requirements.") from e
-
-# app = Flask(__name__, static_folder="static", static_url_path="/")
-
-# # HF checkpoint raw resolve URL (use the 'resolve/main' raw link)
-# CHECKPOINT_URL = "https://huggingface.co/Subh775/Segment-Tulsi-TFs-3/resolve/main/checkpoint_best_total.pth"
-# CHECKPOINT_PATH = os.path.join("/tmp", "checkpoint_best_total.pth")
-
-# MODEL_LOCK = threading.Lock()
-# MODEL = None
-
-# def download_file(url: str, dst: str):
-#     if os.path.exists(dst):
-#         return dst
-#     print(f"[INFO] Downloading weights from {url} ...")
-#     r = requests.get(url, stream=True, timeout=60)
-#     r.raise_for_status()
-#     with open(dst, "wb") as fh:
-#         for chunk in r.iter_content(chunk_size=8192):
-#             if chunk:
-#                 fh.write(chunk)
-#     print("[INFO] Download complete.")
-#     return dst
-
-# def init_model():
-#     global MODEL
-#     with MODEL_LOCK:
-#         if MODEL is None:
-#             # Ensure model checkpoint
-#             try:
-#                 download_file(CHECKPOINT_URL, CHECKPOINT_PATH)
-#             except Exception as e:
-#                 print(f"[WARN] Failed to download checkpoint: {e}. Attempting to init model without weights.")
-#                 # continue; model may fallback to default weights
-#             print("[INFO] Loading RF-DETR model (CPU mode)...")
-#             MODEL = RFDETRSegPreview(pretrain_weights=CHECKPOINT_PATH if os.path.exists(CHECKPOINT_PATH) else None)
-#             try:
-#                 MODEL.optimize_for_inference()
-#             except Exception:
-#                 # optimization may fail on CPU or if not implemented; ignore
-#                 pass
-#             print("[INFO] Model ready.")
-#     return MODEL
-
-# @app.route("/")
-# def index():
-#     return send_from_directory("static", "index.html")
-
-# def decode_data_url(data_url: str) -> Image.Image:
-#     if data_url.startswith("data:"):
-#         header, b64 = data_url.split(",", 1)
-#         data = base64.b64decode(b64)
-#         return Image.open(io.BytesIO(data)).convert("RGB")
-#     else:
-#         # assume plain base64 or path
-#         data = base64.b64decode(data_url)
-#         return Image.open(io.BytesIO(data)).convert("RGB")
-
-# def encode_pil_to_dataurl(pil_img: Image.Image, fmt="PNG"):
-#     buf = io.BytesIO()
-#     pil_img.save(buf, format=fmt)
-#     b = base64.b64encode(buf.getvalue()).decode("ascii")
-#     return f"data:image/{fmt.lower()};base64,{b}"
-
-# def overlay_mask_on_image(pil_img: Image.Image, masks, confidences, threshold=0.01, mask_color=(255,77,166), alpha=0.45):
-#     """
-#     masks: either list of HxW bool arrays or numpy array (N,H,W)
-#     confidences: list of floats
-#     Returns annotated PIL image and list of kept confidences and count.
-#     """
-#     base = pil_img.convert("RGBA")
-#     W, H = base.size
-
-#     # Normalize masks to N,H,W
-#     if masks is None:
-#         return base, []
-
-#     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)
-#         # masks might be (H,W,N) -> transpose
-#         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)
-
-#     # create overlay
-#     overlay = Image.new("RGBA", (W, H), (0,0,0,0))
-#     draw = ImageDraw.Draw(overlay)
-
-#     kept_confidences = []
-#     for i in range(masks_arr.shape[0]):
-#         conf = float(confidences[i]) if confidences is not None and 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").resize((W, H), resample=Image.NEAREST)
-#         # create colored mask image
-#         color_layer = Image.new("RGBA", (W,H), mask_color + (0,))
-#         # put alpha using mask
-#         color_layer.putalpha(mask_img.point(lambda p: int(p * alpha)))
-#         overlay = Image.alpha_composite(overlay, color_layer)
-#         kept_confidences.append(conf)
-
-#     # composite
-#     annotated = Image.alpha_composite(base, overlay)
-
-#     # add confidence text (show highest kept confidence)
-#     if len(kept_confidences) > 0:
-#         best = max(kept_confidences)
-#         draw = ImageDraw.Draw(annotated)
-#         try:
-#             # Try to use a builtin font
-#             font = ImageFont.truetype("DejaVuSans-Bold.ttf", size=max(16, W//30))
-#         except Exception:
-#             font = ImageFont.load_default()
-#         text = f"Confidence: {best:.2f}"
-#         # draw background box for text
-#         tw, th = draw.textsize(text, font=font)
-#         pad = 8
-#         draw.rectangle([6,6, 6+tw+pad, 6+th+pad], fill=(0,0,0,180))
-#         draw.text((6+4,6+2), text, font=font, fill=(255,255,255,255))
-#     return annotated.convert("RGB"), kept_confidences
-
-# @app.route("/predict", methods=["POST"])
-# def predict():
-#     payload = request.get_json(force=True)
-#     if not payload or "image" not in payload:
-#         return jsonify({"error": "Missing image"}), 400
-#     conf = float(payload.get("conf", 0.25))
-
-#     # ensure model ready
-#     model = init_model()
-
-#     # decode image
-#     try:
-#         pil = decode_data_url(payload["image"])
-#     except Exception as e:
-#         return jsonify({"error": f"Invalid image: {e}"}), 400
-
-#     # perform prediction (model.predict expects PIL image)
-#     try:
-#         detections = model.predict(pil, threshold=0.0)  # we filter using conf manually
-#     except Exception as e:
-#         return jsonify({"error": f"Inference failure: {e}"}), 500
-
-#     # extract masks and confidences
-#     masks = getattr(detections, "masks", None)
-#     confidences = []
-#     # attempt to read per-instance confidence
-#     try:
-#         confidences = [float(x) for x in getattr(detections, "confidence", [])]
-#     except Exception:
-#         # fallback: attempt attribute 'scores' or 'scores_' or generate ones
-#         confidences = []
-#         try:
-#             confidences = [float(x) for x in getattr(detections, "scores", [])]
-#         except Exception:
-#             confidences = [1.0] * (masks.shape[0] if masks is not None and hasattr(masks, "shape") and masks.shape[0] else 0)
-
-#     # overlay mask with pink-red color
-#     mask_color = (255, 77, 166)  # pinkish
-#     annotated_pil, kept_conf = overlay_mask_on_image(pil, masks, confidences, threshold=conf, mask_color=mask_color, 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 up model on startup (non-blocking)
-#     try:
-#         init_model()
-#     except Exception as e:
-#         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 traceback
+import gc
 from typing import Optional
 
 from flask import Flask, request, jsonify, send_from_directory
@@ -211,13 +15,21 @@ import torch
 # Set environment variables for CPU-only operation
 os.environ.setdefault("MPLCONFIGDIR", "/tmp/.matplotlib")
 os.environ.setdefault("FONTCONFIG_PATH", "/tmp/.fontconfig")
+os.environ.setdefault("FONTCONFIG_FILE", "/etc/fonts/fonts.conf")
 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")
 
+# Create writable fontconfig cache
+os.makedirs("/tmp/.fontconfig", exist_ok=True)
+os.makedirs("/tmp/.matplotlib", exist_ok=True)
+
 # Limit torch threads
-torch.set_num_threads(4)
+try:
+    torch.set_num_threads(4)
+except Exception:
+    pass
 
 import supervision as sv
 from rfdetr import RFDETRSegPreview
@@ -238,14 +50,18 @@ def download_file(url: str, dst: str, chunk_size: int = 8192):
         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=120)
-    r.raise_for_status()
-    with open(dst, "wb") as fh:
-        for chunk in r.iter_content(chunk_size=chunk_size):
-            if chunk:
-                fh.write(chunk)
-    print("[INFO] Download complete.")
-    return dst
+    try:
+        r = requests.get(url, stream=True, timeout=180)
+        r.raise_for_status()
+        with open(dst, "wb") as fh:
+            for chunk in r.iter_content(chunk_size=chunk_size):
+                if chunk:
+                    fh.write(chunk)
+        print("[INFO] Download complete.")
+        return dst
+    except Exception as e:
+        print(f"[ERROR] Download failed: {e}")
+        raise
 
 
 def init_model():
@@ -253,28 +69,31 @@ def init_model():
     global MODEL
     with MODEL_LOCK:
         if MODEL is not None:
+            print("[INFO] Model already loaded, returning cached instance")
             return MODEL
         try:
             # Ensure checkpoint present
-            try:
+            if not os.path.exists(CHECKPOINT_PATH):
+                print("[INFO] Checkpoint not found, downloading...")
                 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
+            else:
+                print(f"[INFO] Using existing checkpoint at {CHECKPOINT_PATH}")
 
             print("[INFO] Loading RF-DETR model (CPU mode)...")
             MODEL = RFDETRSegPreview(pretrain_weights=CHECKPOINT_PATH)
             
             # Try to optimize for inference
             try:
+                print("[INFO] Optimizing model for inference...")
                 MODEL.optimize_for_inference()
+                print("[INFO] Model optimization complete")
             except Exception as e:
-                print("[WARN] optimize_for_inference() skipped/failed:", e)
+                print(f"[WARN] optimize_for_inference() skipped/failed: {e}")
             
-            print("[INFO] Model ready.")
+            print("[INFO] Model ready for inference")
             return MODEL
-        except Exception:
+        except Exception as e:
+            print(f"[ERROR] Model initialization failed: {e}")
             traceback.print_exc()
             raise
 
@@ -295,7 +114,8 @@ 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)
+    pil_img.save(buf, format=fmt, optimize=False)
+    buf.seek(0)
     return "data:image/{};base64,".format(fmt.lower()) + base64.b64encode(buf.getvalue()).decode("ascii")
 
 
@@ -304,38 +124,53 @@ def annotate_segmentation(image: Image.Image, detections: sv.Detections) -> Imag
     Annotate image with segmentation masks using supervision library.
     This matches the visualization from rfdetr_seg_infer.py script.
     """
-    # 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
+    try:
+        # 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)
+        
+        print(f"[INFO] Creating annotators with text_scale={text_scale}")
+        
+        # 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 confidence scores
+        labels = [
+            f"Tulsi {float(conf):.2f}"
+            for conf in detections.confidence
+        ]
+        
+        print(f"[INFO] Annotating {len(labels)} detections")
+
+        # Apply annotations step by step
+        out = image.copy()
+        print("[INFO] Applying mask annotation...")
+        out = mask_annotator.annotate(out, detections)
+        print("[INFO] Applying polygon annotation...")
+        out = polygon_annotator.annotate(out, detections)
+        print("[INFO] Applying label annotation...")
+        out = label_annotator.annotate(out, detections, labels)
+        
+        print("[INFO] Annotation complete")
+        return out
+        
+    except Exception as e:
+        print(f"[ERROR] Annotation failed: {e}")
+        traceback.print_exc()
+        # Return original image if annotation fails
+        return image
 
 
 @app.route("/", methods=["GET"])
@@ -344,7 +179,18 @@ def index():
     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")
-    return jsonify({"message": "RF-DETR Segmentation API is running."})
+    return jsonify({"message": "RF-DETR Segmentation API is running.", "status": "ready"})
+
+
+@app.route("/health", methods=["GET"])
+def health():
+    """Health check endpoint"""
+    model_loaded = MODEL is not None
+    return jsonify({
+        "status": "healthy",
+        "model_loaded": model_loaded,
+        "checkpoint_exists": os.path.exists(CHECKPOINT_PATH)
+    })
 
 
 @app.route("/predict", methods=["POST"])
@@ -356,10 +202,16 @@ def predict():
     Returns JSON:
       {"annotated": "<data:image/png;base64,...>", "confidences": [..], "count": N}
     """
+    print("\n[INFO] ========== New prediction request ==========")
+    
     try:
+        print("[INFO] Initializing model...")
         model = init_model()
+        print("[INFO] Model ready")
     except Exception as e:
-        return jsonify({"error": f"Model initialization failed: {e}"}), 500
+        error_msg = f"Model initialization failed: {e}"
+        print(f"[ERROR] {error_msg}")
+        return jsonify({"error": error_msg}), 500
 
     # Parse input
     img: Optional[Image.Image] = None
@@ -368,10 +220,13 @@ def predict():
     # Check if file uploaded
     if "file" in request.files:
         file = request.files["file"]
+        print(f"[INFO] Processing uploaded file: {file.filename}")
         try:
             img = Image.open(file.stream).convert("RGB")
         except Exception as e:
-            return jsonify({"error": f"Invalid uploaded image: {e}"}), 400
+            error_msg = f"Invalid uploaded image: {e}"
+            print(f"[ERROR] {error_msg}")
+            return jsonify({"error": error_msg}), 400
         conf_threshold = float(request.form.get("conf", conf_threshold))
     else:
         # Try JSON payload
@@ -379,31 +234,37 @@ def predict():
         if not payload or "image" not in payload:
             return jsonify({"error": "No image provided. Upload 'file' or JSON with 'image' data-url."}), 400
         try:
+            print("[INFO] Decoding image from data URL...")
             img = decode_data_url(payload["image"])
         except Exception as e:
-            return jsonify({"error": f"Invalid image data: {e}"}), 400
+            error_msg = f"Invalid image data: {e}"
+            print(f"[ERROR] {error_msg}")
+            return jsonify({"error": error_msg}), 400
         conf_threshold = float(payload.get("conf", conf_threshold))
 
+    print(f"[INFO] Image size: {img.size}, Confidence threshold: {conf_threshold}")
+
     # 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))
+        print(f"[INFO] Resizing image from {w}x{h} to {new_w}x{new_h}")
         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:
+        print("[INFO] Running inference...")
         with torch.no_grad():
             detections = model.predict(img, threshold=conf_threshold)
         
-        print(f"[INFO] Detected {len(detections)} objects")
+        print(f"[INFO] Raw detections: {len(detections)} objects")
         
         # Check if detections exist
-        if len(detections) == 0:
+        if len(detections) == 0 or not hasattr(detections, 'confidence') or len(detections.confidence) == 0:
             print("[INFO] No detections above threshold")
-            # Return original image with message
+            # Return original image
             data_url = encode_pil_to_dataurl(img, fmt="PNG")
             return jsonify({
                 "annotated": data_url,
@@ -411,15 +272,33 @@ def predict():
                 "count": 0
             })
         
+        print(f"[INFO] Detections have {len(detections.confidence)} confidence scores")
+        print(f"[INFO] Confidence range: {min(detections.confidence):.3f} - {max(detections.confidence):.3f}")
+        
+        # Check if masks exist
+        if hasattr(detections, 'masks') and detections.masks is not None:
+            print(f"[INFO] Masks present: shape={np.array(detections.masks).shape if hasattr(detections.masks, '__len__') else 'unknown'}")
+        else:
+            print("[WARN] No masks found in detections!")
+        
         # Annotate image using supervision library
+        print("[INFO] Starting annotation...")
         annotated_pil = annotate_segmentation(img, detections)
         
         # Extract confidence scores
         confidences = [float(conf) for conf in detections.confidence]
+        print(f"[INFO] Final confidences: {confidences}")
         
         # Encode to data URL
+        print("[INFO] Encoding annotated image...")
         data_url = encode_pil_to_dataurl(annotated_pil, fmt="PNG")
         
+        # Clean up
+        del detections
+        gc.collect()
+        
+        print(f"[INFO] ========== Prediction complete: {len(confidences)} leaves detected ==========\n")
+        
         return jsonify({
             "annotated": data_url,
             "confidences": confidences,
@@ -427,19 +306,25 @@ def predict():
         })
         
     except Exception as e:
+        error_msg = f"Inference failed: {e}"
+        print(f"[ERROR] {error_msg}")
         traceback.print_exc()
-        return jsonify({"error": f"Inference failed: {e}"}), 500
+        return jsonify({"error": error_msg}), 500
 
 
 if __name__ == "__main__":
+    print("\n" + "="*60)
+    print("Starting Tulsi Leaf Segmentation Server")
+    print("="*60 + "\n")
+    
     # Warm model in background thread
     def warm():
         try:
-            print("[INFO] Starting model warmup...")
+            print("[INFO] Starting model warmup in background...")
             init_model()
-            print("[INFO] Model warmup complete")
+            print("[INFO] ✓ Model warmup complete - ready for predictions")
         except Exception as e:
-            print(f"[ERROR] Model warmup failed: {e}")
+            print(f"[ERROR] ✗ Model warmup failed: {e}")
             traceback.print_exc()
     
     threading.Thread(target=warm, daemon=True).start()