Upload folder using huggingface_hub

run.py

@@ -1,65 +1,157 @@
 #!/usr/bin/env python3
 """
-Single-file script (c.py): make prediction 0 or 1 by loading:
-
-  - input_image       path to image (argv[1] or test_image.jpg)
-  - best.pt           YOLO weights
-  - branch_cache.npz  precomputed branches (loaded for reference; prediction uses ViT+YOLO on input_image)
-  - ViT (yit)         natix-network-org/roadwork
-  - efficientnetv2_branch.keras  (via keras_dir)
-  - fusion_head.keras             (via keras_dir)
-  - final_model       final_output.keras (built from the two .keras files if missing)
+Self-contained prediction script for roadwork-miner (no imports from inception_model or parent repo).
+Uses: image (generated in code or from path), best.pt (YOLO), ViT (HuggingFace), final_output.keras.
+Returns 0 or 1 (no roadwork / roadwork).
 
 Usage:
-  cd /root/Workspace/aaa && poetry run python hf_refo/c.py [path/to/input_image.jpg]
+  python run.py image.jpg    # require image path; no default image
 """
 import sys
 from pathlib import Path
 
-REPO_ROOT = Path(__file__).resolve().parents[1]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
+DIR = Path(__file__).resolve().parent
+IMG_SIZE = 224
 
-# Paths: edit these or set via env
-INPUT_IMAGE = REPO_ROOT / "test_image.jpg"  # or pass as argv[1]
-BEST_PT = REPO_ROOT / "best.pt"
-BRANCH_CACHE = REPO_ROOT / "branch_cache.npz"
-VIT_REPO = "natix-network-org/roadwork"  # yit / ViT
-KERAS_DIR = REPO_ROOT / "inception_fusion_keras"  # efficientnetv2_branch.keras + fusion_head.keras
-FINAL_MODEL = Path(__file__).resolve().parent / "final_output.keras"
+# Paths (all under this folder)
+BEST_PT = DIR / "best.pt"
+FINAL_MODEL = DIR / "final_output.keras"
+VIT_REPO = "natix-network-org/roadwork"
 
 
-def main():
-    from PIL import Image
+def _image_to_keras_input(image):
+    """PIL or numpy (224,224,3) -> (1, 224, 224, 3) normalized for EfficientNet."""
     import numpy as np
+    from PIL import Image
 
-    # 1) input_image
-    image_path = Path(sys.argv[1]) if len(sys.argv) > 1 else INPUT_IMAGE
-    if not image_path.exists():
-        print(f"Image not found: {image_path}. Create it or run: python hf_refo/c.py /path/to/image.jpg")
-        sys.exit(1)
-    input_image = Image.open(image_path).convert("RGB")
-
-    # 2) branch_cache.npz
-    branch_cache = None
-    if BRANCH_CACHE.exists():
-        branch_cache = np.load(BRANCH_CACHE)
-        n = len(branch_cache.get("p_vit", branch_cache.get("labels", [])))
-        print(f"Loaded branch_cache.npz ({n} samples).")
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(np.asarray(image).astype(np.uint8))
+    if image.size != (IMG_SIZE, IMG_SIZE):
+        image = image.resize((IMG_SIZE, IMG_SIZE))
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    arr = np.array(image, dtype=np.float32) / 255.0
+    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+    arr = (arr - mean) / std
+    return np.expand_dims(arr, axis=0)
+
+
+def _get_vit_prob(pipe, image_pil):
+    out = pipe(image_pil)
+    if not isinstance(out, list):
+        out = [out]
+    for item in out:
+        if item.get("label") == "Roadwork":
+            return item["score"]
+    return 0.0
+
+
+def _get_yolo_prob(yolo_model, image_pil, roadwork_idx=1):
+    import numpy as np
+    r = yolo_model.predict(source=image_pil, verbose=False, device="cpu")
+    if not r or not hasattr(r[0], "probs") or r[0].probs is None:
+        return 0.0
+    p = r[0].probs.data
+    if hasattr(p, "cpu"):
+        p = p.cpu().numpy()
     else:
-        print(f"branch_cache.npz not found at {BRANCH_CACHE} (optional).")
-
-    # 3) best.pt, ViT (yit), efficientnetv2_branch.keras, fusion_head.keras, final_model -> prediction
-    from hf_refo.inception_model import predict_from_assets
-
-    label = predict_from_assets(
-        input_image,
-        best_pt=BEST_PT,
-        vit_repo=VIT_REPO,
-        keras_dir=KERAS_DIR,
-        final_model_path=FINAL_MODEL,
-        threshold=0.5,
+        p = np.asarray(p)
+    if p.ndim > 1:
+        p = p.ravel()
+    idx = min(roadwork_idx, len(p) - 1)
+    return float(p[idx])
+
+
+def load_pipeline():
+    """Load ViT, YOLO, and final_output.keras from this folder. Returns dict with vit, yolo, model."""
+    from tensorflow import keras
+    from transformers import AutoImageProcessor, AutoModelForImageClassification, pipeline
+    from ultralytics import YOLO
+
+    if not FINAL_MODEL.exists():
+        raise FileNotFoundError(f"final_output.keras not found at {FINAL_MODEL}")
+    if not BEST_PT.exists():
+        raise FileNotFoundError(f"best.pt not found at {BEST_PT}")
+
+    model = keras.models.load_model(FINAL_MODEL)
+    pipe = pipeline(
+        "image-classification",
+        model=AutoModelForImageClassification.from_pretrained(VIT_REPO),
+        feature_extractor=AutoImageProcessor.from_pretrained(VIT_REPO, use_fast=True),
+        device=-1,
     )
+    yolo = YOLO(str(BEST_PT))
+    return {"vit": pipe, "yolo": yolo, "model": model}
+
+
+def predict(image, pipeline, threshold=0.5):
+    """
+    Predict 0 or 1 from one image (PIL or numpy 224x224x3).
+    pipeline: from load_pipeline().
+    """
+    import numpy as np
+    from PIL import Image
+
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image.astype(np.uint8) if image.ndim == 3 else image[0].astype(np.uint8))
+    if image.size != (IMG_SIZE, IMG_SIZE):
+        image = image.resize((IMG_SIZE, IMG_SIZE))
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    p_vit = _get_vit_prob(pipeline["vit"], image)
+    p_yolo = _get_yolo_prob(pipeline["yolo"], image)
+
+    X_img = _image_to_keras_input(image)
+    p_vit_arr = np.array([[float(p_vit)]], dtype=np.float32)
+    p_yolo_arr = np.array([[float(p_yolo)]], dtype=np.float32)
+
+    prob = pipeline["model"].predict([X_img, p_vit_arr, p_yolo_arr], verbose=0)
+    roadwork_prob = float(prob[0, 0])
+    return 1 if roadwork_prob >= threshold else 0
+
+
+def make_demo_image(size=IMG_SIZE):
+    """Create a 224x224 RGB image in code (no file). Simple gradient for demo."""
+    import numpy as np
+    from PIL import Image
+
+    y = np.linspace(0, 1, size).reshape(size, 1)
+    x = np.linspace(0, 1, size).reshape(1, size)
+    r = (0.4 + 0.2 * x).clip(0, 1)   # (1, size) -> broadcast
+    g = (0.5 + 0.2 * y).clip(0, 1)   # (size, 1) -> broadcast
+    b = (0.45 + 0.1 * (x + y)).clip(0, 1)  # (size, size)
+    r = np.broadcast_to(r, (size, size))
+    g = np.broadcast_to(g, (size, size))
+    arr = np.stack([r, g, b], axis=-1)
+    arr = (arr * 255).astype(np.uint8)
+    return Image.fromarray(arr, mode="RGB")
+
+
+def load_image(path):
+    """Load an image from file. Returns PIL Image (RGB)."""
+    from PIL import Image
+
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"Image not found: {path}")
+    return Image.open(path).convert("RGB")
+
+
+def main():
+    if len(sys.argv) < 2:
+        print("Usage: python run.py <path_to_image>")
+        sys.exit(1)
+
+    try:
+        input_image = load_image(sys.argv[1])
+    except FileNotFoundError as e:
+        print(e)
+        sys.exit(1)
+
+    pipeline = load_pipeline()
+    label = predict(input_image, pipeline, threshold=0.5)
     print(f"Prediction: {label} (0=no roadwork, 1=roadwork)")
     return label