Added model

app.py

@@ -3,9 +3,11 @@ import torch
 from PIL import Image
 import numpy as np
 from app.utils import recover_light_sources
+from model.model import model
 
 device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model=torch.load('model/model_epoch_49.pth',map_location=device)
+chk=torch.load('model/model_epoch_49.pth',map_location=device)
+model.load_state_dict(chk['model_state_dict'])
 
 def evaluate(model,image):
     model.eval()

model/model.py

@@ -0,0 +1,9 @@
+import segmentation_models_pytorch as smp
+
+model = smp.Unet(
+    encoder_name="mit_b1",
+    encoder_weights="imagenet",
+    in_channels=3,
+    classes=3,
+    activation=None
+)