import torch, matplotlib.pyplot as plt, matplotlib.patches as patches

def draw_prediction(img, pred, S=20, B=3, thr=0.15):
    img_np = img.permute(1,2,0).numpy()
    H,W,_ = img_np.shape
    fig,ax=plt.subplots(1,figsize=(6,6)); ax.imshow(img_np)
    c=0
    for i in range(S):
        for j in range(S):
            for b in range(B):
                off=b*5
                conf=float(torch.sigmoid(pred[i,j,off+4]))
                if conf>thr:
                    cx,cy,w,h = [float(pred[i,j,off+k]) for k in range(4)]
                    cx*=W; cy*=H; w*=W; h*=H
                    x=cx-w/2; y=cy-h/2
                    ax.add_patch(patches.Rectangle((x,y),w,h,edgecolor='red',facecolor='none',linewidth=2))
                    c+=1
    ax.set_title(f"{c} boxes"); plt.axis('off'); plt.show()

def load_and_predict(model_path, img_tensor, model):
    model.load_state_dict(torch.load(model_path, map_location='cpu'))
    model.eval()
    with torch.no_grad():
        out=model(img_tensor.unsqueeze(0)).squeeze(0)
    draw_prediction(img_tensor, out)