.history/mbench/gpt_ref-ytvos_numbered_cy_20250130183936.py

import os
import sys
from os import path as osp
from io import BytesIO

from mbench.ytvos_ref import build as build_ytvos_ref
import argparse
import opts

import sys
from pathlib import Path
import os
from os import path as osp
import skimage
from io import BytesIO

import numpy as np
import pandas as pd
import regex as re
import json

import cv2
from PIL import Image, ImageDraw
import torch
from torchvision.transforms import functional as F

from skimage import measure                        # (pip install scikit-image)
from shapely.geometry import Polygon, MultiPolygon # (pip install Shapely)

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle
import textwrap


import ipywidgets as widgets
from IPython.display import display, clear_output

from openai import OpenAI
import base64

def number_objects_and_encode(idx, color_mask=False):
    encoded_frames = {}
    contoured_frames = {}  # New dictionary for original images
    vid_cat_cnts = {}

    vid_meta = metas[idx]
    vid_data = train_dataset[idx]
    vid_id = vid_meta['video']
    frame_indx = vid_meta['sample_indx']
    cat_names = set(vid_meta['obj_id_cat'].values())
    imgs = vid_data[0]
    
    for cat in cat_names:
        cat_frames = []
        contour_frames = []
        frame_cat_cnts = {}

        for i in range(imgs.size(0)):
            frame_name = frame_indx[i]
            frame = np.copy(imgs[i].permute(1, 2, 0).numpy())  
            frame_for_contour = np.copy(imgs[i].permute(1, 2, 0).numpy()) 

            frame_data = vid_data[2][frame_name]
            obj_ids = list(frame_data.keys())

            cat_cnt = 0

            for j in range(len(obj_ids)):
                obj_id = obj_ids[j]
                obj_data = frame_data[obj_id]
                obj_bbox = obj_data['bbox']
                obj_valid = obj_data['valid']
                obj_mask = obj_data['mask'].numpy().astype(np.uint8)
                obj_cat = obj_data['category_name']

                if obj_cat == cat and obj_valid:
                    cat_cnt += 1
                    
                    if color_mask == False:
                        contours, _ = cv2.findContours(obj_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                        cv2.drawContours(frame, contours, -1, colors[j], 3)
                        for i, contour in enumerate(contours):
                            # 윤곽선 중심 계산
                            moments = cv2.moments(contour)
                            if moments["m00"] != 0:  # 중심 계산 가능 여부 확인
                                cx = int(moments["m10"] / moments["m00"])
                                cy = int(moments["m01"] / moments["m00"])
                            else:
                                cx, cy = contour[0][0]  # 중심 계산 불가시 대체 좌표 사용
                            
                            # 텍스트 배경 (검은색 배경 만들기)
                            font = cv2.FONT_HERSHEY_SIMPLEX
                            text = obj_id
                            text_size = cv2.getTextSize(text, font, 1, 2)[0]
                            text_w, text_h = text_size
                            
                            # 텍스트 배경 그리기 (검은색 배경)
                            cv2.rectangle(frame, (cx - text_w // 2 - 5, cy - text_h // 2 - 5),
                                        (cx + text_w // 2 + 5, cy + text_h // 2 + 5), (0, 0, 0), -1)
                            
                            # 텍스트 그리기 (흰색 텍스트)
                            cv2.putText(frame, text, (cx - text_w // 2, cy + text_h // 2),
                                        font, 1, (255, 255, 255), 2)

                    else:
                        alpha = 0.08

                        colored_obj_mask = np.zeros_like(frame)  
                        colored_obj_mask[obj_mask == 1] = colors[j]
                        frame[obj_mask == 1] = (
                            (1 - alpha) * frame[obj_mask == 1]
                            + alpha * colored_obj_mask[obj_mask == 1]
                        )


                        contours, _ = cv2.findContours(obj_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                        cv2.drawContours(frame, contours, -1, colors[j], 2)
                        cv2.drawContours(frame_for_contour, contours, -1, colors[j], 2)
                        
                        
                        
                        if len(contours) > 0:
                            largest_contour = max(contours, key=cv2.contourArea)
                            M = cv2.moments(largest_contour)
                            if M["m00"] != 0:  
                                center_x = int(M["m10"] / M["m00"])
                                center_y = int(M["m01"] / M["m00"])
                            else:
                                center_x, center_y = 0, 0

                        font = cv2.FONT_HERSHEY_SIMPLEX
                        text = obj_id

                        font_scale = 0.9  
                        text_size = cv2.getTextSize(text, font, font_scale, 2)[0]                        
                        text_x = center_x - text_size[0] // 1  # 텍스트의 가로 중심
                        text_y = center_y
                        # text_y = center_y + text_size[1] // 2 # 텍스트의 세로 중심

                        # 텍스트 배경 사각형 좌표 계산
                        rect_start = (text_x - 5, text_y - text_size[1] - 5)  # 배경 사각형 좌상단
                        # rect_end = (text_x + text_size[0] + 5, text_y + 5) 
                        rect_end = (text_x + text_size[0] + 5, text_y)

                        cv2.rectangle(frame, rect_start, rect_end, (0, 0, 0), -1)
                        cv2.putText(frame, text, (text_x, text_y), font, 1, (255, 255, 255), 2)

            # plt.figure(figsize=(12, 8))
            # plt.imshow(frame)
            # plt.title(f"frame {frame_name}")
            # plt.tight_layout()
            # plt.axis('off')
            # plt.show()
        
            buffer = BytesIO()
            frame = Image.fromarray(frame)
            frame.save(buffer, format='jpeg')
            buffer.seek(0)
            cat_frames.append(base64.b64encode(buffer.read()).decode("utf-8"))
            frame_cat_cnts[frame_name] = cat_cnt

            buffer.seek(0)  # Reuse buffer instead of creating a new one
            buffer.truncate()
            frame_for_contour = Image.fromarray(frame_for_contour)
            frame_for_contour.save(buffer, format='jpeg')
            buffer.seek(0)
            contour_frames.append(base64.b64encode(buffer.read()).decode("utf-8"))
        
        encoded_frames[cat] = cat_frames
        contoured_frames[cat] = contour_frames
        vid_cat_cnts[cat] = frame_cat_cnts
    
    return encoded_frames, vid_cat_cnts, contoured_frames 

if __name__ == '__main__':
    parser = argparse.ArgumentParser('ReferFormer training and evaluation script', parents=[opts.get_args_parser()])
    args = parser.parse_args()

    #==================데이터 불러오기===================
    # 전체 데이터셋
    train_dataset = build_ytvos_ref(image_set = 'train', args = args)

    # 전체 데이터셋 메타데이터
    metas = train_dataset.metas

    # 색상 후보 8개 (RGB 형식)
    colors = [
        (255, 0, 0),    # Red
        (0, 255, 0),    # Green
        (0, 0, 255),    # Blue
        (255, 255, 0),  # Yellow
        (255, 0, 255),  # Magenta
        (0, 255, 255),  # Cyan
        (128, 0, 128),  # Purple
        (255, 165, 0)   # Orange
    ]