Upload folder using huggingface_hub

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 import gradio as gr
-from ormbg.models.ormbg import ORMBG
+from ormbg import ORMBG
 from PIL import Image
 
 model_path = "models/ormbg.pth"

ormbg.py

@@ -357,6 +357,9 @@ class myrebnconv(nn.Module):
         return self.rl(self.bn(self.conv(x)))
 
 
+bce_loss = nn.BCELoss(size_average=True)
+
+
 class ORMBG(nn.Module):
 
     def __init__(self, in_ch=3, out_ch=1):
@@ -398,6 +401,14 @@ class ORMBG(nn.Module):
 
         # self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
 
+    def compute_loss(self, predictions, ground_truth):
+        loss0, loss = 0.0, 0.0
+        for i in range(0, len(predictions)):
+            loss = loss + bce_loss(predictions[i], ground_truth)
+            if i == 0:
+                loss0 = loss
+        return loss0, loss
+
     def forward(self, x):
 
         hx = x