Update app.py

app.py

@@ -46,6 +46,26 @@ class FeatureLoss(nn.Module):
     def __del__(self): self.hooks.remove()
 
 
+def getNeighbours(i, j, n, m) :
+    arr = []
+    if i-1 >= 0 and j-1 >= 0 :
+        arr.append((i-1, j-1))
+    if i-1 >= 0 :
+        arr.append((i-1, j))
+    if i-1 >= 0 and j+1 < m :
+        arr.append((i-1, j+1))
+    if j+1 < m :
+        arr.append((i, j+1))
+    if i+1 < n and j+1 < m :
+        arr.append((i+1, j+1))
+    if i+1 < n :
+        arr.append((i+1, j))
+    if i+1 < n and j-1 >= 0 :
+        arr.append((i+1, j-1))
+    if j-1 >= 0 :
+        arr.append((i, j-1))
+    return arr
+
 MODEL_URL = "https://www.dropbox.com/s/05ong36r29h51ov/popd.pkl?dl=1"
 urllib.request.urlretrieve(MODEL_URL, "popd.pkl")
 path = Path(".")
@@ -70,4 +90,27 @@ if uploaded_file is not None:
   white_areas = (red == 0) & (blue == 0) & (green == 0)
   data[..., :-1][white_areas.T] = (0,0,0) # Transpose back needed
   im2 = Image.fromarray(data)
-  st.image(im2, caption='PoP ArT')
+  membuf = BytesIO()
+  im2.save(membuf, format="png") 
+  img = Image.open(membuf)
+  bitmap = img.load()
+  n = img.size[0]
+  m = img.size[1]
+  stateMap = []
+  for i in range(n):
+    stateMap.append([False for j in range(m)])  
+  queue = [(0, 0)]
+  while queue:
+    e = queue.pop(0)
+    i = e[0]
+    j = e[1]
+    if not stateMap[i][j]:
+      stateMap[i][j] = True
+      color = int((bitmap[i, j][0] + bitmap[i, j][1] + bitmap[i, j][2])/3)
+      if color > 100:
+        bitmap[i, j] = (0, 0, 0)
+        neigh = getNeighbours(i, j, n, m)
+        for ne in neigh:
+          queue.append(ne)
+
+  st.image(img, caption='PoP ArT')