umuthopeyildirim commited on
Commit
dccb91c
1 Parent(s): 0f1bbf6

Refactor image processing code to use torch.no_grad() context manager

Browse files
Files changed (1) hide show
  1. app.py +13 -12
app.py CHANGED
@@ -93,22 +93,23 @@ with gr.Blocks(css=css) as demo:
93
  image = Normalize(mean=[0.485, 0.456, 0.406], std=[
94
  0.229, 0.224, 0.225])(image)
95
  # image = torch.from_numpy(image).unsqueeze(0)
96
- image = torch.autograd.Variable(image.unsqueeze(0))
 
97
 
98
- pred_depths_r_list, _, _ = model(image)
99
- image_flipped = flip_lr(image)
100
- pred_depths_r_list_flipped, _, _ = model(image_flipped)
101
- pred_depth = post_process_depth(
102
- pred_depths_r_list[-1], pred_depths_r_list_flipped[-1])
103
 
104
- pred_depth = pred_depth.cpu().numpy().squeeze()
105
- colored_depth = cv2.applyColorMap(
106
- pred_depth, cv2.COLORMAP_INFERNO)[:, :, ::-1]
107
 
108
- tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
109
- pred_depth.save(tmp.name)
110
 
111
- return [(original_image, colored_depth), tmp.name]
112
 
113
  submit.click(on_submit, inputs=[input_image], outputs=[
114
  depth_image_slider, raw_file])
 
93
  image = Normalize(mean=[0.485, 0.456, 0.406], std=[
94
  0.229, 0.224, 0.225])(image)
95
  # image = torch.from_numpy(image).unsqueeze(0)
96
+ with torch.no_grad():
97
+ image = torch.autograd.Variable(image.unsqueeze(0))
98
 
99
+ pred_depths_r_list, _, _ = model(image)
100
+ image_flipped = flip_lr(image)
101
+ pred_depths_r_list_flipped, _, _ = model(image_flipped)
102
+ pred_depth = post_process_depth(
103
+ pred_depths_r_list[-1], pred_depths_r_list_flipped[-1])
104
 
105
+ pred_depth = pred_depth.cpu().numpy().squeeze()
106
+ colored_depth = cv2.applyColorMap(
107
+ pred_depth, cv2.COLORMAP_INFERNO)[:, :, ::-1]
108
 
109
+ tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
110
+ pred_depth.save(tmp.name)
111
 
112
+ return [(original_image, colored_depth), tmp.name]
113
 
114
  submit.click(on_submit, inputs=[input_image], outputs=[
115
  depth_image_slider, raw_file])