samurai 旋转分类增强

samurai 旋转分类增强

samurai_xuanzhuan.py

import glob
import cv2
import numpy as np
import json
import os

import time


class Time_Str:
    def __init__(self):
        self.last_time, self.counter = '', 0

    def get_time(self, fmt="%m%d_%H%M_%S"):
        now_time = time.strftime(fmt)
        self.counter += 1
        if now_time != self.last_time:
            self.last_time, self.counter = now_time, 0
        return f"{now_time}_{self.counter}"

    def get_time_sec(self):
        return self.get_time("%m%d_%H%M_%S")

    def get_time_mm(self):
        return self.get_time("%m%d_%H%M")

    def get_time_hh(self):
        return self.get_time("%m%d_%H")
def rotate_image(image, angle, center):
    h, w = image.shape[:2]
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rot = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_LINEAR, borderValue=(0, 0, 0))
    return rot, M

def rotate_points(points, M):
    pts = np.array(points, dtype=np.float32)
    ones = np.ones((pts.shape[0], 1), dtype=np.float32)
    pts_ex = np.hstack([pts, ones])
    rotated = M.dot(pts_ex.T).T
    return rotated

def get_AABB(rotated_contour):
    xs = rotated_contour[:, 0]
    ys = rotated_contour[:, 1]
    x1, y1 = int(np.min(xs)), int(np.min(ys))
    x2, y2 = int(np.max(xs)), int(np.max(ys))
    return (x1, y1, x2, y2)

def crop_AABB(rotated_img, box):
    x1, y1, x2, y2 = box
    x1, y1 = max(0, x1), max(0, y1)
    x2, y2 = min(rotated_img.shape[1], x2), min(rotated_img.shape[0], y2)
    if x2 <= x1 or y2 <= y1:
        return None
    return rotated_img[y1:y2, x1:x2]

def draw_center(img, center, color, text):
    pt = (int(round(center[0])), int(round(center[1])))
    cv2.circle(img, pt, 5, color, -1)
    cv2.putText(img, text, (pt[0]+6, pt[1]-6), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)


def visualize_contour(img, contour, winname="contour_vis", save_path=None):
    """
    在图像上可视化 contour（绿色），并显示每个点的 index（白色）
    """
    vis = img.copy()
    cnt_int = contour.astype(np.int32)

    # 绘制 contour
    cv2.polylines(vis, [cnt_int], True, (0, 255, 0), 2)

    # 标注每个点
    for i, (x, y) in enumerate(cnt_int):
        cv2.circle(vis, (x, y), 4, (0, 0, 255), -1)  # 点
        cv2.putText(vis, str(i), (x + 5, y - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)

    # 显示
    try:
        cv2.imshow(winname, vis)
        cv2.waitKey(0)
        cv2.destroyWindow(winname)
    except:
        pass

if __name__ == '__main__':
    base_dir = r"D:\data\chantu\1120"   # 改成你的路径

    img_files = ['%s/%s' % (i[0].replace("\\", "/"), j) for i in os.walk(base_dir) for j in i[-1] if
                  j.lower().endswith(('jpg', 'xjson', 'xpeg'))]

    out_dir = "rotated_crop_results_viz"
    os.makedirs(out_dir, exist_ok=True)

    name_get=Time_Str()
    for img_path in img_files:
        basename = os.path.splitext(os.path.basename(img_path))[0]
        json_path = os.path.splitext(img_path)[0] + '.json'
        if not os.path.exists(json_path):
            print("No json for", img_path)
            continue
        print(json_path)
        with open(json_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        img = cv2.imread(img_path)

        vis_orig = img.copy()

        # 找到第一个 contour
        contour = None
        for s in data.get("shapes", []):
            if s.get("shape_type") == "contour":
                contour = np.array(s["points"], dtype=np.int32)
                break
        if contour is None :
            print("No valid contour in", json_path)
            continue

        cv2.polylines(vis_orig, [contour], isClosed=True, color=(0,255,0), thickness=2)

        x1, y1, w, h = cv2.boundingRect(contour)
        bbox = [x1, y1, x1 + w, y1 + h]

        # 3. 绘制 bounding box（红色）
        cv2.rectangle(vis_orig, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 2)
        if 0:
            center = contour.mean(axis=0)  # [cx, cy]
            cx, cy = float(center[0]), float(center[1])
            center = (cx, cy)

        center = (x1 + w // 2, y1 + h // 2)

        draw_center(vis_orig, center, (0, 0, 255), "use")
        cv2.imshow('asf',vis_orig)

        # 对每个角度做可视化
        for angle in range(0, 360, 15):
            rotated_img, RM = rotate_image(img, angle, center)
            rotated_contour = rotate_points(contour, RM)
            # AABB
            x1, y1, x2, y2 = get_AABB(rotated_contour)
            crop_img = crop_AABB(rotated_img, (x1, y1, x2, y2))

            # 可视化旋转后图像（contour + AABB + center）
            vis_rot = rotated_img.copy()
            cv2.polylines(vis_rot, [rotated_contour.astype(np.int32)], True, (0,255,0), 2)
            cv2.rectangle(vis_rot, (x1, y1), (x2, y2), (0,0,255), 2)
            draw_center(vis_rot, center, (255,0,255), "center")
            # 在图上写角度信息
            cv2.putText(vis_rot, f"angle: {angle}", (10,30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,255,255), 2, cv2.LINE_AA)

            # 显示 crop（若存在）
            if crop_img is None:
                crop_preview = np.zeros((100,200,3), dtype=np.uint8)
                cv2.putText(crop_preview, "invalid crop", (10,50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,0,255), 2)
            else:
                crop_preview = crop_img

            # 保存文件（包含源名与角度）
            # crop_save = os.path.join(out_dir, f"{basename}_crop_{name_get.get_time_mm()}.jpg")
            crop_save = os.path.join(out_dir, f"{name_get.get_time_mm()}.jpg")
            # cv2.imwrite(rot_save, rotated_img)

            cv2.imshow('rotated_img',rotated_img)
            cv2.imshow('vis_rot',vis_rot)
            if crop_img is not None:
                cv2.imwrite(crop_save, crop_img)

            # 键控：任意键继续，s = 保存（文件已保存）， q = 跳到下张图， e = 退出程序
            k2 = cv2.waitKey(1) & 0xFF
            if k2 == ord('e'):
                print("Exit by user")
                cv2.destroyAllWindows()
                exit(0)
            if k2 == ord('q'):
                # 退出当前图片角度循环，进入下一个图片
                cv2.destroyWindow("rotated_vis")
                cv2.destroyWindow("crop_preview")
                break
            # 其它按键默认继续到下一个角度
        print(f"[{basename}] Done and saved to {out_dir}")
    cv2.destroyAllWindows()

    print("All files processed.")