Update app.py

app.py

@@ -6,15 +6,18 @@ from PIL import Image
 os.system('wget https://github.com/FanChiMao/HWMNet/releases/download/v0.0/LOL_enhancement_HWMNet.pth -P experiments/pretrained_models')
 os.system('wget https://github.com/FanChiMao/HWMNet/releases/download/v0.0/MIT5K_enhancement_HWMNet.pth -P experiments/pretrained_models')
 
-def inference(img):
+def inference(img, model):
     os.system('mkdir test')
     #basewidth = 256
     #wpercent = (basewidth / float(img.size[0]))
     #hsize = int((float(img.size[1]) * float(wpercent)))
     #img = img.resize((basewidth, hsize), Image.ANTIALIAS)
     img.save("test/1.png", "PNG")
-    os.system(
-        'python main_test_HWMNet.py --input_dir test --weights experiments/pretrained_models/LOL_enhancement_HWMNet.pth')
+    if model == 'LOL':
+        os.system('python main_test_HWMNet.py --input_dir test --weights experiments/pretrained_models/LOL_enhancement_HWMNet.pth')
+    elif model == 'MIT-5K':
+        os.system('python main_test_HWMNet.py --input_dir test --weights experiments/pretrained_models/MIT5K_enhancement_HWMNet.pth')
+        
     return 'result/1.png'
 
 
@@ -26,6 +29,7 @@ examples = [['low-light.png']]
 gr.Interface(
     inference,
     [gr.inputs.Image(type="pil", label="Input")],
+    gr.inputs.Dropdown(choices=['LOL', 'MIT-5K'], type="value", default='LOL', label="model")],
     gr.outputs.Image(type="file", label="Output"),
     title=title,
     description=description,