v1

app.py

@@ -3,20 +3,25 @@ import inpaint.infer_model as inpaint
 import sod.infer_model as sod
 import numpy as np
 import torch
+import glob
 # import os
 # cmd = 'sh download.sh'
 # os.system(cmd)
 
 device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+print(device)
 inpaint_model = inpaint.IVModel(device=device)
 sod_model = sod.IVModel(device=torch.device("cpu"))
 def sod_inpaint(img):
     img = img[:,:,::-1]
     res = sod_model.forward(img,None)
+    print('sod complete!')
     res = np.uint8(res)
     res = inpaint_model.forward(res,None)
     res = np.uint8(res)
+    print('inpaint complete!')
     return res[:,:,::-1]
 
-iface = gr.Interface(fn=sod_inpaint, inputs="image", outputs="image", examples='examples', title='显著物体消除', description='这是一个图像API，功能是自动把画面中的显著物体消除', theme='huggingface')
+examples = glob.glob('examples/*.*')
+iface = gr.Interface(fn=sod_inpaint, inputs="image", outputs="image", examples=examples, title='显著物体消除', description='这是一个图像API，功能是自动把画面中的显著物体消除', theme='huggingface')
 iface.launch()

inpaint/infer_model.py

@@ -45,13 +45,13 @@ class IVModel():
         self.__first_forward__()
 
 
-    def __first_forward__(self, input_size=(2048, 4096, 3)):
+    def __first_forward__(self, input_size=(512, 1024, 3)):
         # 调用forward()严格控制最大显存
         print('initialize Inpaint Model...')
         _ = self.forward(np.random.rand(*input_size) * 255, None)
         print('initialize Complete!')
 
-    def __resize_tensor__(self, image, max_size=1024, scale_factor=8):
+    def __resize_tensor__(self, image, max_size=512, scale_factor=8):
         h, w = image.size()[2:]
         if max(h, w) > max_size:
             if h < w:
@@ -68,7 +68,7 @@ class IVModel():
         img_t = img_t.permute(2, 0, 1).unsqueeze(0)
         img_t = img_t / 255.
         img_t_for_net = self.__resize_tensor__(img_t).to(self.device)  # 为了控制最大显存容量
-        img_t_for_out = self.__resize_tensor__(img_t, max_size=2048).to(self.device)  # 为了控制最大显存容量
+        img_t_for_out = self.__resize_tensor__(img_t, max_size=512).to(self.device)  # 为了控制最大显存容量
         return img_t_for_net, img_t_for_out
 
     def forward(self, img, json_data):

sod/infer_model.py

@@ -49,13 +49,13 @@ class IVModel():
         self.__first_forward__()
 
 
-    def __first_forward__(self, input_size=(2048, 2048, 3)):
+    def __first_forward__(self, input_size=(512, 512, 3)):
         # 调用forward()严格控制最大显存
         print('initialize Sod Model...')
         _ = self.forward(np.random.rand(*input_size) * 255, None)
         print('initialize Complete!')
 
-    def __resize_tensor__(self, image, max_size=1024):
+    def __resize_tensor__(self, image, max_size=512):
         h, w = image.size()[2:]
         if max(h, w) > max_size:
             if h < w:
@@ -76,7 +76,7 @@ class IVModel():
         img_t = self.input_preprocess_tensor(img)
         shape = [torch.as_tensor([img_t.shape[2]]), torch.as_tensor([img_t.shape[3]])]
         h, w = img_t.shape[2], img_t.shape[3]
-        img_t_temp = F.interpolate(img_t, (1024, 1024), mode='area')
+        img_t_temp = F.interpolate(img_t, (512, 512), mode='area')
         with torch.no_grad():
             res = self.net(img_t_temp, shape=shape)
         res   = F.interpolate(res[0],size=shape, mode='bilinear')