Views
No views yet
tra-base-v1 is a lightweight, TinyML-focused model designed for vehicle counting and localization through density map estimation. This is a base version (v1) suitable for direct deployment or further fine-tuning on custom datasets.results/ folder:
pip install onnxruntime opencv-python-headless scipy matplotlib psutil --quiet1from huggingface_hub import snapshot_download
2import os
3
4repo_id = "realAABeigi/tra-base-1"
5
6print(f"[INFO] Downloading all files from {repo_id} to root...")
7
8try:
9 # This will download all files from the repo and place them in the current directory
10 snapshot_download(
11 repo_id=repo_id,
12 local_dir="./",
13 local_dir_use_symlinks=False
14 )
15 print("[SUCCESS] All files downloaded to root directory.")
16except Exception as e:
17 print(f"[ERROR] Failed to download: {e}")1!pip install onnxruntime opencv-python-headless scipy matplotlib psutil --quiet
2
3import os
4import cv2
5import numpy as np
6import matplotlib.pyplot as plt
7from scipy.ndimage import maximum_filter
8import gc
9import tracemalloc
10import time
11import psutil
12import onnxruntime as ort
13
14THRESHOLD = 0.4
15IMG_SIZE = 192
16GRID_SIZE = 48
17ONNX_PATH = "tra-base.onnx"
18DATA_PATH = "tra-base.onnx.data"
19TEST_DIR = "/content/Test"
20
21def get_process_memory():
22 process = psutil.Process(os.getpid())
23 return process.memory_info().rss
24
25try:
26 model_size = os.path.getsize(ONNX_PATH)
27 if os.path.exists(DATA_PATH):
28 model_size += os.path.getsize(DATA_PATH)
29
30 session = ort.InferenceSession(ONNX_PATH, providers=['CPUExecutionProvider'])
31 input_name = session.get_inputs()[0].name
32
33 def preprocess(img_path):
34 orig_img = cv2.imread(img_path)
35 if orig_img is None: return None, None
36 img_rgb = cv2.cvtColor(orig_img, cv2.COLOR_BGR2RGB)
37 img_resized = cv2.resize(img_rgb, (IMG_SIZE, IMG_SIZE))
38 img_data = img_resized.astype(np.float32) / 255.0
39 mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
40 std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
41 img_data = (img_data - mean) / std
42 img_data = np.transpose(img_data, (2, 0, 1))
43 img_data = np.expand_dims(img_data, axis=0)
44 return img_data, img_rgb
45
46 def run_onnx_cpu_inference(img_path):
47 gc.collect()
48 tracemalloc.start()
49
50 mem_before = get_process_memory()
51 start_time = time.time()
52
53 img_data, img_rgb = preprocess(img_path)
54 if img_data is None: return
55
56 outputs = session.run(None, {input_name: img_data})
57
58 mem_after = get_process_memory()
59 inference_time = (time.time() - start_time) * 1000
60
61 current, peak = tracemalloc.get_traced_memory()
62 tracemalloc.stop()
63
64 heatmap = outputs[0].squeeze()
65 data_max = maximum_filter(heatmap, size=3)
66 maxima = (heatmap == data_max) & (heatmap > THRESHOLD)
67 y_coords, x_coords = np.where(maxima)
68
69 system_delta = mem_after - mem_before
70 total_footprint_kb = (model_size + peak + abs(system_delta)) / 1024
71
72 print(f"\n--- Image: {os.path.basename(img_path)} ---")
73 print(f"Latency: {inference_time:.2f}ms")
74 print(f"System Memory Change: {system_delta/1024:.2f} KB")
75 print(f"Total RAM Footprint (Est): {total_footprint_kb:.2f} KB")
76
77 plt.figure(figsize=(10, 4))
78 plt.subplot(1, 2, 1)
79 display_img = cv2.resize(img_rgb, (384, 384))
80 for y, x in zip(y_coords, x_coords):
81 cx, cy = int(x * (384/GRID_SIZE)), int(y * (384/GRID_SIZE))
82 cv2.circle(display_img, (cx, cy), 6, (255, 0, 0), -1)
83 plt.imshow(display_img)
84 plt.title(f"Cars: {len(y_coords)} | Time: {inference_time:.1f}ms")
85 plt.axis('off')
86
87 plt.subplot(1, 2, 2)
88 plt.imshow(heatmap, cmap='jet')
89 plt.title(f"Total RAM: {total_footprint_kb:.1f} KB")
90 plt.axis('off')
91 plt.show()
92
93 if os.path.exists(TEST_DIR):
94 image_files = [f for f in os.listdir(TEST_DIR) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
95 for img_file in image_files:
96 run_onnx_cpu_inference(os.path.join(TEST_DIR, img_file))
97 else:
98 print("Folder Test not found.")
99
100except Exception as e:
101 print(f"[ERROR]: {e}")