Update app.py

app.py

@@ -27,6 +27,16 @@ class LABEL_TYPE(Enum):
     WALL = 0
     FLOOR = 3
 
+def query_image(img):
+    target_size = (img.shape[0], img.shape[1])
+    inputs = preprocessor(images=img, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    outputs.class_queries_logits = outputs.class_queries_logits.cpu()
+    outputs.masks_queries_logits = outputs.masks_queries_logits.cpu()
+    results = preprocessor.post_process_segmentation(outputs=outputs, target_size=target_size)[0].cpu().detach()
+    results = torch.argmax(results, dim=0).numpy()
+    return results
 
 def find_boundary(label_value, mask):
     contours = find_contours(mask == label_value, 0.5, fully_connected="high")
@@ -76,6 +86,7 @@ def detect():
         # Assuming base64_str is the string value without 'data:image/jpeg;base64,'
         img = Image.open(io.BytesIO(base64.decodebytes(bytes(base64_str, "utf-8"))))
         numpydata = np.asarray(img)
+        mask = query_image(numpydata)
         window_mask = find_boundary(LABEL_TYPE.WINDOW.value, mask)
         floor_mask = find_boundary(LABEL_TYPE.FLOOR.value, mask)
         window_contours = [np.fliplr(ctr).astype(np.int32) for ctr in window_mask]