deflare.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import torch\n",
    "import torchvision\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from PIL import Image\n",
    "from basicsr.data.flare7k_dataset import Flare_Image_Loader,RandomGammaCorrection\n",
    "from basicsr.archs.uformer_arch import Uformer\n",
    "from basicsr.archs.unet_arch import U_Net\n",
    "from basicsr.utils.flare_util import blend_light_source,get_args_from_json,save_args_to_json,mkdir,annel,predict_flare_from_3_channel\n",
    "from torch.distributions import Normal\n",
    "import torchvision.transforms as transforms\n",
    "import os\n",
    "print(\"  ssasa\")\n",
    "\n",
    "# def mkdir(path):\n",
    "# \tfolder = os.path.exists(path)\n",
    "# \tif not folder:\n",
    "# \t\tos.makedirs(path)\n",
    "\n",
    "# def demo(images_path,output_path,model_type,output_ch,pretrain_dir):\n",
    "#     if not os.path.exists(output_path):\n",
    "#         os.makedirs(output_path)\n",
    "#     test_path=glob.glob(images_path)\n",
    "#     result_path=output_path\n",
    "#     torch.cuda.empty_cache()\n",
    "#     if model_type=='Uformer':\n",
    "#         model=Uformer(img_size=512,img_ch=3,output_ch=output_ch).cuda()\n",
    "#         model.load_state_dict(torch.load(pretrain_dir))\n",
    "#     elif model_type=='U_Net' or model_type=='U-Net':\n",
    "#         model=U_Net(img_ch=3,output_ch=output_ch).cuda()\n",
    "#         model.load_state_dict(torch.load(pretrain_dir))\n",
    "#     else:\n",
    "#         assert False, \"This model is not supported!!\"\n",
    "#     to_tensor=transforms.ToTensor()\n",
    "#     resize=transforms.Resize((512,512)) #The output should in the shape of 128X\n",
    "#     for i,image_path in tqdm(enumerate(test_path)):\n",
    "#         mkdir(result_path+\"deflare/\")\n",
    "#         mkdir(result_path+\"flare/\")\n",
    "#         mkdir(result_path+\"input/\")\n",
    "#         mkdir(result_path+\"blend/\")\n",
    "\n",
    "#         deflare_path = result_path+\"deflare/\"+str(i).zfill(5)+\"_deflare.png\"\n",
    "#         flare_path = result_path+\"flare/\"+str(i).zfill(5)+\"_flare.png\"\n",
    "#         merge_path = result_path+\"input/\"+str(i).zfill(5)+\"_input.png\"\n",
    "#         blend_path = result_path+\"blend/\"+str(i).zfill(5)+\"_blend.png\"\n",
    "\n",
    "#         merge_img = Image.open(image_path).convert(\"RGB\")\n",
    "#         merge_img = resize(to_tensor(merge_img))\n",
    "#         merge_img = merge_img.cuda().unsqueeze(0)\n",
    "\n",
    "#         model.eval()\n",
    "#         with torch.no_grad():\n",
    "#             output_img=model(merge_img)\n",
    "#             #if ch is 6, first three channels are deflare image, others are flare image\n",
    "#             #if ch is 3, unsaturated region of output is the deflare image.\n",
    "#             gamma=torch.Tensor([2.2])\n",
    "#             if output_ch==6:\n",
    "#                 deflare_img,flare_img_predicted,merge_img_predicted=predict_flare_from_6_channel(output_img,gamma)\n",
    "#             elif output_ch==3:\n",
    "#                 flare_mask=torch.zeros_like(merge_img)\n",
    "#                 deflare_img,flare_img_predicted=predict_flare_from_3_channel(output_img,flare_mask,output_img,merge_img,merge_img,gamma)\n",
    "#             else:\n",
    "#                 assert False, \"This output_ch is not supported!!\"\n",
    "\n",
    "#             blend_img= blend_light_source(merge_img, deflare_img, 0.97)\n",
    "\n",
    "#             torchvision.utils.save_image(merge_img, merge_path)\n",
    "#             torchvision.utils.save_image(flare_img_predicted, flare_path)\n",
    "#             torchvision.utils.save_image(deflare_img, deflare_path)\n",
    "#             torchvision.utils.save_image(blend_img, blend_path)\n",
    "\n",
    "\n",
    "# model_type=\"Uformer\"\n",
    "# images_path=\"test/test_images/*.*\"\n",
    "# result_path=\"result/test_images/Uformer/\"\n",
    "# pretrain_dir='experiments/pretrained_models/uformer/net_g_last.pth'\n",
    "# output_ch=6\n",
    "# mask_flag=False\n",
    "# demo(images_path,result_path,model_type,output_ch,pretrain_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "flare_py38_t",
   "language": "python",
   "name": "flare_py38_t"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "e3d435e51e702e7b0995eba7d152ca2d2b0b33732f1590a6d012c8aa76dc59cc"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}