init

app.py

@@ -23,14 +23,17 @@ def dark_inference(img):
     state_dict = torch.load(checkpoint_file_path, map_location='cpu')
     model.load_state_dict(state_dict)
     model.eval()
+    print(f'Load model from {checkpoint_file_path}')
 
     transform = Compose([
         ToTensor(), 
         Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 
         ConvertImageDtype(torch.float) 
     ])
+    input_img = transform(img)
+    print(f'Image shape: {input_img.shape}')
 
-    enhanced_img = model(transform(img).unsqueeze(0))
+    enhanced_img = model(input_img.unsqueeze(0))
     return enhanced_img[0].permute(1, 2, 0).detach().numpy()
 
 
@@ -40,13 +43,16 @@ def exposure_inference(img):
     state_dict = torch.load(checkpoint_file_path, map_location='cpu')
     model.load_state_dict(state_dict)
     model.eval()
+    print(f'Load model from {checkpoint_file_path}')
 
     transform = Compose([
         ToTensor(), 
         ConvertImageDtype(torch.float) 
     ])
+    input_img = transform(img)
+    print(f'Image shape: {input_img.shape}')
 
-    enhanced_img = model(transform(img).unsqueeze(0))
+    enhanced_img = model(input_img.unsqueeze(0))
     return enhanced_img[0].permute(1, 2, 0).detach().numpy()
 
 

test_dark.ipynb

@@ -246,6 +246,39 @@
     "plt.imshow(enhanced_img[0].permute(1, 2, 0).detach().numpy())"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def dark_inference(img_path):\n",
+    "    model = IAT()\n",
+    "    checkpoint_file_path = './checkpoint/best_Epoch_lol.pth'\n",
+    "    state_dict = torch.load(checkpoint_file_path, map_location='cpu')\n",
+    "    model.load_state_dict(state_dict)\n",
+    "    model.eval()\n",
+    "\n",
+    "    img = np.array(Image.open(img_path))\n",
+    "    transform = Compose([\n",
+    "        ToTensor(), \n",
+    "        Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), \n",
+    "        ConvertImageDtype(torch.float) \n",
+    "    ])\n",
+    "\n",
+    "    enhanced_img = model(transform(img).unsqueeze(0))\n",
+    "    return enhanced_img[0].permute(1, 2, 0).detach().numpy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "out = dark_inference('./dark_imgs/1.jpg')"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,