{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70cab2b9",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "from huggingface_hub import snapshot_download\n",
    "import os\n",
    "\n",
    "repo_id = \"realAABeigi/tra-base-1\"\n",
    "\n",
    "print(f\"[INFO] Downloading all files from {repo_id} to root...\")\n",
    "\n",
    "try:\n",
    "    # This will download all files from the repo and place them in the current directory\n",
    "    snapshot_download(\n",
    "        repo_id=repo_id,\n",
    "        local_dir=\"./\",\n",
    "        local_dir_use_symlinks=False\n",
    "    )\n",
    "    print(\"[SUCCESS] All files downloaded to root directory.\")\n",
    "except Exception as e:\n",
    "    print(f\"[ERROR] Failed to download: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aeb2c2a",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "!pip install onnxruntime opencv-python-headless scipy matplotlib psutil --quiet\n",
    "\n",
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.ndimage import maximum_filter\n",
    "import gc\n",
    "import tracemalloc\n",
    "import time\n",
    "import psutil\n",
    "import onnxruntime as ort\n",
    "\n",
    "THRESHOLD = 0.4\n",
    "IMG_SIZE = 192\n",
    "GRID_SIZE = 48\n",
    "ONNX_PATH = \"tiny_heatmap_car.onnx\"\n",
    "DATA_PATH = \"tiny_heatmap_car.onnx.data\"\n",
    "TEST_DIR = \"/content/Test\"\n",
    "\n",
    "def get_process_memory():\n",
    "    process = psutil.Process(os.getpid())\n",
    "    return process.memory_info().rss\n",
    "\n",
    "try:\n",
    "    model_size = os.path.getsize(ONNX_PATH)\n",
    "    if os.path.exists(DATA_PATH):\n",
    "        model_size += os.path.getsize(DATA_PATH)\n",
    "\n",
    "    session = ort.InferenceSession(ONNX_PATH, providers=['CPUExecutionProvider'])\n",
    "    input_name = session.get_inputs()[0].name\n",
    "\n",
    "    def preprocess(img_path):\n",
    "        orig_img = cv2.imread(img_path)\n",
    "        if orig_img is None: return None, None\n",
    "        img_rgb = cv2.cvtColor(orig_img, cv2.COLOR_BGR2RGB)\n",
    "        img_resized = cv2.resize(img_rgb, (IMG_SIZE, IMG_SIZE))\n",
    "        img_data = img_resized.astype(np.float32) / 255.0\n",
    "        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)\n",
    "        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)\n",
    "        img_data = (img_data - mean) / std\n",
    "        img_data = np.transpose(img_data, (2, 0, 1))\n",
    "        img_data = np.expand_dims(img_data, axis=0)\n",
    "        return img_data, img_rgb\n",
    "\n",
    "    def run_onnx_cpu_inference(img_path):\n",
    "        gc.collect()\n",
    "        tracemalloc.start()\n",
    "\n",
    "        mem_before = get_process_memory()\n",
    "        start_time = time.time()\n",
    "\n",
    "        img_data, img_rgb = preprocess(img_path)\n",
    "        if img_data is None: return\n",
    "\n",
    "        outputs = session.run(None, {input_name: img_data})\n",
    "\n",
    "        mem_after = get_process_memory()\n",
    "        inference_time = (time.time() - start_time) * 1000\n",
    "\n",
    "        current, peak = tracemalloc.get_traced_memory()\n",
    "        tracemalloc.stop()\n",
    "\n",
    "        heatmap = outputs[0].squeeze()\n",
    "        data_max = maximum_filter(heatmap, size=3)\n",
    "        maxima = (heatmap == data_max) & (heatmap > THRESHOLD)\n",
    "        y_coords, x_coords = np.where(maxima)\n",
    "\n",
    "        system_delta = mem_after - mem_before\n",
    "        total_footprint_kb = (model_size + peak + abs(system_delta)) / 1024\n",
    "\n",
    "        print(f\"\\n--- Image: {os.path.basename(img_path)} ---\")\n",
    "        print(f\"Latency: {inference_time:.2f}ms\")\n",
    "        print(f\"System Memory Change: {system_delta/1024:.2f} KB\")\n",
    "        print(f\"Total RAM Footprint (Est): {total_footprint_kb:.2f} KB\")\n",
    "\n",
    "        plt.figure(figsize=(10, 4))\n",
    "        plt.subplot(1, 2, 1)\n",
    "        display_img = cv2.resize(img_rgb, (384, 384))\n",
    "        for y, x in zip(y_coords, x_coords):\n",
    "            cx, cy = int(x * (384/GRID_SIZE)), int(y * (384/GRID_SIZE))\n",
    "            cv2.circle(display_img, (cx, cy), 6, (255, 0, 0), -1)\n",
    "        plt.imshow(display_img)\n",
    "        plt.title(f\"Cars: {len(y_coords)} | Time: {inference_time:.1f}ms\")\n",
    "        plt.axis('off')\n",
    "\n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.imshow(heatmap, cmap='jet')\n",
    "        plt.title(f\"Total RAM: {total_footprint_kb:.1f} KB\")\n",
    "        plt.axis('off')\n",
    "        plt.show()\n",
    "\n",
    "    if os.path.exists(TEST_DIR):\n",
    "        image_files = [f for f in os.listdir(TEST_DIR) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]\n",
    "        for img_file in image_files:\n",
    "            run_onnx_cpu_inference(os.path.join(TEST_DIR, img_file))\n",
    "    else:\n",
    "        print(\"Folder Test not found.\")\n",
    "\n",
    "except Exception as e:\n",
    "    print(f\"[ERROR]: {e}\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}