commit from ssq

app.py

@@ -3,16 +3,16 @@ import gradio as gr
 import numpy as np
 import torch.nn.functional as F
 
-from Allweather.util import load_img
+from Allweather.util import load_img, save_img
 from basicsr.models.archs.histoformer_arch import Histoformer
 
 model_restoration = Histoformer()
-checkpoint = torch.load("Allweather/pretrained_models/net_g_real.pth")
-model_restoration.load_state_dict(checkpoint['params'])
-model_restoration.cuda()
-model_restoration = nn.DataParallel(model_restoration)
+model, transform = Histoformer.create_model_and_transforms()
+model = model.to(device)
+model.eval()
 
-def preprocess(file_, factor = 8):
+factor = 8
+def predict(input_img):
     img = np.float32(load_img(file_))/255.
     img = torch.from_numpy(img).permute(2,0,1)
     input_ = img.unsqueeze(0).cuda()
@@ -24,12 +24,15 @@ def preprocess(file_, factor = 8):
     padw = W-w if w%factor!=0 else 0
     input_ = F.pad(input_, (0,padw,0,padh), 'reflect')
     return input_
+    prediction = model_restoration(input_)
+    output_path = "inverse_depth_map.png"
+    restored = restored[:,:,:h,:w]
+    restored = torch.clamp(restored,0,1).cpu().detach().permute(0, 2, 3, 1).squeeze(0).numpy()
+
+    save_img((os.path.join(result_dir, os.path.splitext(os.path.split(file_)[-1])[0]+'.png')), img_as_ubyte(restored))
 
-def predict(input_img):
-    prediction = model_restoration(preprocess(input_img))
-    return input_img, prediction 
 example_images = [
-    "examples/lemur.jpg",
+    "examples/example.jpeg",
 ]
 gradio_app = gr.Interface(
     predict,

requirements.txt

@@ -0,0 +1 @@
+torch