Create meaw.py

meaw.py

@@ -0,0 +1,116 @@
+import os
+import cv2
+import torch
+import gdown
+import base64
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import JSONResponse
+from detectron2.config import get_cfg
+from detectron2 import model_zoo
+from detectron2.engine import DefaultPredictor
+from detectron2.utils.visualizer import Visualizer, ColorMode
+from detectron2.data import MetadataCatalog
+import io
+# === Model Setup ===
+MODEL_PATH = "/tmp/model_final.pth"
+DRIVE_FILE_ID = "1bazIVYG0CYMubDLoMu5pgH6ArC1sayzg"
+
+def download_model():
+    if not os.path.exists(MODEL_PATH):
+        url = f"https://drive.google.com/uc?id={DRIVE_FILE_ID}"
+
+        # Create a writable temporary directory for gdown
+        tmp_dir = "/tmp/gdown"
+        os.makedirs(tmp_dir, exist_ok=True)
+
+        # Set GDOWN cache path to a safe location
+        os.environ["GDOWN_CACHE_DIR"] = tmp_dir
+
+        print("Downloading Detectron2 model...")
+        gdown.download(
+            url,
+            MODEL_PATH,
+            quiet=False,
+            fuzzy=True,
+            use_cookies=False  # Important for Hugging Face
+        )
+        print("Download complete.")
+
+download_model()
+
+
+# Register dummy metadata (for visualization only)
+MetadataCatalog.get(metadata_name).set(thing_classes=class_names)
+
+# === Config Setup ===
+cfg = get_cfg()
+cfg.merge_from_file(
+    model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
+)
+cfg.MODEL.WEIGHTS = MODEL_PATH
+cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+cfg.MODEL.DEVICE = "cpu"  # or "cuda" if GPU is enabled
+
+# === Predictor ===
+predictor = DefaultPredictor(cfg)
+# -----------------------------
+# 3. Helper: Encode mask to Base64
+# -----------------------------
+def encode_mask(mask: np.ndarray) -> str:
+    """Convert mask numpy array to base64 PNG string."""
+    mask_img = Image.fromarray(mask.astype(np.uint8))
+    buf = io.BytesIO()
+    mask_img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode("utf-8")
+# === Predict Function ===
+def predict(image_path: str):
+    im = cv2.imread(image_path)
+    outputs = predictor(im)
+    instances = outputs["instances"].to("cpu")
+
+    # Visualization
+    v = Visualizer(
+        im[:, :, ::-1],
+        metadata=metadata,
+        scale=1.2,
+        instance_mode=ColorMode.IMAGE
+    )
+    out = v.draw_instance_predictions(instances)
+    result_img = out.get_image()[:, :, ::-1]
+
+ # -----------------------------
+# 4. API Endpoint
+# -----------------------------
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    # Read image
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents)).convert("RGB")
+    image = np.array(image)[:, :, ::-1]  # to BGR for OpenCV/Detectron2
+
+    # Run inference
+    outputs = predictor(image)
+    instances = outputs["instances"].to("cpu")
+
+    results = []
+    mask_b64 = None
+
+    if instances.has("pred_masks"):
+        masks = instances.pred_masks.numpy()
+        boxes = instances.pred_boxes.tensor.numpy()
+        scores = instances.scores.numpy()
+
+        # Combine masks into one
+        combined_mask = np.any(masks, axis=0).astype(np.uint8) * 255
+        mask_b64 = encode_mask(combined_mask)
+
+        for i in range(len(masks)):
+            results.append({
+                "box": boxes[i].tolist(),
+                "score": float(scores[i])
+            })
+
+    return JSONResponse({
+        "predictions": results,
+        "mask": mask_b64  # base64 string (PNG)
+    })