agrobot_base/Python/gal5000/dataset_caputre.py

423 lines
13 KiB
Python

import os
import time
import json
import math
import ctypes as C
from ctypes import wintypes as W
from datetime import datetime
import numpy as np
import cv2
# =========================
# CONFIG
# =========================
SDK_DIR = os.path.join(os.path.dirname(__file__), "dlls")
DLL_NAME = "VT_SDK64.dll"
# Onde salvar o dataset
OUT_ROOT = os.path.join(os.path.dirname(__file__), "dataset")
SESSION_DIR = os.path.join(OUT_ROOT, datetime.now().strftime("%Y%m%d"))
os.makedirs(SESSION_DIR, exist_ok=True)
# Camera scan/open
DEVICE_UDEF = 0
DEVICE_INDEX = 0
DATA_RAW = 0
# RAW geometry (se mudar no futuro, ajuste)
RAW_W = 2592
RAW_H = 2056
TIMEOUT_MS = 2000
WINDOW_NAME = "GAL5000 Dataset Capture (C/SPACE=save | A=auto-save | E=AE toggle | Q=quit)"
# Preview
UPSCALE = 2
# Auto-save
CAPTURE_INTERVAL_S = 1.0
# Param IDs (VT_Param.h)
BUF_SIZE = 256
PARAM_ID_SENSOR_EXPOSURETIMERAW = 0x00003010
PARAM_ID_SENSOR_GAINANALOGRAW = 0x00003020
PARAM_ID_SENSOR_GAINDIGITRAW = 0x0000302A
# PARAM_VALUETYPE
VALUE_INT = 0
VALUE_FLOAT = 1
VALUE_STR = 2
# Exposure/Gain limits (ajuste depois conforme o sensor aceitar)
EXP_MIN = 1
EXP_MAX = 20000
GAIN_A_MIN, GAIN_A_MAX = 0, 255
GAIN_D_MIN, GAIN_D_MAX = 0, 255
# =========================
# Helpers
# =========================
def ck(ret: int, name: str):
if ret != 0:
raise RuntimeError(f"{name} falhou, ret={ret}")
def ts_name() -> str:
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
def clamp(v, lo, hi):
return lo if v < lo else hi if v > hi else v
def norm8(x, p_lo=2, p_hi=98):
lo = np.percentile(x, p_lo)
hi = np.percentile(x, p_hi)
if hi <= lo + 1:
return x.astype(np.uint8)
y = (x.astype(np.float32) - lo) * (255.0 / (hi - lo))
return np.clip(y, 0, 255).astype(np.uint8)
def make_rgb_preview(raw: np.ndarray, upscale=2) -> np.ndarray:
# pattern:
# R G
# IR B
R = raw[0::2, 0::2]
G = raw[0::2, 1::2]
B = raw[1::2, 1::2]
Rn, Gn, Bn = norm8(R), norm8(G), norm8(B)
bgr = np.dstack([Bn, Gn, Rn]) # OpenCV usa BGR
if upscale and upscale != 1:
bgr = cv2.resize(bgr, (bgr.shape[1]*upscale, bgr.shape[0]*upscale), interpolation=cv2.INTER_NEAREST)
return bgr
def measure_raw_g_metrics(raw: np.ndarray):
"""
Mede brilho no canal G cru usando uma ROI na base (mais parecido com chão).
Retorna p90/p95 e fração saturada.
"""
G = raw[0::2, 1::2] # H/2 x W/2
h2, w2 = G.shape
# ROI: base da imagem, cortando laterais
y0, y1 = int(h2 * 0.55), int(h2 * 0.95)
x0, x1 = int(w2 * 0.15), int(w2 * 0.85)
roi = G[y0:y1, x0:x1]
p90 = float(np.percentile(roi, 90))
p95 = float(np.percentile(roi, 95))
sat = float(np.mean(roi >= 250))
return p90, p95, sat
class RobustAE:
"""
Controle soft de exposure (sem depender do GET da camera):
- mede p95 do canal G cru em ROI
- usa EMA + deadband (pra não ficar "descendo até 16" como você viu)
- passo multiplicativo em log, com limite de passo
"""
def __init__(self,
exp_min=EXP_MIN, exp_max=EXP_MAX,
target_p95=140.0,
deadband=6.0,
k=0.12,
max_step=0.10,
ema_alpha=0.20,
sat_limit=0.01):
self.exp_min = exp_min
self.exp_max = exp_max
self.target = target_p95
self.deadband = deadband
self.k = k
self.max_step = max_step
self.ema_alpha = ema_alpha
self.sat_limit = sat_limit
self.p95_ema = None
def step(self, raw, exp_raw):
p90, p95, sat = measure_raw_g_metrics(raw)
# EMA do p95 (estabiliza)
if self.p95_ema is None:
self.p95_ema = p95
else:
self.p95_ema = (1 - self.ema_alpha) * self.p95_ema + self.ema_alpha * p95
e = self.target - self.p95_ema # erro em nível de pixel
# deadband: segura a mão perto do alvo
if abs(e) <= self.deadband and sat <= self.sat_limit:
return exp_raw, {"p90": p90, "p95": p95, "p95_ema": self.p95_ema, "sat": sat, "hold": True}
# saturou: garante redução
if sat > self.sat_limit:
step = -min(self.max_step, 0.12)
else:
ratio = (self.target + 1e-6) / (self.p95_ema + 1e-6)
step = self.k * math.log(ratio)
step = max(-self.max_step, min(self.max_step, step))
new_exp = int(round(exp_raw * math.exp(step)))
new_exp = max(self.exp_min, min(self.exp_max, new_exp))
return new_exp, {"p90": p90, "p95": p95, "p95_ema": self.p95_ema, "sat": sat, "step": step, "hold": False}
# =========================
# STRUCTS + Param API
# =========================
class VT_FRAMEINFO(C.Structure):
_fields_ = [
("lFrameID", W.DWORD),
("lBufSize", W.DWORD),
("lWidth", W.DWORD),
("lHeight", W.DWORD),
("lPixBits", C.c_ubyte),
("_pad0", C.c_ubyte * 3),
("pBufPtr", C.POINTER(C.c_ubyte)),
("lFrameStatus", W.DWORD),
("lPixType", W.DWORD),
("lTimeStamp", W.DWORD),
("_reserve", W.DWORD * 8),
]
class VT_DEVPARAM(C.Structure):
_fields_ = [
("bUseName", W.BOOL),
("lParamByID", W.DWORD),
("lParamByName", C.c_char * BUF_SIZE),
]
def devparam_by_id(pid: int) -> VT_DEVPARAM:
p = VT_DEVPARAM()
p.bUseName = False
p.lParamByID = pid
p.lParamByName = b""
return p
# =========================
# DLL LOAD + prototypes
# =========================
os.add_dll_directory(SDK_DIR)
dll = C.WinDLL(os.path.join(SDK_DIR, DLL_NAME))
print("DLL carregada OK:", dll)
dll.VT_DeviceScan.argtypes = [C.POINTER(C.c_ubyte), C.c_int]
dll.VT_DeviceScan.restype = C.c_int
dll.VT_DeviceOpen.argtypes = [C.c_void_p, C.POINTER(W.HANDLE), C.c_int, C.c_int]
dll.VT_DeviceOpen.restype = C.c_int
dll.VT_SingleFrameCapture.argtypes = [W.HANDLE, C.POINTER(VT_FRAMEINFO), C.c_int, C.c_int, W.BOOL]
dll.VT_SingleFrameCapture.restype = C.c_int
dll.VT_DeviceClose.argtypes = [C.POINTER(W.HANDLE)]
dll.VT_DeviceClose.restype = C.c_int
dll.VT_ParamGetValue.argtypes = [W.HANDLE, VT_DEVPARAM, C.c_void_p, C.c_int]
dll.VT_ParamGetValue.restype = C.c_int
dll.VT_ParamSetValue.argtypes = [W.HANDLE, VT_DEVPARAM, C.c_void_p, C.c_int]
dll.VT_ParamSetValue.restype = C.c_int
def param_set_int(h: W.HANDLE, pid: int, value: int):
p = devparam_by_id(pid)
v = C.c_int(int(value))
ret = dll.VT_ParamSetValue(h, p, C.byref(v), VALUE_INT)
ck(ret, f"VT_ParamSetValue({hex(pid)})")
def capture_raw8(h: W.HANDLE) -> np.ndarray:
fi = VT_FRAMEINFO()
ret = dll.VT_SingleFrameCapture(h, C.byref(fi), DATA_RAW, TIMEOUT_MS, True)
ck(ret, "VT_SingleFrameCapture")
w, hh = int(fi.lWidth), int(fi.lHeight)
buf = C.string_at(fi.pBufPtr, fi.lBufSize)
arr = np.frombuffer(buf, dtype=np.uint8)
needed = w * hh
if arr.size < needed:
arr = np.pad(arr, (0, needed - arr.size), mode="constant", constant_values=0)
arr = arr[:needed].reshape(hh, w)
return arr
def overlay_hud(img_bgr, lines):
y = 28
for s in lines:
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 3, cv2.LINE_AA)
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255,255,255), 2, cv2.LINE_AA)
y += 28
def save_sample(raw: np.ndarray, bgr_preview: np.ndarray, meta: dict):
name = ts_name()
raw_path = os.path.join(SESSION_DIR, f"{name}.raw")
png_path = os.path.join(SESSION_DIR, f"{name}.png")
json_path = os.path.join(SESSION_DIR, f"{name}.json")
raw.tofile(raw_path)
cv2.imwrite(png_path, bgr_preview)
with open(json_path, "w", encoding="utf-8") as f:
json.dump(meta, f, ensure_ascii=False, indent=2)
return raw_path, png_path, json_path
def main():
# scan
n = C.c_ubyte(0)
ck(dll.VT_DeviceScan(C.byref(n), DEVICE_UDEF), "VT_DeviceScan")
if n.value == 0:
raise RuntimeError("Nenhuma câmera encontrada.")
# open
idx = C.c_ubyte(0)
h = W.HANDLE()
ck(dll.VT_DeviceOpen(C.byref(idx), C.byref(h), DEVICE_INDEX, DEVICE_UDEF), "VT_DeviceOpen")
print("DeviceOpen OK, handle=", h.value)
print("Saving to:", SESSION_DIR)
cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
# Estado local (não dependemos de GET)
exp_raw = 1500
gain_a = 0
gain_d = 0
# Aplica estado inicial
try:
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
param_set_int(h, PARAM_ID_SENSOR_GAINANALOGRAW, gain_a)
param_set_int(h, PARAM_ID_SENSOR_GAINDIGITRAW, gain_d)
except Exception as e:
print("[WARN] Falhou set inicial:", e)
ae = RobustAE(target_p95=140.0, deadband=6.0, k=0.12, max_step=0.10, ema_alpha=0.20, sat_limit=0.01)
ae_on = True
auto_save = False
last_auto_t = 0.0
# FPS
t0 = time.time()
frames = 0
fps = 0.0
last_msg = ""
last_msg_t = 0.0
try:
while True:
raw = capture_raw8(h)
# soft AE
ae_dbg = {}
if ae_on:
new_exp, ae_dbg = ae.step(raw, exp_raw)
if new_exp != exp_raw:
exp_raw = new_exp
try:
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
except Exception as e:
# se set falhar, desliga AE pra não ficar insistindo
print("[ERR] set exposure:", e)
ae_on = False
# preview RGB bonitão
rgb_clean = make_rgb_preview(raw, upscale=UPSCALE)
bgr = rgb_clean.copy()
# FPS
frames += 1
dt = time.time() - t0
if dt >= 1.0:
fps = frames / dt
frames = 0
t0 = time.time()
# HUD
lines = [
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Interval: {CAPTURE_INTERVAL_S:.1f}s",
f"exp_raw={exp_raw} gain_a={gain_a} gain_d={gain_d} | FPS={fps:.1f}",
f"AEdbg: p95={ae_dbg.get('p95_ema', ae_dbg.get('p95', 0)):.1f} sat={ae_dbg.get('sat', 0):.3f} hold={ae_dbg.get('hold', False)}",
"Keys: C/SPACE=save | A=toggle autosave | E=toggle AE | +/- exp | Q/ESC quit",
]
overlay_hud(bgr, lines)
# msg pós-save
if last_msg and (time.time() - last_msg_t) < 2.0:
cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18),
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2, cv2.LINE_AA)
cv2.imshow(WINDOW_NAME, bgr)
# autosave
now = time.time()
if auto_save and (now - last_auto_t) >= CAPTURE_INTERVAL_S:
meta = {
"ts": datetime.now().isoformat(timespec="milliseconds"),
"raw_w": RAW_W, "raw_h": RAW_H,
"exp_raw": int(exp_raw),
"gain_a": int(gain_a),
"gain_d": int(gain_d),
"ae_on": bool(ae_on),
"note": "autosave",
}
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta)
last_msg = f"SAVED: {os.path.basename(raw_path)}"
last_msg_t = now
last_auto_t = now
k = cv2.waitKey(1) & 0xFF
if k in (ord('q'), ord('Q'), 27):
break
elif k in (ord('a'), ord('A')):
auto_save = not auto_save
last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
last_msg_t = time.time()
elif k in (ord('e'), ord('E')):
ae_on = not ae_on
last_msg = f"AE -> {'ON' if ae_on else 'OFF'}"
last_msg_t = time.time()
elif k in (ord('c'), ord('C'), 32): # C ou SPACE
meta = {
"ts": datetime.now().isoformat(timespec="milliseconds"),
"raw_w": RAW_W, "raw_h": RAW_H,
"exp_raw": int(exp_raw),
"gain_a": int(gain_a),
"gain_d": int(gain_d),
"ae_on": bool(ae_on),
"note": "manual",
}
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta)
last_msg = f"SAVED: {os.path.basename(raw_path)}"
last_msg_t = time.time()
elif k in (ord('+'), ord('=')):
exp_raw = clamp(exp_raw + 200, EXP_MIN, EXP_MAX)
try:
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
except Exception as e:
print("[ERR] manual exp +:", e)
elif k in (ord('-'), ord('_')):
exp_raw = clamp(exp_raw - 200, EXP_MIN, EXP_MAX)
try:
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
except Exception as e:
print("[ERR] manual exp -:", e)
finally:
try:
ret = dll.VT_DeviceClose(C.byref(h))
if ret != 0:
print("VT_DeviceClose retornou:", ret)
except Exception as e:
print("Erro ao fechar:", e)
cv2.destroyAllWindows()
print("Fim.")
if __name__ == "__main__":
main()