Create speclab.py

speclab.py

@@ -0,0 +1,23 @@
+
+
+def speclab(img):
+
+    # initialize the model
+    model = torch.hub.load('Nano1337/SpecLab', 'srdetect', force_reload=True) # for some reasons loads the model in src rather than demo
+    model.eval()
+
+    # preprocess image to be used as input
+    transforms = A.Compose([
+        A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+        ToTensorV2()
+    ])
+    input = transforms(image=img)['image']
+    input = input.unsqueeze(0)
+
+    # model prediction
+    output = model(input)
+
+    # overlay output onto original image
+    img[output==255] = [0, 255, 0]
+
+    return img