File size: 4,738 Bytes
1c8e113 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | from PIL import Image
from tqdm import tqdm
import sys
import os
import json
import cv2
import torch
from torch.autograd import Variable
import torchvision
from torch.utils.data import DataLoader, Dataset
from torchvision.datasets import CocoCaptions
from torchvision.transforms import ToTensor
import torchvision.transforms as transforms
from torchvision.utils import save_image
import numpy as np
# from net_canny import Net
from ControlNet.annotator.canny import CannyDetector
from ControlNet.annotator.util import resize_image, HWC3
# from pytorch_msssim import ssim, ms_ssim, SSIM
class ResultFolderDataset(Dataset):
def __init__(self, data_dir, results_dir, n, transform=None):
self.data_dir = data_dir
self.results_dir = results_dir
self.n = n
self.image_paths = sorted([f for f in os.listdir(data_dir) if f.lower().endswith(('.png'))])
# self.image_paths = sorted([os.path.join(folder, f) for f in os.listdir(folder) if f.lower().endswith('_{}.png'.format(n))])
self.transform = transform
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
image_name = self.image_paths[idx]
source_path = os.path.join(self.data_dir, image_name)
base_name = image_name.split('_')[1].split('.')[0] # Extract 'x' from 'image_x.png'
image_name2 = f'result_{base_name}_{self.n}.png'
result_path = os.path.join(self.results_dir, image_name2)
source_image = Image.open(source_path) #.convert('RGB')
result_image = Image.open(result_path) #.convert('RGB')
if self.transform:
source_image = self.transform(source_image)
result_image = self.transform(result_image)
return source_image, result_image
def calculate_metrics(pred, target):
intersection = np.logical_and(pred, target)
union = np.logical_or(pred, target)
iou_score = np.sum(intersection) / np.sum(union)
accuracy = np.sum(pred == target) / target.size
tp = np.sum(intersection) # True positive
fp = np.sum(pred) - tp # False positive
fn = np.sum(target) - tp # False negative
f1_score = (2 * tp) / (2 * tp + fp + fn)
return iou_score, accuracy, f1_score
if __name__ == '__main__':
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
low_threshold = 100
high_threshold = 200
n = 0
# epoch = 10
# experiment = './log/image_log_oft_COCO_canny_eps_1-3_r_4_cayley_4gpu'
experiment = 'log/image_log_householder_gramschmidt_COCO_canny_eps_7e-06_pe_diff_mlp_l_8_8gpu_2024-05-19-21-22-24-466032'
if 'train_with_norm' in experiment:
epoch = 4
else:
if 'COCO' in experiment:
epoch = 10
else:
epoch = 19
data_dir = os.path.join(experiment, 'source', str(epoch))
result_dir = os.path.join(experiment, 'results', str(epoch))
json_file = os.path.join(experiment, 'results.json')
# Define the transforms to apply to the images
transform = transforms.Compose([
# transforms.Resize((512, 512)),
# transforms.CenterCrop(512),
transforms.ToTensor()
])
dataset = ResultFolderDataset(data_dir, result_dir, n=n, transform=transform)
data_loader = DataLoader(dataset, batch_size=1, shuffle=False)
apply_canny = CannyDetector()
loss = 0
iou_score_mean = 0
accuracy_mean = 0
f1_score_mean = 0
ssim_mean = 0
for i, data in tqdm(enumerate(data_loader), total=len(data_loader)):
source_image, result_image = data
# Convert the tensor to a numpy array and transpose it to have the channels last (H, W, C)
source_image_np = source_image.squeeze(0).permute(1, 2, 0).numpy()
result_image_np = result_image.squeeze(0).permute(1, 2, 0).numpy()
# Convert the image to 8-bit unsigned integers (0-255)
source_image_np = (source_image_np * 255).astype(np.uint8)
result_image_np = (result_image_np * 255).astype(np.uint8)
source_detected_map = apply_canny(source_image_np, low_threshold, high_threshold) / 255
result_detected_map = apply_canny(result_image_np, low_threshold, high_threshold) / 255
iou_score, accuracy, f1_score = calculate_metrics(result_detected_map, source_detected_map)
iou_score_mean = iou_score_mean + iou_score
accuracy_mean = accuracy_mean + accuracy
f1_score_mean = f1_score_mean + f1_score
iou_score_mean = iou_score_mean / len(dataset)
accuracy_mean = accuracy_mean / len(dataset)
f1_score_mean = f1_score_mean / len(dataset)
print(experiment)
print('[Canny]', '[IOU]:', iou_score_mean, '[F1 Score]:', f1_score_mean) |