File size: 6,184 Bytes
524b7e1 3f48fe9 524b7e1 a1a7b41 524b7e1 a1a7b41 524b7e1 440da7b 524b7e1 a1a7b41 524b7e1 440da7b 524b7e1 34b6cc1 524b7e1 34b6cc1 524b7e1 800efd9 524b7e1 |
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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
from tqdm import tqdm
#from einops import rearrange
from PIL import Image
from copy import deepcopy
from typing import List, Optional, Union
from torch import autocast
#from torchvision import utils as vutils
from utils.util import EditingJsonDataset, EditingSingleImageDataset, plot_images
from lr_schedule import WarmupLinearLRSchedule
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.tensorboard import SummaryWriter
from models.model import RGN
from models.utils import visualize_images, read_image_from_url, draw_image_with_bbox_new, Bbox
from utils.util2 import compose_text_with_templates, get_augmentations_template
#from torchvision.utils import draw_bounding_boxes
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from torchvision import datasets, transforms
from engine import *
from vis import *
import os, jax, cv2, pdb
import numpy as np
import argparse, torch, inspect
import PIL, time, json, datetime
import random
import torch.nn as nn
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import utils.misc as misc
import torchvision.transforms as T
import torch.distributed as dist
from tqdm import tqdm
from einops import rearrange
from PIL import Image
from copy import deepcopy
from typing import List, Optional, Union
from torch import autocast
#from torchvision import utils as vutils
from utils.util import build_dataset, plot_images
from lr_schedule import WarmupLinearLRSchedule
from torch.utils.tensorboard import SummaryWriter
from models.model import RGN
from models.utils import visualize_images, read_image_from_url, draw_image_with_bbox_new, Bbox
from utils.util2 import compose_text_with_templates, get_augmentations_template
from torchvision.utils import draw_bounding_boxes
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from torchvision import datasets, transforms
from engine import *
from utils.post_process import get_final_img
import random
import os, jax, cv2, pdb
import numpy as np
import argparse, torch, inspect
import PIL, time, json, datetime
import torch.nn as nn
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import utils.misc as misc
import torchvision.transforms as T
import torch.distributed as dist
from tqdm import tqdm
#from einops import rearrange
from PIL import Image
from copy import deepcopy
from typing import List, Optional, Union
from torch import autocast
from torchvision import utils as vutils
from utils.util import EditingJsonDataset, EditingSingleImageDataset, plot_images
from lr_schedule import WarmupLinearLRSchedule
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.tensorboard import SummaryWriter
from models.model import RGN
from models.utils import visualize_images, read_image_from_url, draw_image_with_bbox_new, Bbox
from utils.util2 import compose_text_with_templates, get_augmentations_template
from torchvision.utils import draw_bounding_boxes
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from torchvision import datasets, transforms
from engine import *
from vis import *
import os, jax, cv2, pdb
import numpy as np
import argparse, torch, inspect
import PIL, time, json, datetime
import random
import torch.nn as nn
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import utils.misc as misc
import torchvision.transforms as T
import torch.distributed as dist
def map_cooridates(bbox, min_num=0, max_num=255):
# input feat size: 32 x 32
min_num2, max_num2 = 0, 31
return (max_num-min_num)/(max_num2-min_num2) * \
(bbox-min_num2) + min_num
def get_mask_imgs(imgs, bboxs):
imgs = imgs.repeat_interleave(bboxs.shape[0]//imgs.shape[0], 0)
mask_imgs = torch.zeros(imgs.shape, dtype=torch.uint8)
for i in range(imgs.shape[0]):
mask_imgs[i][:, bboxs[i][1].int().item():bboxs[i][3].int().item(), \
bboxs[i][0].int().item():bboxs[i][2].int().item()] = 1
return imgs, mask_imgs.float()
def save_img(args, batch, results, bboxs, imgs, mask_imgs, editing_rompt):
transform = T.Resize(512)
for i in range(results.shape[0]):
img = (imgs[i]*255.0).to(dtype=torch.uint8)
bbox = bboxs[i].to(dtype=torch.uint8).unsqueeze(0)
draw_img = draw_bounding_boxes(img, bbox, width=3, colors=(255,255,0))
img_name = '-'.join(str(editing_rompt).split(' '))
ori_img_path = os.path.join(new_path, 'input_image.png')
if i == 0:
save_image(transform(imgs[i]), ori_img_path)
save_image(res[i], os.path.join(new_path2, str(batch) + '_' +str(img_name) + 'anchor'+ str(i)+'.png'))
if args.draw_box:
bbox = bboxs[i].to(dtype=torch.uint8).unsqueeze(0)
draw_img = draw_bounding_boxes(img, bbox, width=3, colors=(255,255,0))
draw_img_path = os.path.join(new_path3, str(batch) + '_' + str(img_name) + 'anchor' + str(i)+'_ori_draw.png')
save_image(transform((draw_img/255.0).float()), draw_img_path)
get_final_img(args, editing_rompt, ori_img_path, new_path2)
template = get_augmentations_template()
device_id = 'cuda:1'
model = RGN(image_size=args.image_size, device=device_id, args=args).to(device_id)
# 使用 OpenCV 读取图像 (BGR 格式)
image_cv = cv2.imread("images/1.png")
image_cv = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB) # 转换为 RGB
transform = transforms.Compose([
transforms.Resize((224, 224)), # 调整大小
transforms.ToTensor(), # 转换为 PyTorch 张量
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 归一化
])
imgs = transform(image_pil)
imgs = imgs.to(device=device_id, non_blocking=True)[0].unsqueeze(0)
e_prompt = "Put some birds in the sky and some flowers around the trees"
e_prompt = compose_text_with_templates(e_prompt, template)
bboxs = torch.ceil(map_cooridates(model.module.get_anchor_box(imgs)))
imgs = imgs.repeat_interleave(bboxs.shape[0]//imgs.shape[0], 0)
_, mask_imgs = get_mask_imgs(imgs, bboxs)
results = model.module.generate_result(imgs, mask_imgs.to(device_id), e_prompt)
results.save('ans.png')
|