379 lines
13 KiB
Python
379 lines
13 KiB
Python
|
# infer.py (batendo com o treino: RAW->uint8 HWC, LetterBox, forward direto, NMS, masks, unletterbox)
|
||
|
|
import cv2
|
||
|
|
import torch
|
||
|
|
import numpy as np
|
||
|
|
from pathlib import Path
|
||
|
|
import yaml
|
||
|
|
|
||
|
|
from ultralytics import YOLO
|
||
|
|
from ultralytics.utils import ops
|
||
|
|
|
||
|
|
# ============================
|
||
|
|
# CONFIG
|
||
|
|
# ============================
|
||
|
|
with Path("data.yaml").open("r", encoding="utf-8") as f:
|
||
|
|
cfg = yaml.safe_load(f) or {}
|
||
|
|
|
||
|
|
CHANNELS = int(cfg.get("channels", 4)) # 4 ou 5
|
||
|
|
SCALE = cfg.get("scale", "m")
|
||
|
|
IMG_SZ = int(cfg.get("size", 800))
|
||
|
|
|
||
|
|
names_dict = cfg.get("names", {})
|
||
|
|
colors_dict = cfg.get("colors", {})
|
||
|
|
|
||
|
|
CLASS_NAMES = [names_dict[k] for k in sorted(names_dict.keys())] if names_dict else []
|
||
|
|
CLASS_COLORS = [
|
||
|
|
tuple(map(int, colors_dict[k].strip("()").split(",")))
|
||
|
|
for k in sorted(colors_dict.keys())
|
||
|
|
] if colors_dict else []
|
||
|
|
|
||
|
|
# Ajuste conforme seu run
|
||
|
|
WEIGHTS_PATH = f"C:/ZendionInc/agrobot_base/Python/yolov8-seg/runs/segment/ch_{CHANNELS}/sc_{SCALE}/sz_{IMG_SZ}/weed_detector_segformer/weights/best.pt"
|
||
|
|
TEST_DIR = Path(r"C:\ZendionInc\agrobot_base\Python\yolov8-seg\images\val")
|
||
|
|
|
||
|
|
# ============================
|
||
|
|
# RAW settings (iguais do treino)
|
||
|
|
# ============================
|
||
|
|
RAW_DTYPE = np.float32
|
||
|
|
RAW_ALREADY_01 = True
|
||
|
|
RAW_MAX_VAL = 1.0 if RAW_ALREADY_01 else 65535.0 # se algum dia usar uint16
|
||
|
|
|
||
|
|
# ============================
|
||
|
|
# Letterbox (multi-canal)
|
||
|
|
# ============================
|
||
|
|
def letterbox_multi(
|
||
|
|
img: np.ndarray,
|
||
|
|
new_shape: int = 640,
|
||
|
|
color: int = 114
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
img: HWC uint8 (H,W,C)
|
||
|
|
Retorna:
|
||
|
|
img_lb: letterboxed HWC uint8 (new_shape, new_shape, C)
|
||
|
|
ratio: (r, r)
|
||
|
|
pad: (dw, dh) em pixels (float)
|
||
|
|
new_unpad: (w_new, h_new)
|
||
|
|
"""
|
||
|
|
assert img.ndim == 3, "img deve ser HWC"
|
||
|
|
h0, w0 = img.shape[:2]
|
||
|
|
|
||
|
|
# scale ratio (new / old)
|
||
|
|
r = min(new_shape / h0, new_shape / w0)
|
||
|
|
w_new = int(round(w0 * r))
|
||
|
|
h_new = int(round(h0 * r))
|
||
|
|
|
||
|
|
# compute padding
|
||
|
|
dw = (new_shape - w_new) / 2
|
||
|
|
dh = (new_shape - h_new) / 2
|
||
|
|
|
||
|
|
# resize
|
||
|
|
if (w0, h0) != (w_new, h_new):
|
||
|
|
img = cv2.resize(img, (w_new, h_new), interpolation=cv2.INTER_LINEAR)
|
||
|
|
|
||
|
|
# pad (top, bottom, left, right)
|
||
|
|
top = int(round(dh - 0.1))
|
||
|
|
bottom = int(round(dh + 0.1))
|
||
|
|
left = int(round(dw - 0.1))
|
||
|
|
right = int(round(dw + 0.1))
|
||
|
|
|
||
|
|
# padding constante para qualquer número de canais
|
||
|
|
img_lb = np.pad(
|
||
|
|
img,
|
||
|
|
pad_width=((top, bottom), (left, right), (0, 0)),
|
||
|
|
mode="constant",
|
||
|
|
constant_values=color
|
||
|
|
).astype(img.dtype, copy=False)
|
||
|
|
|
||
|
|
ratio = (r, r)
|
||
|
|
pad = (dw, dh)
|
||
|
|
return img_lb, ratio, pad, (w_new, h_new), (top, bottom, left, right)
|
||
|
|
|
||
|
|
def unletterbox_mask_to_original(
|
||
|
|
masks_lb: torch.Tensor, # (N, Hlb, Wlb) bool/0-1
|
||
|
|
orig_hw: tuple, # (h0, w0)
|
||
|
|
new_unpad_wh: tuple, # (w_new, h_new)
|
||
|
|
pads_tblr: tuple # (top, bottom, left, right)
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Remove padding e volta masks para tamanho original.
|
||
|
|
"""
|
||
|
|
h0, w0 = orig_hw
|
||
|
|
w_new, h_new = new_unpad_wh
|
||
|
|
top, bottom, left, right = pads_tblr
|
||
|
|
|
||
|
|
# crop padding
|
||
|
|
masks_crop = masks_lb[:, top:top + h_new, left:left + w_new] # (N, h_new, w_new)
|
||
|
|
|
||
|
|
# resize para original com interpolate (torch)
|
||
|
|
masks_crop = masks_crop.unsqueeze(1).float() # (N,1,h_new,w_new)
|
||
|
|
masks_up = torch.nn.functional.interpolate(
|
||
|
|
masks_crop, size=(h0, w0), mode="bilinear", align_corners=False
|
||
|
|
)
|
||
|
|
return (masks_up[:, 0] > 0.5) # (N,h0,w0) bool
|
||
|
|
|
||
|
|
# ============================
|
||
|
|
# RAW loading (igual treino)
|
||
|
|
# ============================
|
||
|
|
def load_raw_as_uint8_hwc(img_path: Path, canais: int = CHANNELS) -> np.ndarray:
|
||
|
|
"""
|
||
|
|
Lê RAW (4 canais físicos R,G,IR,B) como CHW no arquivo, normaliza fixo e
|
||
|
|
retorna HWC uint8 com C=4 ou C=5 (com NDVI em 0..255).
|
||
|
|
"""
|
||
|
|
# usa o preview só pra pegar H,W (igual treino)
|
||
|
|
preview = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)
|
||
|
|
if preview is None:
|
||
|
|
raise FileNotFoundError(f"Não consegui ler preview: {img_path}")
|
||
|
|
h0, w0 = preview.shape[:2]
|
||
|
|
|
||
|
|
raw_path = img_path.with_suffix(".raw")
|
||
|
|
if not raw_path.exists():
|
||
|
|
raise FileNotFoundError(f"RAW não encontrado: {raw_path}")
|
||
|
|
|
||
|
|
num_raw_channels = 4
|
||
|
|
arr_flat = np.fromfile(str(raw_path), dtype=RAW_DTYPE)
|
||
|
|
expected_size = num_raw_channels * h0 * w0
|
||
|
|
if arr_flat.size != expected_size:
|
||
|
|
raise ValueError(f"RAW size errado {raw_path}: esperado {expected_size}, veio {arr_flat.size}")
|
||
|
|
|
||
|
|
raw_chw = arr_flat.reshape(num_raw_channels, h0, w0) # (4,H,W)
|
||
|
|
arr = np.transpose(raw_chw, (1, 2, 0)).astype(np.float32, copy=False) # (H,W,4)
|
||
|
|
|
||
|
|
# normalização fixa
|
||
|
|
if RAW_ALREADY_01:
|
||
|
|
arr_norm = arr
|
||
|
|
else:
|
||
|
|
arr_norm = arr / float(RAW_MAX_VAL)
|
||
|
|
arr_norm = np.clip(arr_norm, 0.0, 1.0)
|
||
|
|
|
||
|
|
# NDVI opcional (5ch)
|
||
|
|
if canais == 5:
|
||
|
|
R = arr_norm[..., 0]
|
||
|
|
IR = arr_norm[..., 2]
|
||
|
|
eps = 1e-6
|
||
|
|
ndvi = (IR - R) / (IR + R + eps) # [-1,1]
|
||
|
|
ndvi = np.clip(ndvi, -1.0, 1.0)
|
||
|
|
ndvi01 = (ndvi + 1.0) / 2.0 # [0,1]
|
||
|
|
arr_norm = np.concatenate([arr_norm, ndvi01[..., None]], axis=-1) # (H,W,5)
|
||
|
|
elif canais != 4:
|
||
|
|
raise ValueError(f"canais deve ser 4 ou 5, veio {canais}")
|
||
|
|
|
||
|
|
im_uint8 = (arr_norm * 255.0).round().clip(0, 255).astype(np.uint8)
|
||
|
|
return im_uint8 # HWC uint8
|
||
|
|
|
||
|
|
def make_bgr_preview_from_raw_uint8(raw_hwc_uint8: np.ndarray) -> np.ndarray:
|
||
|
|
"""
|
||
|
|
Preview BGR (OpenCV) a partir de HWC uint8:
|
||
|
|
assume ordem [R,G,IR,B,(NDVI)].
|
||
|
|
Usa B=canal 3, G=canal 1, R=canal 0.
|
||
|
|
"""
|
||
|
|
R = raw_hwc_uint8[..., 0]
|
||
|
|
G = raw_hwc_uint8[..., 1]
|
||
|
|
B = raw_hwc_uint8[..., 3]
|
||
|
|
return np.dstack([B, G, R])
|
||
|
|
|
||
|
|
# ============================
|
||
|
|
# Drawing
|
||
|
|
# ============================
|
||
|
|
def draw_result(vis_bgr: np.ndarray, boxes_xyxy: np.ndarray, cls_ids: np.ndarray, confs: np.ndarray, masks: np.ndarray | None):
|
||
|
|
out = vis_bgr.copy()
|
||
|
|
|
||
|
|
# máscaras (se houver)
|
||
|
|
if masks is not None:
|
||
|
|
for i in range(masks.shape[0]):
|
||
|
|
cls = int(cls_ids[i])
|
||
|
|
color = CLASS_COLORS[cls] if cls < len(CLASS_COLORS) else (0, 255, 0)
|
||
|
|
m = masks[i].astype(bool)
|
||
|
|
overlay = np.array(color, dtype=np.float32)
|
||
|
|
out[m] = (out[m] * 0.4 + overlay * 0.6).astype(np.uint8)
|
||
|
|
|
||
|
|
# boxes + labels
|
||
|
|
for i in range(len(boxes_xyxy)):
|
||
|
|
x1, y1, x2, y2 = boxes_xyxy[i].astype(int)
|
||
|
|
cls = int(cls_ids[i])
|
||
|
|
conf = float(confs[i])
|
||
|
|
color = CLASS_COLORS[cls] if cls < len(CLASS_COLORS) else (0, 255, 0)
|
||
|
|
name = CLASS_NAMES[cls] if cls < len(CLASS_NAMES) else str(cls)
|
||
|
|
|
||
|
|
cv2.rectangle(out, (x1, y1), (x2, y2), color, 2)
|
||
|
|
txt = f"{name} {conf:.2f}"
|
||
|
|
cv2.putText(out, txt, (x1, max(0, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2, cv2.LINE_AA)
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
# ============================
|
||
|
|
# Inference core
|
||
|
|
# ============================
|
||
|
|
@torch.no_grad()
|
||
|
|
def infer_one(model: YOLO, raw_uint8_hwc: np.ndarray, conf_thres=0.25, iou_thres=0.7, max_det=300):
|
||
|
|
"""
|
||
|
|
raw_uint8_hwc: HWC uint8 com C=4 ou 5
|
||
|
|
Retorna boxes/masks já no tamanho original (h0,w0)
|
||
|
|
"""
|
||
|
|
device = next(model.model.parameters()).device
|
||
|
|
|
||
|
|
h0, w0 = raw_uint8_hwc.shape[:2]
|
||
|
|
|
||
|
|
# 1) LetterBox igual treino (mantém proporção)
|
||
|
|
lb_img, ratio, pad, new_unpad_wh, pads_tblr = letterbox_multi(raw_uint8_hwc, new_shape=IMG_SZ, color=114)
|
||
|
|
|
||
|
|
# 2) Tensor (B,C,H,W) float32 0..1
|
||
|
|
x = torch.from_numpy(lb_img).to(device) # (H,W,C) uint8
|
||
|
|
x = x.permute(2, 0, 1).contiguous().float() / 255.0 # (C,H,W) float
|
||
|
|
x = x.unsqueeze(0) # (1,C,H,W)
|
||
|
|
|
||
|
|
# 3) Forward direto
|
||
|
|
out = model.model(x)
|
||
|
|
|
||
|
|
# 4) separa preds/proto
|
||
|
|
if isinstance(out, (list, tuple)):
|
||
|
|
preds_raw = out[0]
|
||
|
|
proto = out[1] if len(out) > 1 else None
|
||
|
|
else:
|
||
|
|
preds_raw = out
|
||
|
|
proto = None
|
||
|
|
|
||
|
|
if preds_raw.ndim == 2:
|
||
|
|
preds_raw = preds_raw.unsqueeze(0)
|
||
|
|
|
||
|
|
nc = len(model.names) if hasattr(model, "names") else len(CLASS_NAMES)
|
||
|
|
det_list = ops.non_max_suppression(
|
||
|
|
preds_raw,
|
||
|
|
conf_thres=conf_thres,
|
||
|
|
iou_thres=iou_thres,
|
||
|
|
classes=None,
|
||
|
|
agnostic=False,
|
||
|
|
max_det=max_det,
|
||
|
|
nc=nc
|
||
|
|
)
|
||
|
|
det = det_list[0]
|
||
|
|
if det is None or len(det) == 0:
|
||
|
|
return None
|
||
|
|
|
||
|
|
# det: [x1,y1,x2,y2,conf,cls, ...mask coeffs]
|
||
|
|
boxes_lb = det[:, :4]
|
||
|
|
confs = det[:, 4]
|
||
|
|
cls_ids = det[:, 5].to(torch.int64)
|
||
|
|
|
||
|
|
# 5) Máscaras no espaço letterboxed
|
||
|
|
masks_orig = None
|
||
|
|
if proto is not None and det.shape[1] > 6:
|
||
|
|
# proto pode vir com batch
|
||
|
|
proto_t = proto
|
||
|
|
while isinstance(proto_t, (list, tuple)):
|
||
|
|
proto_t = proto_t[0] if len(proto_t) else None
|
||
|
|
if proto_t is None:
|
||
|
|
break
|
||
|
|
|
||
|
|
if isinstance(proto_t, torch.Tensor):
|
||
|
|
if proto_t.ndim == 4:
|
||
|
|
proto_t = proto_t[0] # (C,Hm,Wm)
|
||
|
|
proto_t = proto_t.to(det.device)
|
||
|
|
|
||
|
|
mask_coeffs = det[:, 6:].to(det.device)
|
||
|
|
|
||
|
|
# garante C_proto == C_mask
|
||
|
|
c_proto = proto_t.shape[0]
|
||
|
|
c_mask = mask_coeffs.shape[1]
|
||
|
|
if c_proto != c_mask:
|
||
|
|
if c_proto > c_mask:
|
||
|
|
proto_t = proto_t[:c_mask]
|
||
|
|
else:
|
||
|
|
mask_coeffs = mask_coeffs[:, :c_proto]
|
||
|
|
|
||
|
|
# gera máscaras no tamanho do input (letterboxed)
|
||
|
|
masks_lb = ops.process_mask(proto_t, mask_coeffs, boxes_lb, x.shape[2:], upsample=True) # (N,Hlb,Wlb) bool
|
||
|
|
# desfaz letterbox -> original
|
||
|
|
masks_orig_t = unletterbox_mask_to_original(masks_lb, (h0, w0), new_unpad_wh, pads_tblr)
|
||
|
|
masks_orig = masks_orig_t.detach().cpu().numpy().astype(np.uint8) # 0/1
|
||
|
|
|
||
|
|
# 6) Boxes: letterbox -> original (usa ratio_pad)
|
||
|
|
# ratio_pad esperado: (ratio, pad) onde pad é (dw, dh)
|
||
|
|
ratio_pad = (ratio, pad)
|
||
|
|
boxes_scaled = ops.scale_boxes(x.shape[2:], boxes_lb.clone(), (h0, w0), ratio_pad=ratio_pad)
|
||
|
|
boxes_xyxy = boxes_scaled.detach().cpu().numpy()
|
||
|
|
|
||
|
|
return {
|
||
|
|
"boxes": boxes_xyxy,
|
||
|
|
"confs": confs.detach().cpu().numpy(),
|
||
|
|
"cls": cls_ids.detach().cpu().numpy(),
|
||
|
|
"masks": masks_orig
|
||
|
|
}
|
||
|
|
|
||
|
|
# ============================
|
||
|
|
# Main loop viewer
|
||
|
|
# ============================
|
||
|
|
def main():
|
||
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
|
|
print(f"Usando dispositivo: {device}")
|
||
|
|
|
||
|
|
model = YOLO(WEIGHTS_PATH)
|
||
|
|
model.to(device)
|
||
|
|
model.fuse()
|
||
|
|
|
||
|
|
# sanity check do patch
|
||
|
|
first_conv = model.model.model[0].conv
|
||
|
|
print("First conv weight shape:", tuple(first_conv.weight.shape)) # deve ser [48, C, 3, 3]
|
||
|
|
|
||
|
|
img_paths = sorted([p for p in TEST_DIR.iterdir() if p.suffix.lower() in [".jpg", ".jpeg", ".png"]])
|
||
|
|
if not img_paths:
|
||
|
|
print(f"Nenhuma imagem encontrada em {TEST_DIR}")
|
||
|
|
return
|
||
|
|
|
||
|
|
idx = 0
|
||
|
|
n = len(img_paths)
|
||
|
|
|
||
|
|
CONF = 0.45
|
||
|
|
IOU = 0.7
|
||
|
|
|
||
|
|
cv2.namedWindow(f"YOLO {CHANNELS}ch Viewer", cv2.WINDOW_NORMAL)
|
||
|
|
|
||
|
|
while True:
|
||
|
|
img_path = img_paths[idx]
|
||
|
|
print("\n======================================")
|
||
|
|
print(f"Processando: {img_path}")
|
||
|
|
|
||
|
|
try:
|
||
|
|
raw_uint8 = load_raw_as_uint8_hwc(img_path, canais=CHANNELS) # HWC uint8
|
||
|
|
preview_bgr = make_bgr_preview_from_raw_uint8(raw_uint8)
|
||
|
|
|
||
|
|
res = infer_one(model, raw_uint8, conf_thres=CONF, iou_thres=IOU)
|
||
|
|
if res is None:
|
||
|
|
vis = preview_bgr.copy()
|
||
|
|
cv2.putText(vis, "SEM DETECCOES", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,0,255), 2, cv2.LINE_AA)
|
||
|
|
else:
|
||
|
|
vis = draw_result(
|
||
|
|
preview_bgr,
|
||
|
|
boxes_xyxy=res["boxes"],
|
||
|
|
cls_ids=res["cls"],
|
||
|
|
confs=res["confs"],
|
||
|
|
masks=res["masks"]
|
||
|
|
)
|
||
|
|
|
||
|
|
txt = f"[{idx+1}/{n}] {img_path.name} | A/D navega | +/- conf {CONF:.2f} | Q/ESC sai"
|
||
|
|
cv2.putText(vis, txt, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2, cv2.LINE_AA)
|
||
|
|
cv2.imshow(f"YOLO {CHANNELS}ch Viewer", vis)
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Erro ao processar {img_path}: {e}")
|
||
|
|
blank = np.zeros((480, 900, 3), dtype=np.uint8)
|
||
|
|
cv2.putText(blank, f"Erro em {img_path.name}", (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,0,255), 2)
|
||
|
|
cv2.putText(blank, str(e)[:120], (10, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,255), 2)
|
||
|
|
cv2.imshow(f"YOLO {CHANNELS}ch Viewer", blank)
|
||
|
|
|
||
|
|
k = cv2.waitKey(0) & 0xFF
|
||
|
|
if k in (ord("q"), 27): # q ou ESC
|
||
|
|
break
|
||
|
|
elif k == ord("a"):
|
||
|
|
idx = (idx - 1) % n
|
||
|
|
elif k == ord("d"):
|
||
|
|
idx = (idx + 1) % n
|
||
|
|
elif k in (ord("+"), ord("=")):
|
||
|
|
CONF = min(0.99, CONF + 0.05)
|
||
|
|
elif k in (ord("-"), ord("_")):
|
||
|
|
CONF = max(0.01, CONF - 0.05)
|
||
|
|
|
||
|
|
cv2.destroyAllWindows()
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|