컴퓨터 비전 완전 정복: CNN부터 ViT, YOLO, Stable Diffusion까지

import torch
import torchvision.transforms as T
from PIL import Image

# 이미지 로드 및 텐서 변환
img = Image.open("sample.jpg").convert("RGB")
transform = T.Compose([
    T.Resize((224, 224)),
    T.ToTensor(),          # [0,255] -> [0.0,1.0], HWC -> CHW
    T.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225])
])
tensor = transform(img)   # shape: [3, 224, 224]
print(f"Shape: {tensor.shape}, dtype: {tensor.dtype}")

import torch
import torch.nn.functional as F

# 3x3 Sobel 엣지 검출 커널 (수평 방향)
kernel = torch.tensor([
    [-1, 0, 1],
    [-2, 0, 2],
    [-1, 0, 1]
], dtype=torch.float32).unsqueeze(0).unsqueeze(0)  # [1, 1, 3, 3]

# 그레이스케일 이미지에 컨볼루션 적용
gray_tensor = T.ToTensor()(T.Grayscale()(img)).unsqueeze(0)
edges = F.conv2d(gray_tensor, kernel, padding=1)

import albumentations as A
from albumentations.pytorch import ToTensorV2
import cv2

train_transform = A.Compose([
    A.RandomResizedCrop(224, 224, scale=(0.7, 1.0)),
    A.HorizontalFlip(p=0.5),
    A.ColorJitter(brightness=0.4, contrast=0.4,
                  saturation=0.4, hue=0.1, p=0.8),
    A.GaussianBlur(blur_limit=(3, 7), p=0.2),
    A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
                       rotate_limit=15, p=0.5),
    A.Normalize(mean=(0.485, 0.456, 0.406),
                std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

image = cv2.cvtColor(cv2.imread("sample.jpg"), cv2.COLOR_BGR2RGB)
augmented = train_transform(image=image)["image"]

import torch
import torch.nn as nn
import torchvision.models as models

class ResNetClassifier(nn.Module):
    def __init__(self, num_classes: int, pretrained: bool = True):
        super().__init__()
        weights = models.ResNet50_Weights.IMAGENET1K_V2 if pretrained else None
        self.backbone = models.resnet50(weights=weights)
        in_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(in_features, num_classes)
        )

    def forward(self, x):
        return self.backbone(x)

def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss, correct = 0.0, 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * images.size(0)
        correct += (outputs.argmax(1) == labels).sum().item()
    return total_loss / len(loader.dataset), correct / len(loader.dataset)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ResNetClassifier(num_classes=10).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-2)
criterion = nn.CrossEntropyLoss()

import timm

# EfficientNet-B4 파인튜닝
model = timm.create_model(
    "efficientnet_b4",
    pretrained=True,
    num_classes=100,
    drop_rate=0.3
)

# 백본 고정 (특징 추출 모드)
for name, param in model.named_parameters():
    if "classifier" not in name:
        param.requires_grad = False

model = timm.create_model(
    "convnext_large",
    pretrained=True,
    num_classes=1000
)
# 7x7 depthwise conv, LayerNorm, GELU 활성화 사용

from ultralytics import YOLO

# 사전 학습 모델 로드
model = YOLO("yolov8n.pt")   # nano: 속도 최우선
# model = YOLO("yolov8x.pt") # extra-large: 정확도 최우선

# 단일 이미지 추론
results = model("image.jpg", conf=0.25, iou=0.45)
for result in results:
    for box in result.boxes:
        cls = int(box.cls[0])
        conf = float(box.conf[0])
        x1, y1, x2, y2 = box.xyxy[0].tolist()
        print(f"Class: {result.names[cls]}, Conf: {conf:.2f}")

# 커스텀 데이터셋 파인튜닝
model.train(
    data="custom.yaml",
    epochs=100,
    imgsz=640,
    batch=16,
    lr0=0.01,
    device=0
)

# 평가
metrics = model.val()
print(f"mAP50: {metrics.box.map50:.3f}")
print(f"mAP50-95: {metrics.box.map:.3f}")

from ultralytics import YOLO

model = YOLO("yolov10n.pt")
# NMS 없이 바로 최종 결과 반환 — 지연 시간 단축
results = model.predict("video.mp4", stream=True)
for frame_result in results:
    print(frame_result.boxes)

import torch
from transformers import DetrImageProcessor, DetrForObjectDetection
from PIL import Image

processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")

img = Image.open("image.jpg")
inputs = processor(images=img, return_tensors="pt")
outputs = model(**inputs)

target_sizes = torch.tensor([img.size[::-1]])
results = processor.post_process_object_detection(
    outputs, target_sizes=target_sizes, threshold=0.9
)[0]

for score, label, box in zip(
    results["scores"], results["labels"], results["boxes"]
):
    name = model.config.id2label[label.item()]
    print(f"{name}: {score:.3f} at {[round(i,2) for i in box.tolist()]}")

import torch
import numpy as np
import cv2
from segment_anything import sam_model_registry, SamPredictor

sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
sam.to("cuda")
predictor = SamPredictor(sam)

# 이미지 설정
image = cv2.cvtColor(cv2.imread("image.jpg"), cv2.COLOR_BGR2RGB)
predictor.set_image(image)

# 포인트 프롬프트로 세그멘테이션
input_point = np.array([[500, 375]])  # 클릭 위치 (x, y)
input_label = np.array([1])           # 1=전경, 0=배경

masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True   # 여러 마스크 후보 반환
)

best_mask = masks[scores.argmax()]
print(f"Mask shape: {best_mask.shape}, Score: {scores.max():.3f}")

# 박스 프롬프트
input_box = np.array([100, 100, 400, 400])
masks_box, _, _ = predictor.predict(
    box=input_box[None, :],
    multimask_output=False
)

import torch
import torchvision.models.segmentation as seg_models

model = seg_models.deeplabv3_resnet101(
    weights=seg_models.DeepLabV3_ResNet101_Weights.COCO_WITH_VOC_LABELS_V1
)
model.eval()

with torch.no_grad():
    output = model(tensor.unsqueeze(0))["out"]  # [1, num_classes, H, W]
    pred = output.argmax(dim=1).squeeze()        # [H, W]
    print(f"Segmentation map shape: {pred.shape}")

from torchvision.models.detection import maskrcnn_resnet50_fpn_v2
from torchvision.models.detection import MaskRCNN_ResNet50_FPN_V2_Weights
from torchvision.utils import draw_segmentation_masks
import torchvision.transforms.functional as TF
from PIL import Image

weights = MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT
model = maskrcnn_resnet50_fpn_v2(weights=weights, box_score_thresh=0.75)
model.eval()

img = Image.open("people.jpg").convert("RGB")
inp = weights.transforms()(img).unsqueeze(0)

with torch.no_grad():
    predictions = model(inp)

pred = predictions[0]
masks = (pred["masks"] > 0.5).squeeze(1)
img_uint8 = (TF.to_tensor(img) * 255).byte()
result = draw_segmentation_masks(img_uint8, masks, alpha=0.5)

import timm
import torch

# ViT-Base/16 파인튜닝
model = timm.create_model(
    "vit_base_patch16_224",
    pretrained=True,
    num_classes=10,
    img_size=224
)

# ViT는 더 강한 증강 + AdamW + cosine schedule이 효과적
data_config = timm.data.resolve_model_data_config(model)
transforms_train = timm.data.create_transform(**data_config, is_training=True)
transforms_val = timm.data.create_transform(**data_config, is_training=False)

optimizer = torch.optim.AdamW(
    model.parameters(), lr=1e-3, weight_decay=0.05
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=100
)

import torch

# DINOv2 특징 추출기 (학습 없이 바로 사용)
dinov2 = torch.hub.load("facebookresearch/dinov2", "dinov2_vitl14")
dinov2.eval().cuda()

# 임의 크기 이미지 처리
import torchvision.transforms as T
preprocess = T.Compose([
    T.Resize(518),
    T.CenterCrop(518),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]),
])

with torch.no_grad():
    # [B, 3, 518, 518] -> [B, 1024] 특징
    features = dinov2(preprocess(img).unsqueeze(0).cuda())
    print(f"Feature dim: {features.shape}")  # [1, 1024]

# DINOv2는 k-NN 분류기로도 강력한 성능을 보임

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch

# SDXL 1.0 텍스트-이미지 생성
pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    use_safetensors=True,
)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config
)
pipe.to("cuda")

image = pipe(
    prompt="a photorealistic cat on a desk, 8k, studio lighting",
    negative_prompt="blurry, low quality, cartoon",
    num_inference_steps=25,
    guidance_scale=7.5,
    width=1024,
    height=1024
).images[0]
image.save("generated.png")

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import cv2
import numpy as np
from PIL import Image

controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-canny",
    torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    torch_dtype=torch.float16
).to("cuda")

# Canny 엣지맵 생성
ref_image = np.array(load_image("reference.jpg"))
edges = cv2.Canny(ref_image, 100, 200)
edges_rgb = np.stack([edges] * 3, axis=-1)

result = pipe(
    "a beautiful landscape painting",
    image=Image.fromarray(edges_rgb),
    num_inference_steps=20
).images[0]
result.save("controlnet_output.png")

from torch.utils.data import Dataset, DataLoader
from pathlib import Path
import albumentations as A
from albumentations.pytorch import ToTensorV2
import cv2

class CustomDataset(Dataset):
    def __init__(self, root: str, split: str = "train"):
        self.root = Path(root)
        self.image_paths = list(
            (self.root / split / "images").glob("*.jpg")
        )
        self.label_paths = [
            self.root / split / "labels" / p.with_suffix(".txt").name
            for p in self.image_paths
        ]
        self.transform = self._get_transforms(split)

    def _get_transforms(self, split: str):
        if split == "train":
            return A.Compose([
                A.RandomResizedCrop(224, 224, scale=(0.7, 1.0)),
                A.HorizontalFlip(p=0.5),
                A.ColorJitter(brightness=0.4, contrast=0.4,
                              saturation=0.4, hue=0.1, p=0.8),
                A.GaussianBlur(blur_limit=(3, 7), p=0.2),
                A.Normalize(mean=(0.485, 0.456, 0.406),
                            std=(0.229, 0.224, 0.225)),
                ToTensorV2(),
            ])
        return A.Compose([
            A.Resize(256, 256),
            A.CenterCrop(224, 224),
            A.Normalize(mean=(0.485, 0.456, 0.406),
                        std=(0.229, 0.224, 0.225)),
            ToTensorV2(),
        ])

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image = cv2.cvtColor(
            cv2.imread(str(self.image_paths[idx])), cv2.COLOR_BGR2RGB
        )
        label = int(self.label_paths[idx].read_text().strip())
        augmented = self.transform(image=image)
        return augmented["image"], label

train_dataset = CustomDataset("data/", split="train")
train_loader = DataLoader(
    train_dataset, batch_size=32, shuffle=True,
    num_workers=4, pin_memory=True, persistent_workers=True
)

import torch
import torch.onnx

model.eval()
dummy_input = torch.randn(1, 3, 224, 224, device="cuda")

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input":  {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print("ONNX 변환 완료: model.onnx")

# ONNX Runtime으로 검증
import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession(
    "model.onnx", providers=["CUDAExecutionProvider"]
)
input_name = sess.get_inputs()[0].name
result = sess.run(None, {input_name: dummy_input.cpu().numpy()})
print(f"ONNX 추론 결과 shape: {result[0].shape}")

# TensorRT 변환 (trtexec 사용)
trtexec \
  --onnx=model.onnx \
  --saveEngine=model_fp16.engine \
  --fp16 \
  --workspace=4096 \
  --optShapes=input:8x3x224x224

import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit

def trt_inference(engine_path: str, input_data: np.ndarray) -> np.ndarray:
    logger = trt.Logger(trt.Logger.WARNING)
    with open(engine_path, "rb") as f:
        engine = trt.Runtime(logger).deserialize_cuda_engine(f.read())
    context = engine.create_execution_context()

    input_mem = cuda.mem_alloc(input_data.nbytes)
    output = np.empty((input_data.shape[0], 1000), dtype=np.float32)
    output_mem = cuda.mem_alloc(output.nbytes)

    cuda.memcpy_htod(input_mem, input_data)
    context.execute_v2([int(input_mem), int(output_mem)])
    cuda.memcpy_dtoh(output, output_mem)
    return output