Spaces:
Runtime error
Runtime error
File size: 1,952 Bytes
358321e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import torch
import torch.nn as nn
import torchvision
import torch.backends.cudnn as cudnn
import torch.optim
import os
import sys
import argparse
import time
import dataloader
import model
import numpy as np
from torchvision import transforms
from PIL import Image
import glob
import time
import gradio as gr
def lowlight(image_path):
os.environ['CUDA_VISIBLE_DEVICES']='0'
data_lowlight = Image.open(image_path)
data_lowlight = (np.asarray(data_lowlight)/255.0)
data_lowlight = torch.from_numpy(data_lowlight).float()
data_lowlight = data_lowlight.permute(2,0,1)
data_lowlight = data_lowlight.cuda().unsqueeze(0)
DCE_net = model.enhance_net_nopool().cuda()
DCE_net.load_state_dict(torch.load('snapshots/Epoch99.pth'))
start = time.time()
_,enhanced_image,_ = DCE_net(data_lowlight)
end_time = (time.time() - start)
print(end_time)
image_path = image_path.replace('test_data','result')
result_path = image_path
if not os.path.exists(image_path.replace('/'+image_path.split("/")[-1],'')):
os.makedirs(image_path.replace('/'+image_path.split("/")[-1],''))
torchvision.utils.save_image(enhanced_image, result_path)
def predict(img):
data_lowlight = (np.asarray(img)/255.0)
data_lowlight = torch.from_numpy(data_lowlight).float()
data_lowlight = data_lowlight.permute(2,0,1)
data_lowlight = data_lowlight.cuda().unsqueeze(0)
DCE_net = model.enhance_net_nopool().cuda()
DCE_net.load_state_dict(torch.load('snapshots/Epoch99.pth'))
_,enhanced_image,_ = DCE_net(data_lowlight)
return enhanced_image
if __name__ == '__main__':
# test_images
with torch.no_grad():
# filePath = 'data/test_data/'
# file_list = os.listdir(filePath)
# for file_name in file_list:
# test_list = glob.glob(filePath+file_name+"/*")
# for image in test_list:
# # image = image
# print(image)
# lowlight(image)
interface = gr.Interface(fn=predict, inputs='image', outputs='image')
interface.launch()
|