Anonymous-123/ImageNet-Editing
1
1#!/usr/bin/python2#****************************************************************#3# ScriptName: analysis_data.py4# Author: Anonymous_1235# Create Date: 2022-07-25 19:546# Modify Author: Anonymous_1237# Modify Date: 2022-09-25 12:048# Function: 9#***************************************************************#10 11import os12import sys13import numpy as np14import cv215import torch16from tqdm import tqdm17import shutil18import pdb19 20import argparse21 22parser = argparse.ArgumentParser(description='resize object')23parser.add_argument('--scale', type=float, default=None, help='object scale')24parser.add_argument('--img_path', type=str, help='image path')25parser.add_argument('--mask_path', type=str, help='mask path')26 27 28def get_bbox_and_rate(mask):29 gray = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)30 ret, thresh = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)31 contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)32 if len(contours) == 0:33 return None, None34 max_area = 035 max_idx = 036 for i, cnt in enumerate(contours):37 x,y,w,h = cv2.boundingRect(cnt)38 if w*h > max_area:39 max_idx = i40 max_area = w*h41 # 外接矩形42 x,y,w,h = cv2.boundingRect(contours[max_idx])43 mask_new = np.zeros(mask.shape, dtype='uint8')44 mask_new[y:y+h, x:x+w, :] = mask[y:y+h, x:x+w, :]45 46 rate = (mask_new[:,:,0]>127.5).sum()/mask.shape[0]/mask.shape[1]47 48 return (x,y,w,h), rate49 50def resize_around_the_center(img, mask, bbox, operation, scale_step=1.2):51 x,y,w,h = bbox52 H,W,C = mask.shape53 obj_mask = mask[y:y+h, x:x+w, :].copy()54 # obj_mask = cv2.resize(obj_mask, (int(w*scale_step),int(h*scale_step)) if operation == 'upsample' else (int(w/scale_step), int(h/scale_step)))55 obj_mask = cv2.resize(obj_mask, (int(w*scale_step),int(h*scale_step)))56 start_point_x = max(x+w//2 - obj_mask.shape[1]//2, 0) # center - w57 start_point_y = max(y+h//2 - obj_mask.shape[0]//2, 0) # center - h58 end_point_x = min(x+w//2 + obj_mask.shape[1]//2, W) # center+w59 end_point_y = min(y+h//2 + obj_mask.shape[0]//2, H) # center+h60 61 start_point_x_obj = max(0,obj_mask.shape[1]//2-(x+w//2))62 start_point_y_obj = max(0, obj_mask.shape[0]//2-(y+h//2))63 mask[:] = 064 mask[start_point_y:end_point_y, start_point_x:end_point_x] = obj_mask[start_point_y_obj:start_point_y_obj+(end_point_y-start_point_y), start_point_x_obj:start_point_x_obj+(end_point_x-start_point_x)]65 66 obj_img = img[y:y+h, x:x+w, :].copy()67 # obj_img = cv2.resize(obj_img, (int(w*scale_step),int(h*scale_step)) if operation == 'upsample' else (int(w/scale_step), int(h/scale_step)))68 obj_img = cv2.resize(obj_img, (int(w*scale_step),int(h*scale_step)))69 img = cv2.GaussianBlur(img, (49, 49), 0)70 img[start_point_y:end_point_y, start_point_x:end_point_x] = obj_img[start_point_y_obj:start_point_y_obj+(end_point_y-start_point_y), start_point_x_obj:start_point_x_obj+(end_point_x-start_point_x)]71 72 return img, mask73 74def resize_around_the_center_padding(img, mask, bbox, scale_step=1.2):75 x,y,w,h = bbox76 H,W,C = mask.shape77 mask_new = np.zeros((int(H/scale_step), int(W/scale_step), 3), dtype='uint8')78 mask_new_full = np.zeros((int(H/scale_step), int(W/scale_step), 3), dtype='uint8')79 # img_new = np.zeros((int(H/scale_step), int(W/scale_step), 3), dtype='uint8')80 img_new = cv2.resize(img, (int(W/scale_step), int(H/scale_step)))81 82 if scale_step < 1:83 mask_new[int((y+h/2)*(1/scale_step-1)):int((y+h/2)*(1/scale_step-1)+H), int((x+w/2)*(1/scale_step-1)):int((x+w/2)*(1/scale_step-1)+W)] = mask84 mask_new_full[int((y+h/2)*(1/scale_step-1)):int((y+h/2)*(1/scale_step-1)+H), int((x+w/2)*(1/scale_step-1)):int((x+w/2)*(1/scale_step-1)+W)] = mask.max()*np.ones(mask.shape, dtype='uint8')85 86 img_new[int((y+h/2)*(1/scale_step-1)):int((y+h/2)*(1/scale_step-1)+H), int((x+w/2)*(1/scale_step-1)):int((x+w/2)*(1/scale_step-1)+W)] = img87 88 else:89 mask_new = mask[int((y+h/2)*(1-1/scale_step)):int((y+h/2)*(1-1/scale_step))+int(H/scale_step), int((x+w/2)*(1-1/scale_step)):int((x+w/2)*(1-1/scale_step))+int(W/scale_step)]90 mask_new_full = mask[int((y+h/2)*(1-1/scale_step)):int((y+h/2)*(1-1/scale_step))+int(H/scale_step), int((x+w/2)*(1-1/scale_step)):int((x+w/2)*(1-1/scale_step))+int(W/scale_step)]91 img_new = img[int((y+h/2)*(1-1/scale_step)):int((y+h/2)*(1-1/scale_step))+int(H/scale_step), int((x+w/2)*(1-1/scale_step)):int((x+w/2)*(1-1/scale_step))+int(W/scale_step)]92 93 img_new = cv2.resize(img_new, (W,H))94 mask_new = cv2.resize(mask_new, (W,H))95 mask_new_full = cv2.resize(mask_new_full, (W,H))96 97 return img_new, mask_new, mask_new_full98 99def rescale(img, mask, scale=None, max_steps=50):100 bbox, rate = get_bbox_and_rate(mask)101 if bbox is None:102 return None, None, None103 num_steps = 0104 mask_full = mask.copy()105 while np.floor(rate*100) != scale*100. and abs(rate-scale) > 0.015:106 # while not (abs(bbox[0]-0)<10 or abs(bbox[1]-0)<10 or abs(bbox[0]+bbox[2]-img.shape[1])<10 or abs(bbox[1]+bbox[3]-img.shape[0])<10):107 operation = 'upsample' if np.floor(rate*100) < scale*100. else 'downsample'108 scale_step = np.sqrt(scale/rate)109 # img, mask = resize_around_the_center(img, mask, bbox, operation, scale_step=scale_step)110 img, mask, mask_full = resize_around_the_center_padding(img, mask, bbox, scale_step=scale_step)111 bbox, rate_ = get_bbox_and_rate(mask)112 if (operation == 'upsample' and rate_ < rate) or (operation == 'downsample' and rate_ > rate):113 return None, None, None114 num_steps += 1115 rate = rate_116 print(rate)117 if num_steps > max_steps:118 return None, None, None119 return img, mask_full, mask120 121 122def rescale_maximum(img, mask, scale=None, max_steps=50):123 bbox, rate = get_bbox_and_rate(mask)124 if bbox is None:125 return None, None, None126 x,y,w,h = bbox127 H,W,C = img.shape128 if H/h < W/w:129 y_start, y_end = y, y+h130 new_w = w/H*h131 c_x = x + w//2132 c_x_new = new_w*c_x/W133 x_start = c_x - c_x_new134 x_end = x_start + new_w135 else:136 x_start, x_end = x, x+w137 new_h = h/W*w138 c_y = y+h//2139 c_y_new = new_h*c_y/H140 y_start = c_y - c_y_new141 y_end = y_start + new_h142 img_new = img[min(y, int(y_start)):max(int(y_end), y+h), min(x, int(x_start)):max(int(x_end),x+w), :]143 mask_new = mask[min(y, int(y_start)):max(int(y_end),y+h),min(x, int(x_start)):max(int(x_end),x+w),:]144 145 img_new = cv2.resize(img_new, (W,H))146 mask_new = cv2.resize(mask_new, (W,H))147 148 return img_new, mask_new, mask_new149 150 151if __name__ == '__main__':152 args = parser.parse_args()153 scale = args.scale154 img_path_save = 'results/img_rescaled.png'155 mask_path_save = 'results/mask_rescaled.png'156 if scale == None:157 shutil.copy(args.img_path, img_path_save)158 shutil.copy(args.mask_path, mask_path_save)159 else:160 try:161 finals = []162 img = cv2.imread(args.img_path)163 mask = cv2.imread(args.mask_path)164 165 img_rescale, mask_rescale, mask_obj = rescale_maximum(img.copy(), mask.copy(), scale=scale)166 bbox, max_rate = get_bbox_and_rate(mask_obj)167 if scale < max_rate:168 img_rescale, mask_rescale, mask_obj = rescale(img.copy(), mask.copy(), scale=scale)169 if img_rescale is None:170 print('Invalid size')171 shutil.copy(args.img_path, img_path_save)172 shutil.copy(args.mask_path, mask_path_save)173 sys.exit()174 final = [img, img_rescale, mask, mask_rescale, mask_obj]175 # cv2.imwrite('tmp.png', cv2.hconcat(final))176 177 cv2.imwrite(img_path_save, img_rescale)178 cv2.imwrite(mask_path_save, mask_obj)179 # cv2.imwrite(mask_path_save_full, mask_rescale)180 except:181 print('Invalid size, using the original one')182 shutil.copy(args.img_path, img_path_save)183 shutil.copy(args.mask_path, mask_path_save)184 185 186 187 188 189