Create app.py

app.py

@@ -0,0 +1,42 @@
+import numpy as np
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models.arch.revcol_sig import FullNet,FullNet_NLP
+from models.arch.classifier import PretrainedConvNext
+import torchvision.transforms.functional as TF
+channels = [64, 128, 256, 512]
+layers = [2, 2, 4, 2]
+num_subnet = 4
+net_i = FullNet_NLP(channels, layers, num_subnet, 4,num_classes=1000, drop_path=0,save_memory=True, inter_supv=True, head_init_scale=None,kernel_size=3).cpu()
+net_i.load_state_dict(torch.load('./merge_stem_reg_014_00055524.pt')['icnn'])
+net_c = PretrainedConvNext("convnext_small_in22k").cpu()
+net_c.load_state_dict(torch.load('./cls_newdis_058_00014384.pt')['icnn'])
+net_i.eval()
+net_c.eval()
+def align(x1):
+    x2 = x1
+    h, w = x1.height, x1.width
+    h, w = h // 32 * 32, w // 32 * 32
+    x1 = x1.resize((w, h))
+    x2 = x2.resize((w, h))
+    return x1
+def predict(img):
+    with torch.no_grad():
+        img=align(img)
+        image_tensor=TF.to_tensor(img).cpu()
+        image_tensor=image_tensor.unsqueeze(0).cuda()
+        ipt=net_c(image_tensor)
+        output_i, output_j=net_i(image_tensor,ipt,prompt=True)
+        output_j_out=[]
+        for i in range(4):
+            out_reflection, out_clean = output_j[i][:, :3, ...], output_j[i][:, 3:, ...]
+            output_j_out.append(out_clean)  
+            output_j_out.append(out_reflection)
+        clean = output_j_out[6]
+        clean=torch.clamp(clean, 0, 1)
+    return clean
+demo=gr.Interface(predict, gr.Image(), "image")
+
+demo.launch()