[사내 해커톤] Stable Diffusion을 이용한 영상 가공

 

이번에도 어김없이 회사에서 해커톤이 열렸다. 주제는 "생성형 AI를 사용한 회사에 도움이 되는 서비스를 만들어" 였다.

일단 아이디어 도출에서 텍스트로 하는건 재미가 없으므로 패스했고, (프롬프트 튜닝은 너무 많이 해서 재미가 없다)

영상과 음성쪽으로 해보고 싶어서 가닥을 광고쪽으로 컨셉을 잡고 진행했다.

 

영상의 경우 기존 영상을 일부분만 수정하는 형태로 아이디어를 냈다. 예를 들어 사과를 소개하는 영상에서 사과 대신 배, 복숭아로 바꾸는 형태로 영상을 바꿔서 원하는 영상으로 교체하는것으로 생각했다.

 

간단한 예시는 다음과 같다.

https://tilnote.io/pages/640adfbef4ea08b9071cc823

https://platform.stability.ai/docs/features/inpainting#Python

Stability AI - Developer Platform

 

원본 이미지를 준비한다!

 

변경하려는 부분을 mask 이미지로 준비한다 (검정색 부분을 바꾼다)

 

최종 결과물은 다음과 같이 나오게 된다

이러한 연속된 이미지들을 무수히 많이 만들어서 합친다 == 영상!

 

 

절차는 이렇게 된다.

1. 비디오를 프레임마다 이미지로 저장

2. 이미지에서 변경할 부분 detect

3. mask 이미지 생성

4. stable diffusion을 사용해서 이미지 수정

5. 이미지를 비디오로 변환

 

으로 진행된다.

---

해당 블로그의 내용은

* 아래의 코드는 사내에서만 사용가능한 버드락을 사용하여, 모델 부분의 코드 사용시 일부분을 수정해서 사용하셔야 합니다. 흐름만 보시는것을 추천합니다.

*  음성의 경우 RVC를 사용해서 노래 합성쪽으로 했다. 해당 코드는 유튜브에 검색하면 바로 나와서 따로 블로그에 쓰진 않았다.

---

 

원본 동영상

take2.mp4

3.06MB

 

 

1. 비디오 to 이미지 변환

일단 비디오를 이미지로 변환해야 한다. 해당 영상은 총 97프레임이 나왔다.

import cv2
vidcap = cv2.VideoCapture('/content/take2.mp4')
success,image = vidcap.read()
count = 0
while success:
  cv2.imwrite("/content/frame/take2/frame%d.jpg" % count, image)     # save frame as JPEG file      
  success,image = vidcap.read()
  print('Read a new frame: ', success)
  count += 1

 

 

 

2. frame detect (label)

해당 영상은 배를 다른 과일로 변환하는 영상을 만들고 싶었다.

프레임마다 해당 과일의 위치를 파악하고, 해당 위치에만 다른 색상(검정색)으로 칠해야 한다.

detect를 위해 yolo를 사용해서 과일 추적하는 코드를 작성했다.

마지막 파라미터인 class는 해당 라벨만 사용하겠다는 뜻이다. (사과와 오렌지만 추적한다)

ys = YOLO('yolov8m.pt')

for i in range(0, 97):
  scores = ys.predict(source=f"/content/frame/take2/frame{str(i)}.jpg", conf=0.5, save=True, save_txt=True, classes=[47, 49])

 

완료되면 해당 라벨들의 위치는 text과 라벨링된 이미지가 저장된다.

 

 

3. 라벨의 위치를 읽고 mask 이미지 생성

라벨은 다음과 같이 저장됩니다.

47 0.495016 0.486613 0.330306 0.492391

첫번째는 라벨번호로 오렌지를 뜻한다.

나머지 4개의 숫자는 x_center, y_center, width, height 를 의미합니다.

 

해당 표기법은 albumentations식으로 다음을 참조하시면 됩니다.

https://albumentations.ai/docs/getting_started/bounding_boxes_augmentation/

albumentations과 표현 방식이 같으며 전체 가로와 세로를 기점으로 %로 표기되어 있다고 표시면 됩니다.

아래의 코드는 yolo 좌표를 sacla로 변경해 줍니다.

# yolo 좌표를 scala로 변경
def get_bounding_box_corrdinates(w, h, label_position):
    print("label_position", label_position.rstrip().split(" "))
    label, x_center, y_center, width, height = label_position.rstrip().split(" ")
    x_min = int(w * max(float(x_center) - float(width) / 2, 0))
    x_max = int(w * min(float(x_center) + float(width) / 2, 1))
    y_min = int(h * max(float(y_center) - float(height) / 2, 0))
    y_max = int(h * min(float(y_center) + float(height) / 2, 1))
    
    return x_min, y_min, x_max, y_max

 

 

 

모든 프레임마다 해당 좌표값을 얻은 다음 mask 이미지를 만들어야 합니다.

모든 라벨이 포함된 파일을 돌면서 해당 위치에 다른 색상의 백그라운드 이미지 파일을 만듭니다.

검정색은 변환되는 부분 바탕인 흰색은 변경되지 않는 부분으로 인식됩니다.

# 라벨 text를 읽어서 이미지를 생성

import numpy as np
import os.path
import tqdm 
from google.colab.patches import cv2_imshow


x_min, y_min, x_max, y_max = 0, 0, 0, 0

for i in range(97):
  file_path = f"/content/runs/detect/predict/labels/frame{i}.txt"

  # 라벨정보 파일 읽기
  with open(file_path, "r") as f:
    label_position = f.readlines()
    x_min, y_min, x_max, y_max = get_bounding_box_corrdinates(w, h, label_position[0])

  # 백그라운드 이미지 생성 
  white_img = np.full((h, w), 255, dtype=np.uint8)
  cv2.rectangle(white_img, (x_min, y_min), (x_max, y_max), (0, 255, 0), -1)
  cv2.imwrite(f'/content/frame/take2/background/frame{i}.jpg',white_img)

 

 

 

 

4. 이미지 리사이즈

사용하는 스테이블 디퓨전 모델에 따라 지원하는 이미지 사이즈가 정해져 있습니다.  (저의 경우 768, 512 사이즈로 정했습니다)

사이즈에 맞게 모든 이미지를 리사이즈 합니다.

import os 
import cv2 

for i in range(1, 32):
    origin_img = cv2.imread(f"./content/frame/take2/frame{str(i).zfill(3)}.jpg")
    resize_img = cv2.resize(origin_img, (768, 512))
    cv2.imwrite(f'./resize/take2/frame{i}.jpg',resize_img)
    
    background_img = cv2.imread(f"./content/frame/take2/background/frame{str(i).zfill(3)}.jpg")
    resize_img = cv2.resize(background_img, (768, 512))
    cv2.imwrite(f'./resize/take2/background/frame{i}.jpg',resize_img)

 

 

5. 모델을 사용해서 이미지 생성

정해진 모델인 aws의 버드락을 사용해서 스테이블 디퓨전을 실행했습니다. (알맞게 사용하시는 모델로 변경하시면 됩니다.)

버드락의 경우 이미지를 base64로 읽고, 쓰기 때문에 변환 코드를 작성해 줍니다.

def image_to_base64(img) -> str:
    """Convert a PIL Image or local image file path to a base64 string for Amazon Bedrock"""
    if isinstance(img, str):
        if os.path.isfile(img):
            print(f"Reading image from file: {img}")
            with open(img, "rb") as f:
                return base64.b64encode(f.read()).decode("utf-8")
        else:
            raise FileNotFoundError(f"File {img} does not exist")
    elif isinstance(img, Image.Image):
        print("Converting PIL Image to base64 string")
        buffer = io.BytesIO()
        img.save(buffer, format="PNG")
        return base64.b64encode(buffer.getvalue()).decode("utf-8")
    else:
        raise ValueError(f"Expected str (filename) or PIL Image. Got {type(img)}")

 

 

프레임별 변환을 하고 저장합니다.

프롬프트는 생성과 negative로 나눠집니다. 그리고 수많은 조정 파라미터도 있습니다.

해당 프롬프트와 여러 셋팅값을 자신에 맞게 조정한 후 최종 결과 이미지를 얻습니다. (엄청나게 많이 테스트 해봐야 합니다.)

import time 
negative_prompts = [
    "poorly rendered",
    "poor background details",
    "poorly drawn animal",
    "disfigured animal features",
]
style_preset = "digital-art"  # (e.g. photographic, digital-art, cinematic, ...)
clip_guidance_preset = "FAST_BLUE" # (e.g. FAST_BLUE FAST_GREEN NONE SIMPLE SLOW SLOWER SLOWEST)
sampler = "DDIM" # (e.g. DDIM, DDPM, K_DPMPP_SDE, K_DPMPP_2M, K_DPMPP_2S_ANCESTRAL, K_DPM_2, K_DPM_2_ANCESTRAL, K_EULER, K_EULER_ANCESTRAL, K_HEUN, K_LMS)
width = 768

inpaint_prompt = "Dancing otters, colourful background, highly detailed, digital painting"

for i in range(1, 32):
    print(i)
    origin_image = Image.open(f'./resize-bag/frame{i}.jpg').convert("RGB")
    mask = Image.open(f'./resize-bag-background/frame{i}.jpg').convert("RGB")
    
    request = json.dumps({
        "text_prompts":(
            [{"text": inpaint_prompt, "weight": 1.0}]
            + [{"text": negprompt, "weight": -1.0} for negprompt in negative_prompts]
        ),
        "init_image": image_to_base64(origin_image),
        "mask_source": "MASK_IMAGE_BLACK",
        "mask_image": image_to_base64(mask),
        "cfg_scale": 16,
        "seed": 1,
        "style_preset": style_preset,
    })


    modelId = "stability.stable-diffusion-xl-v0"

    response = boto3_bedrock.invoke_model(body=request, modelId=modelId)
    response_body = json.loads(response.get("body").read())

    image_2_b64_str = response_body["artifacts"][0].get("base64")

    inpaint = Image.open(io.BytesIO(base64.decodebytes(bytes(image_2_b64_str, "utf-8"))))
    inpaint.save(f'./final/frame{i}.jpg')

 

 

 

6. 이미지 to mp4

최종 결과물 이미지들을 하나의 영상으로 합칩니다. (해상도많은 차이가 있으니 전문 툴을 사용하시는것을 추천드립니다.)

def img2mp4(paths, pathOut , fps = 10 ) :
    frame_array = []
    for idx , path in enumerate(paths) : 
        img = cv2.imread(path)
        height, width, layers = img.shape
        size = (width,height)
        frame_array.append(img)
    out = cv2.VideoWriter(pathOut,cv2.VideoWriter_fourcc(*'DIVX'), fps, size)
    for i in range(len(frame_array)):
        out.write(frame_array[i])
    out.release()
    print("done")



import os
paths = [f"/content/runs/detect/predict/frame{i}.jpg" for i in range(0, 145)]
img2mp4(paths , "/content/output_path/test.mp4", fps=20)

 

 

 

최종 결과물

완성.mp4

2.53MB

 

원본영상

합성영상

 

 

 

 

보안해야 할점

* 마스크 영역이 사각형으로 되어 있어서 이미지 변환시 너무 크게 변환이 되는 단점이 있다. 정확하게 영역만 마스크 처리 해야한다.

* 생각보다 실사 이미지생성시 눈에 띄게 이질감이 느껴진다. 프롬프트를 바꾸고, 모델 버전을 바꿔도 똑같다. 적절한 lora를 선택해야 한다.

* 모델에 따라 해상도가 확연하게 떨어진다. 그리고 높히려 하면 돈이 많이든다. 테스트시 주의가 필요하다

 

끝!