def content_loss(img1_path, img2_path, content_layers=['r42']): img1 = ToTensor()(Image.open(img1_path)) img2 = ToTensor()(Image.open(img2_path)) if torch.cuda.is_available(): img1 = img1.cuda() img2 = img2.cuda() img1 = Variable(img1.unsqueeze(0)) img2 = Variable(img2.unsqueeze(0)) img1_fm = vgg(img1, content_layers)[0] img2_fm = vgg(img2, content_layers)[0] return nn.MSELoss()(img1_fm, img2_fm)
def main_worker(args, use_gpu=True): device = torch.device('cuda') if use_gpu else torch.device('cpu') # Model and version net = importlib.import_module('model.' + args.model) model = net.InpaintGenerator(args).cuda() model.load_state_dict(torch.load(args.pre_train, map_location='cuda')) model.eval() # prepare dataset image_paths = [] with open(os.path.join(args.dir_image, args.data_test, 'val.txt')) as f: images_list = f.read().splitlines() for path in images_list: image_paths.append(os.path.join(args.dir_image, args.data_test, path)) # image_paths.sort() mask_paths = glob(os.path.join(args.dir_mask, args.mask_type, '*.png')) os.makedirs(args.outputs, exist_ok=True) trans = transforms.Compose([ transforms.Resize((args.image_size, args.image_size), interpolation=transforms.InterpolationMode.NEAREST), ]) j = 0 # iteration through datasets for ipath, mpath in tqdm(zip(image_paths, mask_paths)): if j >= args.num_test: exit() image = ToTensor()(Image.open(ipath).convert('RGB')) image = trans(image) image = (image * 2.0 - 1.0).unsqueeze(0) mask = ToTensor()(Image.open(mpath).convert('L')) mask = trans(mask) mask = mask.unsqueeze(0) image, mask = image.cuda(), mask.cuda() image_masked = image * (1 - mask.float()) + mask with torch.no_grad(): pred_img = model(image_masked, mask) comp_imgs = (1 - mask) * image + mask * pred_img visualize_test(j, image, image * (1 - mask), pred_img.detach(), comp_imgs.detach()) j += 1
def yolo_detector(self, img): model = Darknet(self.config_path, img_size=self.img_size) model.load_weights(self.weights_path) if self.use_cuda: model.cuda() model.eval() #classes = load_classes(self.class_path) # Extracts class labels from file loader = transforms.Compose([transforms.Scale([self.img_size, self.img_size])]) image = loader(img) image = ToTensor()(image).unsqueeze(0) image = Variable(image) image=image.cuda() #image = image[None, :, :, :] #print ('dete') detections = model(image) detections = non_max_suppression(detections, 80, self.conf_thres, self.nms_thres) return detections
def np_to_tensor(nparray): x = image.astype('float') x.transpose((2, 0, 1)) x = ToTensor()(x) x -= 0.5 return x.cuda()
def resnet34(self, img): bbx = ToTensor()(img) bbx = Variable(bbx) if self.cuda: bbx = bbx.cuda() bbx = bbx.view(-1, bbx.size(0), bbx.size(1), bbx.size(2)) return bbx
def content_loss_images( img1_path, img2_path, content_layers=[ 'r42' ]): # takes the path of two images and determines their content loss img1 = ToTensor()(Image.open(img1_path)) img2 = ToTensor()(Image.open(img2_path)) if torch.cuda.is_available(): img1 = img1.cuda() img2 = img2.cuda() img1 = Variable(img1.unsqueeze(0)) img2 = Variable(img2.unsqueeze(0)) img1_fm = vgg(img1, content_layers)[0] img2_fm = vgg(img2, content_layers)[0] return nn.MSELoss()(img1_fm, img2_fm)
def resnet34(self, img): bbx = ToTensor()(img) bbx = Variable(bbx, volatile=True) if self.cuda: bbx = bbx.cuda() bbx = bbx.view(-1, bbx.size(0), bbx.size(1), bbx.size(2)) ret = self.Appearance(bbx) ret = ret.view(1, -1) return ret
def inference(model, args): image_folder = "/media/pandongwei/ExtremeSSD/work_relative/extract_img/2020.10.16_1/" video_save_path = "/home/pandongwei/work_repository/erfnet_pytorch/eval/" # parameters about saving video fourcc = cv2.VideoWriter_fourcc(*'XVID') out = cv2.VideoWriter(video_save_path + 'output_new.avi', fourcc, 10.0, (640, 480)) cuda = True model.eval() paths = [] for root, dirs, files in os.walk(image_folder, topdown=True): for file in files: image_path = os.path.join(image_folder, file) paths.append(image_path) paths.sort() font = cv2.FONT_HERSHEY_SIMPLEX angle_pre = 0 for i, path in enumerate(paths): start_time = time.time() image = cv2.imread(path) image = (image / 255.).astype(np.float32) image = ToTensor()(image).unsqueeze(0) if (cuda): image = image.cuda() input = Variable(image) with torch.no_grad(): output = model(input) label = output[0].max(0)[1].byte().cpu().data # label_cityscapes = cityscapes_trainIds2labelIds(label.unsqueeze(0)) label_color = Colorize()(label.unsqueeze(0)) label_save = label_color.numpy() label_save = label_save.transpose(1, 2, 0) # 加上路径规划 label_save, angle = perception_to_angle(label_save, angle_pre) # 加一个滤波以防止角度突然跳变 TODO if abs(angle - angle_pre) > 10: angle = angle_pre angle_pre = angle # label_save.save(filenameSave) image = image.cpu().numpy().squeeze(axis=0).transpose(1, 2, 0) image = (image * 255).astype(np.uint8) output = cv2.addWeighted(image, 0.5, label_save, 0.5, 0) cv2.putText(output,str(round(angle,3)),(50,50),cv2.FONT_HERSHEY_SIMPLEX,2,(0,0,0),2) #output = np.hstack([label_save, image]) out.write(output) print(i, " time: %.2f s" % (time.time() - start_time)) out.release()
def to_tensor(x, path): """Converts an object (image or tensor or list) to a Pytorch Variable""" path_to_x = path + str(x) img = Image.open(path_to_x) img = ToTensor()(img).unsqueeze(0) img = img.squeeze(1) if torch.cuda.is_available(): img = img.cuda() return img
def callback(self, oimg): try: #if you want to save images and labels ,please uncomment following codes(No.1 to No.4). #NO.1 write_image_name = "image_" + str(self.count) + ".jpg" #No.2 write_label_name = "label_" + str(self.count) + ".jpg" oimg_b = bytes(oimg.data) np_arr = np.fromstring(oimg_b, np.uint8) img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) #No.3 #cv2.imwrite("/home/amsl/images/output/seg_pub_2/image/" + write_image_name, img) image = PIL_Image.fromarray(img) #image = image.crop((0, 120, 640, 480)) image = image.crop((0, 10, 640, 330)) #image = image.resize((1024,512),PIL_Image.NEAREST) image = image.resize((1024,512),PIL_Image.NEAREST) #img_size = image.shape image = ToTensor()(image) image = torch.Tensor(np.array([image.numpy()])) image = image.cuda() input_image = Variable(image) with torch.no_grad(): output_image = self.model(input_image) label = output_image[0].max(0)[1].byte().cpu().data label_color = Colorize()(label.unsqueeze(0)) label_pub = ToPILImage()(label_color) #label_pub = label_pub.resize((1024, 512),PIL_Image.NEAREST) #label_pub = label_pub.resize((1024, 512),PIL_Image.LANCZOS) label_pub = np.asarray(label_pub) #show label. #plt.imshow(label_pub) #plt.pause(0.001) #No.4 #cv2.imwrite("/home/amsl/images/output/seg_pub_2/label/" + write_label_name, label_pub) self.pub_seg.publish(self.bridge.cv2_to_imgmsg(label_pub, "bgr8")) print("published") self.count += 1 except CvBridgeError as e: print(e)
def view_filters(net, img): img = ToTensor()(img) img = Variable(img.unsqueeze(0), volatile=True) img = img.cuda() output = net.conv1(img) output = output.cpu().data.numpy()[0] fig, axes = plt.subplots(1, len(output)) for i in range(len(output)): axes[i].imshow(output[i]) axes[i].set_xticks([]) axes[i].set_yticks([]) plt.show()
def view_filters(net, img): with torch.no_grad(): img = ToTensor()(img).unsqueeze(0) img = img.cuda() output = net.conv1(img) output = output.cpu().data.numpy()[0] fig, axes = plt.subplots(1, len(output)) for i in range(len(output)): axes[i].imshow(output[i]) axes[i].set_xticks([]) axes[i].set_yticks([]) plt.show()
def load_model(model_path, data, cuda, iter = 1): from model_structure.unet_standard import NestedUNet use_padding = False unet = NestedUNet() unet = unet.eval() if cuda: unet = unet.cuda() if not cuda: unet.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage)) else: unet.load_state_dict(torch.load(model_path)) ori_tensor = ToTensor()(data) if cuda: ori_tensor = ori_tensor.cuda() ori_tensor = torch.unsqueeze(ori_tensor,0) padding_left = 0 padding_right = 0 padding_top = 0 padding_bottom = 0 ori_height = ori_tensor.size()[2] ori_width = ori_tensor.size()[3] if ori_height % 4: padding_top = (4 - ori_height % 4)//2 padding_bottom = 4 - ori_height % 4 - padding_top use_padding = True if ori_width % 4: padding_left = (4 - ori_width % 4)//2 padding_right = 4 - ori_width % 4 - padding_left use_padding = True if use_padding: padding_transform = torch.nn.ConstantPad2d((padding_left, padding_right, padding_top, padding_bottom), 0) ori_tensor = padding_transform(ori_tensor) output = ori_tensor with torch.no_grad(): for _ in range(iter): output = unet(output) if use_padding: output = output[:,:,padding_top : (padding_top + ori_height), padding_left : (padding_left + ori_width)] result = (output.data).cpu().numpy() result = result[0,0,:,:] return result
def main_worker(args, use_gpu=True): device = torch.device('cuda') if use_gpu else torch.device('cpu') # Model and version net = importlib.import_module('model.'+args.model) model = net.InpaintGenerator(args).cuda() model.load_state_dict(torch.load(args.pre_train, map_location='cuda')) model.eval() # prepare dataset image_paths = [] for ext in ['.jpg', '.png']: image_paths.extend(glob(os.path.join(args.dir_image, '*'+ext))) image_paths.sort() mask_paths = sorted(glob(os.path.join(args.dir_mask, '*.png'))) os.makedirs(args.outputs, exist_ok=True) # iteration through datasets for ipath, mpath in zip(image_paths, mask_paths): image = ToTensor()(Image.open(ipath).convert('RGB')) image = (image * 2.0 - 1.0).unsqueeze(0) mask = ToTensor()(Image.open(mpath).convert('L')) mask = mask.unsqueeze(0) image, mask = image.cuda(), mask.cuda() image_masked = image * (1 - mask.float()) + mask with torch.no_grad(): pred_img = model(image_masked, mask) comp_imgs = (1 - mask) * image + mask * pred_img image_name = os.path.basename(ipath).split('.')[0] postprocess(image_masked[0]).save(os.path.join(args.outputs, f'{image_name}_masked.png')) postprocess(pred_img[0]).save(os.path.join(args.outputs, f'{image_name}_pred.png')) postprocess(comp_imgs[0]).save(os.path.join(args.outputs, f'{image_name}_comp.png')) print(f'saving to {os.path.join(args.outputs, image_name)}')
def gesture_recognize(self, hand): hand[:, 0] /= 640 hand[:, 1] /= 480 hand = ToTensor()(hand) if torch.cuda.is_available(): hand = hand.cuda() hand = hand.view(1, -1) out = self.gesture_model(hand) out = F.softmax(out, 1) value, index = torch.max(out, 1) if value.item() > self.gesture_threshold: print(self.idx_to_gesture[index.item()], value.item()) return self.idx_to_gesture[index.item()] else: return None
def main(): print '!' parser = argparse.ArgumentParser(description='PyTorch LapSRN') parser.add_argument('--test_folder', type=str, default=PACKAGE_DIR + '/dataset/', help='input image to use') # TODO: fix parser.add_argument('--model', type=str, default='model/best_rgb.pth', help='model file to use') parser.add_argument('--result_folder', type=str, default=PACKAGE_DIR + '/results/', help='where to save the output image') #TODO: fix opt = parser.parse_args() images_list = glob(os.path.join(opt.test_folder, '*.jpg')) print 'images_list', len(images_list) model = torch.load(opt.model) print 'Loaded' model = model.cuda() for image_path in images_list: img_name = image_path.split('/')[-1].split('.')[0] img = Image.open(image_path) test_img = img.resize((img.size[0] / 2, img.size[1] / 2), PIL.Image.BICUBIC) print img.size, 'after', test_img tensor_img = ToTensor()(test_img).unsqueeze(0) print tensor_img.size() input = Variable(tensor_img.cuda()) HR_2, HR_4, HR_8 = model(input) HR_2 = HR_2.cpu() HR_4 = HR_4.cpu() HR_8 = HR_8.cpu() test = HR_2.data.squeeze(0) img = ToPILImage()(test) img.save(os.path.join(opt.result_folder, os.path.basename(image_path)))
def sr(im, scale): im = im.convert('YCbCr') im, cb, cr = im.split() h, w = im.size im = ToTensor()(im) im = Variable(im).view(1, -1, w, h) im = im.cuda() with torch.no_grad(): im = espcn(im) im = torch.clamp(im, 0., 1.) im = im.cpu() im = im.data[0] im = ToPILImage()(im) cb = cb.resize(im.size, Image.BICUBIC) cr = cr.resize(im.size, Image.BICUBIC) im = Image.merge('YCbCr', [im, cb, cr]) im = im.convert('RGB') return im
def inference(model, image, angle_pre): # pre-process image = (image / 255.).astype(np.float32) image = ToTensor()(image).unsqueeze(0) image = image.cuda() input = Variable(image) # inference with torch.no_grad(): output = model(input) # post-process label = output[0].max(0)[1].byte().cpu().data label_color = Colorize()(label.unsqueeze(0)) label_save = label_color.numpy() label_save = label_save.transpose(1, 2, 0) # 加上路径规划 label_save, angle = perception_to_angle(label_save, angle_pre) return angle
def Test(MODEL_NAME, UPSCALE_FACTOR, is_save=False, IMAGE_DIR=r'data\testing_lr_images', TEST_MODE=True): if type(MODEL_NAME) is EDSR or type(MODEL_NAME) is WDSR or type( MODEL_NAME) is SRResnet: model = MODEL_NAME else: model = EDSR(UPSCALE_FACTOR).eval() if TEST_MODE: model.cuda() model.load_state_dict(torch.load('epochs/' + MODEL_NAME)) else: model.load_state_dict( torch.load('epochs/' + MODEL_NAME, map_location=lambda storage, loc: storage)) print('\n----------------------------------------------------------') imgs = [] with torch.no_grad(): model.eval() for image_name in glob.glob(os.path.join(IMAGE_DIR, '*.*')): image = Image.open(image_name) image = ToTensor()(image).unsqueeze(0) if TEST_MODE: image = image.cuda() start = time.time() out = model(image) elapsed = (time.time() - start) out_img = ToPILImage()(torch.clip(out[0], 0, 1)) if is_save: out_img.save( f'data/testing_sr_images/{os.path.basename(image_name)}') sr_img = gpu_to_numpy(out[0], is_squeeze=False) imgs.append(sr_img) plot_hr_lr(sr_img, image) print('cost time: ' + str(elapsed) + 's') return imgs
def test_img(model_geoMat, model_freiburgForest, cfg): img_path = "/home/pandongwei/work_repository/erfnet_pytorch/test_img/row.png" cuda = cfg['cuda'] model_geoMat.eval() model_freiburgForest.eval() coef = np.ones([512, 1024, 5]).astype(np.float32) traversability = (0, 1.0, 0.6, 0.8, -1) for i in range(5): coef[:, :, i] = coef[:, :, i] * traversability[i] image = cv2.imread(img_path).astype(np.float32) image = cv2.resize(image, (1024, 512), interpolation=cv2.INTER_LINEAR) image = image / 255. image = ToTensor()(image).unsqueeze(0) if (cuda): image = image.cuda() input = Variable(image) # start_time = time.time() with torch.no_grad(): outputs_1 = model_freiburgForest(input) outputs_2 = model_geoMat(input) outputs_2 = outputs_2[0] # TODO #print("model inference time: %.2f s" % (time.time() - start_time)) # start_time = time.time() outputs_1 = outputs_1[0].max(0)[1].byte().cpu().data.unsqueeze(0) outputs_1 = outputs_1.numpy().transpose(1, 2, 0).astype(np.int64) outputs_1 = np.take(coef, outputs_1) outputs_2 = outputs_2[0, 0, :, :].cpu().data.unsqueeze(0) # print(outputs_2.shape) outputs_2 = outputs_2.numpy().transpose(1, 2, 0) outputs_combine = (outputs_1 * outputs_2 * 255).astype(np.int16) cv2.imwrite('test.png', outputs_combine)
def main(): CUDA = torch.cuda.is_available() opts = opt() pose_model = model.poseprosalnet(mpii_dataset.KEYPOINT_NAMES, mpii_dataset.EDGES, opts.local_grid_size, opts.insize) pose_model.load_state_dict(torch.load(opts.checkpoint)) pil_image = Image.open(opts.img) pil_image = Resize(opts.insize)(pil_image) img = ToTensor()(pil_image) img = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img) img = img.unsqueeze(0) pose_model.eval() print(mpii_dataset.KEYPOINT_NAMES) with torch.no_grad(): if CUDA: pose_model.cuda() img = img.cuda() pre = pose_model(img) humans = utils.get_humans_by_feature( pre, opts.insize, pose_model.outsize, opts.local_grid_size, ) # print(humans[0]) # print(humans[1]) pil_image = utils.draw_humans(mpii_dataset.KEYPOINT_NAMES, mpii_dataset.EDGES, pil_image, humans, visbbox=False) pil_image.save("result.png", "PNG")
def view_filters(net, img): """ Affiche le résultat des filtres à convolution sur une image. On suppose que le réseau a une couche `conv1` qui peut être appliqué sur la batch d'images passée en paramètre. Args: network (nn.Module): Un réseau de neurones PyTorch avec une couche `conv1`. img (Union[PILImage, torch.Tensor]): Une image. """ with torch.no_grad(): if not torch.is_tensor(img): img = ToTensor()(img) img = img.unsqueeze(0) img = img.cuda() output = net.conv1(img) output = output.cpu().numpy()[0] _, axes = plt.subplots(1, len(output)) for i, out in enumerate(output): axes[i].imshow(out) axes[i].set_xticks([]) axes[i].set_yticks([]) plt.show()
class_names = idx2cls(class_txt) img_paths = glob(f"{img_dir}/*.jpg") print(len(img_paths)) for img_path in img_paths: img = cv2.imread(img_path) img_cv = cv2.resize(img, dsize=(config["image_size"], config["image_size"])) imshow("i", img, 1) img_cv = bright_adjust(img_cv) img_pil = cv2pil(img_cv) img_tensor = ToTensor()(img_pil).unsqueeze(0) img_tensor = (img_tensor-0.5) / 0.5 if cuda: img_tensor = img_tensor.cuda() t = time.time() classifications, bbox_regression, landmark_regression = net(img_tensor) if priors == None: feature_map_sizes = [i.shape[1:3] for i in bbox_regression] priors = [get_prior(feature_map_size, config["image_size"], min_size, max_size, ratio)for feature_map_size, min_size, max_size, ratio in zip(feature_map_sizes, config["min_sizes"], config["max_sizes"], config["ratios"])] priors = torch.cat(priors, dim=0) if cuda: priors = priors.cuda() bbox_regression = torch.cat([i.reshape(-1, 4) for i in bbox_regression], dim=0) landmark_regression = torch.cat([i.reshape(-1, 8) for i in landmark_regression], dim=0) classifications = torch.cat([i.reshape(-1) for i in classifications], dim=0) classifications = classifications.reshape(landmark_regression.shape[0], -1)
dimensions=2, in_channels=4, out_channels=10, channels=(64, 128, 256, 512, 1024), strides=(2, 2, 2, 2), num_res_units=2, ) model.load_state_dict(torch.load(model_ckpt)) model.eval() model.cuda() # ---------------------------------------------------------------------------------------------- for imgName in imagesList: imgPath = os.path.join(dataPath, imgName) image = np.array(Image.open(imgPath), dtype=np.float32) img_clone = image.copy() image = image / 255.0 # norm to 0 -1 image = ToTensor()(image) image = image.unsqueeze(0) predictions = model(image.cuda()) predictions = nn.Softmax(dim=1)(predictions) predictions = predictions.squeeze(0).cpu().detach().numpy() result = np.argmax(predictions, axis=0) + 1 result_labelFormat = Image.fromarray(np.uint8(result)) (imgNameWE, ext) = os.path.splitext(imgName) imgSaveName = imgNameWE + ".png" saveNamePath = os.path.join(savePath, imgSaveName) result_labelFormat.save(saveNamePath) print(imgName)
def save_video_and_path_planning(model_geoMat, model_freiburgForest, cfg): image_folder = cfg["image_folder_raw"] video_save_path = "/home/pandongwei/work_repository/erfnet_pytorch/eval/" # parameters about saving video fourcc = cv2.VideoWriter_fourcc(*'XVID') out = cv2.VideoWriter(video_save_path + 'output.avi', fourcc, 10.0, (640, 480)) cuda = cfg['cuda'] model_geoMat.eval() model_freiburgForest.eval() coef = np.ones([512, 1024, 5]).astype(np.float32) traversability = (0, 1.0, 0.6, 0.8, -1) for i in range(5): coef[:, :, i] = coef[:, :, i] * traversability[i] paths = [] for root, dirs, files in os.walk(image_folder, topdown=True): for file in files: image_path = os.path.join(image_folder, file) paths.append(image_path) paths.sort() font = cv2.FONT_HERSHEY_SIMPLEX for i, path in enumerate(paths): start_time = time.time() image = cv2.imread(path) image = cv2.resize(image, (512, 256), interpolation=cv2.INTER_LINEAR) image = image / 255. image = ToTensor()(image).unsqueeze(0) if (cuda): image = image.cuda() input = Variable(image) with torch.no_grad(): outputs_1 = model_freiburgForest(input) outputs_2 = model_geoMat(input) outputs_2 = outputs_2[0] # TODO outputs_1 = outputs_1[0].max(0)[1].byte().cpu().data.unsqueeze(0) # print(outputs_1.shape) outputs_1 = outputs_1.numpy().transpose(1, 2, 0).astype(np.int64) # print(outputs_1.shape) outputs_1 = np.take(coef, outputs_1) outputs_2 = outputs_2[0, 0, :, :].cpu().data.unsqueeze(0) # print(outputs_2.shape) outputs_2 = outputs_2.numpy().transpose(1, 2, 0) # 加上momenton来稳定算法 TODO if i == 0: momenton = 0.9 output_pre = outputs_2.copy() else: outputs_2 = (momenton * output_pre + (1 - momenton) * outputs_2) output_pre = outputs_2.copy() outputs_combine = (outputs_1 * outputs_2 * 255).astype(np.uint8) outputs_combine_color = Colorize()(outputs_combine).transpose( 1, 2, 0).astype(np.uint8) outputs_combine_color = cv2.cvtColor(outputs_combine_color, cv2.COLOR_RGB2BGR) image = image.cpu().numpy() image = (image[0].transpose(1, 2, 0) * 255).astype(np.uint8) image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) # print(outputs_combine.shape) # output = np.hstack([outputs_combine_color, image]) # print(" post process time: %.2f s" % (time.time() - start_time)) # start_time = time.time() outputs_combine = outputs_combine[:, :, 0] output = path_planning(outputs_combine, outputs_combine_color) # print(" path_planning time: %.2f s" % (time.time() - start_time)) # start_time = time.time() output = np.hstack([output, image]) cv2.putText(output, str(round(1 / (time.time() - start_time), 1)) + " Hz", (50, 50), font, 1, (0, 0, 255), 2) out.write(image) print(i, " time: %.2f s" % (time.time() - start_time)) out.release()
d = collections.OrderedDict() checkpoint = torch.load(model_path) for key, value in checkpoint.items(): tmp = key[7:] d[tmp] = value ssd.load_state_dict(d) else: print(f"{model_path} 不存在") detect = Detect(config) img = Image.open(img_path) img = img.resize((300, 300)) img = ToTensor()(img).unsqueeze(0) if torch.cuda.is_available() and config["cuda"]: img = img.cuda() ssd = ssd.cuda() with torch.no_grad(): predictions = ssd(img) obj_confes, obj_clses, obj_boxes = detect(predictions, threshold) obj_boxes *= 300 img_cv = cv2.imread(img_path) h, w = img_cv.shape[0:2] h /= 300 w /= 300 keep_index = nms(obj_confes, obj_boxes, iou_threshold) for idx in keep_index: box = obj_boxes[idx]
default='netG_epoch_4_10.pth', type=str, help='generator model epoch name') opt = parser.parse_args() UPSCALE_FACTOR = opt.upscale_factor TEST_MODE = True if opt.test_mode == 'GPU' else False IMAGE_NAME = opt.image_name MODEL_NAME = opt.model_name model = Generator(UPSCALE_FACTOR).eval() if TEST_MODE: model.cuda() model.load_state_dict(torch.load('epochs/' + MODEL_NAME)) else: model.load_state_dict( torch.load('epochs/' + MODEL_NAME, map_location=lambda storage, loc: storage)) image = Image.open(IMAGE_NAME) with torch.no_grad(): image = ToTensor()(image).unsqueeze(0) if TEST_MODE: image = image.cuda() start = time.clock() out = model(image) elapsed = (time.clock() - start) print('cost' + str(elapsed) + 's') out_img = ToPILImage()(out[0].data.cpu()) out_img.save('out_srf_' + str(UPSCALE_FACTOR) + '_' + IMAGE_NAME)
parser.add_argument("--test", type=str, help="path to load test images") opt = parser.parse_args() print(opt) net = Net(opt.rb) net.load_state_dict(torch.load(opt.checkpoint)['state_dict']) net.eval() net = nn.DataParallel(net, device_ids=[0, 1, 2, 3]).cuda() print(net) images = utils.load_all_image(opt.test) for im_path in tqdm(images): filename = im_path.split('/')[-1] print(filename) im = Image.open(im_path) h, w = im.size print(h, w) im = ToTensor()(im) im = Variable(im).view(1, -1, w, h) im = im.cuda() with torch.no_grad(): im = net(im) im = torch.clamp(im, 0., 1.) im = im.cpu() im = im.data[0] im = ToPILImage()(im) im.save('output/%s' % filename)
model.eval() with torch.no_grad(): in_files = sorted(os.listdir(opt['test_dir'])) for file in tqdm(in_files, ncols=80): img = Image.open(os.path.join( opt['test_dir'], file)).convert('RGB').convert('YCbCr') img = np.array(img) # HWC # split to YCbCr ndarray y, cb, cr = img[:, :, 0], img[:, :, 1], img[:, :, 2] # only Y channel goes through the network y = ToTensor()(Image.fromarray(y)) # 0-1, CHW y = y.unsqueeze(0) # add a batch dimension y = y.cuda() # get Y channel output and convert to ndarray CHW out_y = model(y) out_y = out_y.detach().cpu().clamp( 0.0, 1.0).numpy().squeeze() # HW, 0-1 out_y = np.array(out_y * 255.0, dtype=np.uint8) # HW, 0-255 # Cb and Cr channels are upsampled by BICUBIC interplotation directly out_cb = Image.fromarray(cb).resize( (out_y.shape[1], out_y.shape[0]), resample=Image.BICUBIC) out_cb = np.array(out_cb, dtype=np.uint8) # HW, 0-255 out_cr = Image.fromarray(cr).resize( (out_y.shape[1], out_y.shape[0]), resample=Image.BICUBIC) out_cr = np.array(out_cr, dtype=np.uint8) # HW, 0-255
def load_model(model_path, data, cuda, iter=1): if os.path.basename(model_path) == "Gen1-noNoiseNoBackgroundSuperresolution.pth" or \ os.path.basename(model_path) == "Gen1-circularMask.pth" or \ os.path.basename(model_path) == "Gen1-gaussianMask.pth" or \ os.path.basename(model_path) == "Gen1-noBackgroundNonoise.pth" or \ os.path.basename(model_path) == "Gen1-noNoise.pth": from mypackage.model.unet_standard import NestedUNet net = NestedUNet() if cuda: net = net.cuda() if cuda: net = torch.nn.DataParallel(net) net.load_state_dict(torch.load(model_path)) else: net.load_state_dict({ k.replace('module.', ''): v for k, v in torch.load(model_path).items() }) transform = ToTensor() ori_tensor = transform(data) ori_tensor = torch.unsqueeze(ori_tensor, 0) padding_left = 0 padding_right = 0 padding_top = 0 padding_bottom = 0 ori_height = ori_tensor.size()[2] ori_width = ori_tensor.size()[3] use_padding = False if ori_width >= ori_height: padsize = ori_width else: padsize = ori_height if np.log2(padsize) >= 7.0: if (np.log2(padsize) - 7) % 1 > 0: padsize = 2**(np.log2(padsize) // 1 + 1) else: padsize = 2**7 if ori_height < padsize: padding_top = int(padsize - ori_height) // 2 padding_bottom = int(padsize - ori_height - padding_top) use_padding = True if ori_width < padsize: padding_left = int(padsize - ori_width) // 2 padding_right = int(padsize - ori_width - padding_left) use_padding = True if use_padding: padding_transform = torch.nn.ConstantPad2d((padding_left, \ padding_right, \ padding_top, \ padding_bottom), 0) ori_tensor = padding_transform(ori_tensor) if cuda: ori_tensor = ori_tensor.cuda() output = net(ori_tensor) if cuda: result = (output.data).cpu().numpy() else: result = (output.data).numpy() padsize = int(padsize) result = result[0, 0, padding_top:(padsize - padding_bottom), padding_left:(padsize - padding_right)] else: from model_structure.unet_sigmoid import UNet use_padding = False unet = UNet() unet = unet.eval() if cuda: unet = unet.cuda() if not cuda: unet.load_state_dict( torch.load(model_path, map_location=lambda storage, loc: storage)) else: unet.load_state_dict(torch.load(model_path)) #transform = ToTensor() ori_tensor = ToTensor()(data) if cuda: #ori_tensor = Variable(ori_tensor.cuda()) ori_tensor = ori_tensor.cuda() #else: #ori_tensor = Variable(ori_tensor) #ori_tensor = ori_tensor.cuda() ori_tensor = torch.unsqueeze(ori_tensor, 0) padding_left = 0 padding_right = 0 padding_top = 0 padding_bottom = 0 ori_height = ori_tensor.size()[2] ori_width = ori_tensor.size()[3] if ori_height % 4: padding_top = (4 - ori_height % 4) // 2 padding_bottom = 4 - ori_height % 4 - padding_top use_padding = True if ori_width % 4: padding_left = (4 - ori_width % 4) // 2 padding_right = 4 - ori_width % 4 - padding_left use_padding = True if use_padding: padding_transform = torch.nn.ConstantPad2d( (padding_left, padding_right, padding_top, padding_bottom), 0) ori_tensor = padding_transform(ori_tensor) output = ori_tensor with torch.no_grad(): for _ in range(iter): output = unet(output) if use_padding: output = output[:, :, padding_top:(padding_top + ori_height), padding_left:(padding_left + ori_width)] # # if cuda: # result = (output.data).cpu().numpy() # else: # result = (output.data).numpy() result = (output.data).cpu().numpy() result = result[0, 0, :, :] return result