Attention based CNN
import torch
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from torchvision.utils import save_image

import urllib.request
import zipfile
import glob
import os
from PIL import Image, UnidentifiedImageError, ImageFile

import multiprocessing as mp
from multiprocessing import freeze_support

import matplotlib.pyplot as plt


class ImageAug():
    def imtransform(self, opt, transform):
        batch_size = 16
        original_images = self.originalset()
        # 이미지 폴더로부터 데이터를 로드합니다.
        transform_dataset = ImageFolder(root='C:\\Users\\DH\\Desktop\\tmp\\PetImages',  # 다운로드 받은 폴더의 root 경로를 지정합니다.
                                        transform=transform)

        # 데이터 로더를 생성합니다.
        transform_loader = DataLoader(transform_dataset,  # 이전에 생성한 transform_dataset를 적용합니다.
                                      batch_size=batch_size,  # 배치사이즈
                                      shuffle=False,  # 셔플 여부
                                      num_workers=1
                                      )

        transform_images, labels = next(iter(transform_loader))

        fig, axes = plt.subplots(3, 2)
        fig.set_size_inches(4, 6)

        if opt == "show":
            for idx in range(3):
                axes[idx, 0].imshow(original_images[idx].permute(1, 2, 0))
                axes[idx, 0].set_axis_off()
                axes[idx, 0].set_title('Original')
                axes[idx, 1].imshow(transform_images[idx].permute(1, 2, 0))
                axes[idx, 1].set_axis_off()
                axes[idx, 1].set_title('Transformed')
            fig.tight_layout()
            plt.show()
        elif opt == "save":
            os.makedirs('tmp', exist_ok=True)
            # 폴더 생성 및 이미지 저장
            for idx in range(batch_size):
                save_image(transform_images[idx], 'tmp/test'+str(idx)+'.png')
        elif opt == "get_tensor":
            return transform_images
        else:
            raise Exception("not correct command")

        # print(type(transform_images))
        # print(type(transform_dataset))

    def originalset(self):
        # 랜덤 시드 설정
        torch.manual_seed(321)
        # 이미지 크기를 224 x 224 로 조정합니다
        IMAGE_SIZE = 224

        original_dataset = ImageFolder(root='C:\\Users\\DH\\Desktop\\tmp\\PetImages',  # 다운로드 받은 폴더의 root 경로를 지정합니다.
                                       transform=transforms.Compose([  # Resize 후 정규화(0~1)를 수행합니다.
                                           transforms.Resize(
                                               (IMAGE_SIZE, IMAGE_SIZE)),
                                           # 개와 고양이 사진 파일의 크기가 다르므로, Resize로 맞춰줍니다.
                                           transforms.ToTensor()
                                       ]))

        original_loader = DataLoader(original_dataset,  # 이전에 생성한 original_dataset를 로드 합니다.
                                     batch_size=16,  # 배치사이즈
                                     shuffle=False,  # 셔플 여부
                                     num_workers=1
                                     )

        # 1개의 배치를 추출합니다.
        original_images, labels = next(iter(original_loader))

        # 이미지의 shape을 확인합니다. 224 X 224 RGB 이미지 임을 확인합니다.
        # (batch_size, channel, height, width)
        original_images.shape
        return original_images

    def trans_colorJitter(self, brightness, contrast, saturation, hue):
        image_transform = transforms.Compose([
            transforms.Resize((256, 256)),
            # ColorJitter 적용
            transforms.ColorJitter(brightness=brightness,
                                   contrast=contrast,
                                   saturation=saturation,
                                   hue=hue,
                                   ),
            transforms.ToTensor()
        ])
        return image_transform

    def trans_horizontal_flip(self, p):
        image_transform = transforms.Compose([
            transforms.Resize((256, 256)),
            # RandomHorizontalFlip 적용
            transforms.RandomHorizontalFlip(p=p),
            transforms.ToTensor()
        ])
        return image_transform

    def trans_gaussian_blur(self, kernel_size, sigma):
        image_transform = transforms.Compose([
            transforms.Resize((256, 256)),
            # GaussianBlur 적용
            transforms.GaussianBlur(kernel_size=kernel_size, sigma=sigma),
            transforms.ToTensor()
        ])
        return image_transform

    def trans_rotate(self, degrees, fill):
        image_transform = transforms.Compose([
            transforms.Resize((256, 256)),
            # RandomRotation 적용
            transforms.RandomRotation(
                degrees=degrees, interpolation=transforms.InterpolationMode.BILINEAR, fill=fill),
            transforms.ToTensor()
        ])
        return image_transform

    def trans_padding(self, padding, fill, padding_mode):
        image_transform = transforms.Compose([
            transforms.Resize((256, 256)),
            # Pad 적용
            transforms.Pad(padding=padding, fill=fill,
                           padding_mode=padding_mode),
            transforms.ToTensor()
        ])
        return image_transform

    def autotransform(self, policy):
        image_transform = transforms.Compose([
            transforms.Resize((256, 256)),
            # AutoAugment 적용
            transforms.AutoAugment(policy=policy,
                                   interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.ToTensor()
        ])
        return image_transform

    def autotransform(self, policy):
        image_transform = transforms.Compose([
            transforms.Resize((256, 256)),
            # AutoAugment 적용
            transforms.AutoAugment(policy=transforms.autoaugment.AutoAugmentPolicy.IMAGENET,
                                   interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.ToTensor()
        ])
        return image_transform


if __name__ == '__main__':
    freeze_support()

    IA = ImageAug()
    # IA.download_dataset(
    #    'https://download.microsoft.com/download/3/E/1/3E1C3F21-ECDB-4869-8368-6DEBA77B919F/kagglecatsanddogs_5340.zip',
    #    'PetImages') # 데이터셋 zip 경로, 생성할 폴더 이름
    # transform_info = IA.trans_colorJitter((0.5, 0.9), (0.4, 0.8), (0.7, 0.9), (-0.2, 0.2))
    # transform_info = IA.trans_horizontal_flip(0.8)
    # transform_info = IA.trans_gaussian_blur((19, 19), (1.0, 2.0))
    # transform_info = IA.trans_rotate((-30, 30), 0)
    transform_info = IA.trans_padding((100, 50, 100, 200), 255, 'symmetric')
    # transform_info = IA.autotransform(
    # transforms.autoaugment.AutoAugmentPolicy.IMAGENET)

    IA.imtransform("show", transform_info)

AI/Augment.py
icanSeeo/i-can-see
import torch
from utils import *
from model import *
from torchvision.datasets import CIFAR100
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from tqdm import tqdm
import torch.nn as nn
import matplotlib.pyplot as plt

# Define functions for inference and interpretation


def perform_inference(model, input_batch):
    with torch.no_grad():
        output, attention_map = model(input_batch)
    return output, attention_map


def interpret_output(output, class_names):
    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    top5_prob, top5_catid = torch.topk(probabilities, 5)

    for i in range(top5_prob.size(0)):
        print(class_names[top5_catid[i]], top5_prob[i].item())


# Define transformations and load CIFAR-100 dataset
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

model = ResNetWithCBAM(Bottleneck, [3, 4, 6, 3], num_classes=1000)
model.apply_pretrained()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

num_epochs = 200
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()

# Load the CIFAR-100 dataset
train_dataset = CIFAR100(root='./data', train=True,
                         download=True, transform=transform)

# Define a DataLoader
batch_size = 200
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# Initialize lists to store accuracy and loss
accuracy_list = []
loss_list = []

for epoch in tqdm(range(num_epochs)):
    model.train()
    running_loss = 0.0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs, aux_output = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}")

    # Calculate accuracy after each epoch
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            outputs, aux_output = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    accuracy_list.append(accuracy)
    loss_list.append(running_loss / len(train_loader))

    print(
        f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}, Accuracy: {accuracy}%")

# Plot the accuracy graph
plt.plot(range(1, num_epochs+1), accuracy_list, label='Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy (%)')
plt.title('Accuracy over Epochs')
plt.legend()
plt.show()

# Plot the loss graph
plt.plot(range(1, num_epochs+1), loss_list, label='Loss', color='red')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Loss over Epochs')
plt.legend()
plt.show()

# Save the model weights
model.save_weights('cbam-finetunned-200-notpretrained(200batch).pth')

AI/fine-tunning.py
icanSeeo/i-can-see
import torch
from utils import *
from model import *

# 추론을 수행하는 함수 정의


def perform_inference(model, input_batch):
    with torch.no_grad():
        output, attention_map = model(input_batch)
    return output, attention_map

# 결과를 해석하는 함수 정의


def interpret_output(output, class_names):
    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    top5_prob, top5_catid = torch.topk(probabilities, 5)

    for i in range(top5_prob.size(0)):
        print(class_names[top5_catid[i]], top5_prob[i].item())


# 이미지 경로 설정
image_path = 'data/cifar-100-python/train/buildings'
# 이미지 전처리 및 모델 준비
input_batch = preprocess_image(image_path)

model = ResNetWithCBAM.load_weights(
    'cbam-finetunned.pth', Bottleneck, [3, 4, 6, 3], num_classes=1000)

model.eval()
model, input_batch = prepare_model_and_input(model, input_batch)

# 추론 수행
output, attention_map = perform_inference(model, input_batch)

# 클래스 인덱스를 클래스명으로 변환
class_names = load_class_names()

# 결과 해석 및 출력
interpret_output(output, class_names)

original_image = cv2.imread('piano.jpg')
print(original_image.shape)
overlay_attention_map(
    original_image, attention_map.cpu().numpy(), 'output_image.jpg')

AI/inference.py
icanSeeo/i-can-see
from .cbam_resnet import *
from .baseline_resnet import *
AI/model/__init__.py
icanSeeo/i-can-see
import torch
import torchvision.transforms as transforms
from torchvision.models import resnet50
from PIL import Image
import json

# Load class labels from JSON file
with open('imagenet_classes.json') as f:
    class_names = json.load(f)

# Load pre-trained ResNet model
model = resnet50(pretrained=True)
model.eval()

# Image preprocessing
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


if __name__ == "__main__":
    # Load and preprocess the image
    image_path = 'piano.jpg'
    input_image = Image.open(image_path).convert('RGB')
    input_tensor = preprocess(input_image)
    input_batch = input_tensor.unsqueeze(0)

    # Use GPU if available
    if torch.cuda.is_available():
        model.to('cuda')
        input_batch = input_batch.to('cuda')

    # Inference
    with torch.no_grad():
        output = model(input_batch)

    # Convert output to probabilities
    probabilities = torch.nn.functional.softmax(output[0], dim=0)

    # Get top 5 predictions
    top5_prob, top5_catid = torch.topk(probabilities, 5)

    # Print top 5 predictions
    for i in range(top5_prob.size(0)):
        print(class_names[top5_catid[i]], top5_prob[i].item())

AI/model/baseline_resnet.py
icanSeeo/i-can-see
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.models import resnet50
from PIL import Image
import json
import cv2
import numpy as np

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class ChannelAttention(nn.Module):
    def __init__(self, in_channels, reduction_ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction_ratio),
            nn.ReLU(),
            nn.Linear(in_channels // reduction_ratio, in_channels)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc(self.avg_pool(x).view(x.size(0), -1))
        max_out = self.fc(self.max_pool(x).view(x.size(0), -1))
        out = avg_out + max_out
        return self.sigmoid(out).unsqueeze(2).unsqueeze(3)

class SpatialAttention(nn.Module):
    def __init__(self):
        super(SpatialAttention, self).__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size=7, padding=3, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv(x)
        return self.sigmoid(x)

class CBAMModule(nn.Module):
    def __init__(self, in_channels):
        super(CBAMModule, self).__init__()
        self.channel_attention = ChannelAttention(in_channels)
        self.spatial_attention = SpatialAttention()

    def forward(self, x):
        x_out = x * self.channel_attention(x)
        x_out = x_out * self.spatial_attention(x_out)
        return x_out

class ResNetWithCBAM(nn.Module):
    def __init__(self, block, layers, num_classes=1000, pretrained=True):
        super(ResNetWithCBAM, self).__init__()
        self.in_channels = 64
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)
        self.cbam1 = CBAMModule(256)
        self.cbam2 = CBAMModule(512)
        self.cbam3 = CBAMModule(1024)
        self.cbam4 = CBAMModule(2048)
        
    def save_weights(self, filepath):
        torch.save(self.state_dict(), filepath)

    @classmethod
    def load_weights(cls, filepath, block, layers, num_classes=1000):
        model = cls(block, layers, num_classes=num_classes, pretrained=False)
        model.load_state_dict(torch.load(filepath))
        model.eval()
        return model
        
    def apply_pretrained(self):
        pretrained_resnet = resnet50(pretrained=True)
        self.load_state_dict(pretrained_resnet.state_dict(), strict=False)

    def _make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if (stride != 1) or (self.in_channels != out_channels * block.expansion):
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channels, out_channels * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels * block.expansion),
            )
        layers = []
        layers.append(block(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.in_channels, out_channels))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        
        # print('x shape', x.shape)
        
        # print('x layer1 shape', self.layer1(x).shape)
        x = self.cbam1(self.layer1(x))
        # print('x cbam1 shape', x.shape)
        
        # print('x layer2 shape', self.layer2(x).shape)
        x = self.cbam2(self.layer2(x))
        # print('x cbam2 shape', x.shape)
        
        # print('x layer3 shape', self.layer3(x).shape)
        x = self.cbam3(self.layer3(x))
        # print('x cbam3 shape', x.shape)
        
        # print('x layer4 shape', self.layer4(x).shape)
        x = self.cbam4(self.layer4(x))
        # print('x cbam4 shape', x.shape)
        # x = self.layer1(x)
        # x = self.layer2(x)
        # x = self.layer3(x)
        # x = self.layer4(x)
        
        # Calculate attention map (channel-wise and spatial-wise)
        channel_attention = x.mean(dim=(2, 3), keepdim=True)  # Calculate channel-wise attention
        spatial_attention = x.sum(dim=1, keepdim=True)  # Calculate spatial-wise attention

        # Combine channel-wise and spatial-wise attention
        attention_map = channel_attention * spatial_attention

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x, attention_map

AI/model/cbam_resnet.py
icanSeeo/i-can-see
import torch
from utils import *
from model import *
from torchvision.datasets import CIFAR100
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from tqdm import tqdm
import torch.nn as nn
import matplotlib.pyplot as plt

# Define functions for inference and interpretation


def perform_inference(model, input_batch):
    with torch.no_grad():
        output, attention_map = model(input_batch)
    return output, attention_map


def interpret_output(output, class_names):
    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    top5_prob, top5_catid = torch.topk(probabilities, 5)

    for i in range(top5_prob.size(0)):
        print(class_names[top5_catid[i]], top5_prob[i].item())


# Define transformations and load CIFAR-100 dataset
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

model = ResNetWithCBAM(Bottleneck, [3, 4, 6, 3], num_classes=1000)
model.apply_pretrained()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

num_epochs = 200
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()

# Load the CIFAR-100 dataset
train_dataset = CIFAR100(root='./data', train=True,
                         download=True, transform=transform)

# Define a DataLoader
batch_size = 200
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# Initialize lists to store accuracy and loss
accuracy_list = []
loss_list = []

for epoch in tqdm(range(num_epochs)):
    model.train()
    running_loss = 0.0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs, aux_output = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}")

    # Calculate accuracy after each epoch
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            outputs, aux_output = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    accuracy_list.append(accuracy)
    loss_list.append(running_loss / len(train_loader))

    print(
        f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}, Accuracy: {accuracy}%")

# Plot the accuracy graph
plt.plot(range(1, num_epochs+1), accuracy_list, label='Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy (%)')
plt.title('Accuracy over Epochs')
plt.legend()
plt.show()

# Plot the loss graph
plt.plot(range(1, num_epochs+1), loss_list, label='Loss', color='red')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Loss over Epochs')
plt.legend()
plt.show()

# Save the model weights
model.save_weights('cbam-finetunned-200-notpretrained(200batch).pth')

AI/pre-train.py
icanSeeo/i-can-see
from .preprocess import *
from .visualize import *
AI/utils/__init__.py
icanSeeo/i-can-see
import cv2
import numpy as np

def shake_image(image_path, max_shift=10):
    # Load the image
    image = cv2.imread(image_path)
    
    # Generate random shift values
    shift_x = np.random.randint(-max_shift, max_shift)
    shift_y = np.random.randint(-max_shift, max_shift)
    
    # Define the transformation matrix for shifting
    M = np.float32([[1, 0, shift_x], [0, 1, shift_y]])
    
    # Apply the transformation
    shifted_image = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))
    
    return shifted_image

if __name__ == "__main__":
    # Example usage
    image_path = 'utils/curby.jpg'
    max_shift = 50

    shake_image(image_path, max_shift)

AI/utils/noise_generators.py
icanSeeo/i-can-see
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.models import resnet50
from PIL import Image
import json
import cv2
import numpy as np

# 이미지를 전처리하는 함수 정의
def preprocess_image(image_path):
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    input_image = Image.open(image_path).convert('RGB')
    input_tensor = preprocess(input_image)
    input_batch = input_tensor.unsqueeze(0)

    return input_batch

# 클래스 인덱스를 클래스명으로 변환하는 함수 정의
def load_class_names():
    with open('imagenet_classes.json') as f:
        class_names = json.load(f)
    return class_names

# GPU를 사용할 수 있는 경우 GPU로 모델 및 입력 데이터를 이동하는 함수 정의
def prepare_model_and_input(model, input_batch):
    if torch.cuda.is_available():
        model.to('cuda')
        input_batch = input_batch.to('cuda')
    return model, input_batch

AI/utils/preprocess.py
icanSeeo/i-can-see
import cv2

def overlay_attention_map(original_image, attention_map, output_path, alpha=0.5):
    attention_map = (attention_map - attention_map.min()) / (attention_map.max() - attention_map.min()) * 255
    
    # Resize attention map to match the original image size
    attention_map = cv2.resize(attention_map, (original_image.shape[1], original_image.shape[0]))

    # Apply colormap to attention map for better visualization (optional)
    attention_colormap = cv2.applyColorMap(attention_map, cv2.COLORMAP_JET)

    # Blend attention map with the original image
    overlaid_image = cv2.addWeighted(original_image, 1-alpha, attention_colormap, alpha, 0)

    # Save the overlaid image
    cv2.imwrite(output_path, overlaid_image)

AI/utils/visualize.py
icanSeeo/i-can-see