Update app.py

app.py

@@ -8,12 +8,38 @@ import matplotlib.pyplot as plt
 import os, io, base64
 from PIL import Image 
 
+
+# @title Data retrieval
+def download_weights():
+    import os, requests
+    
+    fname = 'cpsam'
+    
+    url = ["https://osf.io/d7c8e/download"]
+    
+    for j in range(len(url)):
+      if not os.path.isfile(fname[j]):
+        try:
+          r = requests.get(url[j])
+        except requests.ConnectionError:
+          print("!!! Failed to download data !!!")
+        else:
+          if r.status_code != requests.codes.ok:
+            print("!!! Failed to download data !!!")
+          else:
+            with open(fname[j], "wb") as fid:
+              fid.write(r.content)
+
 try:
-    model = models.CellposeModel(gpu=True, pretrained_model="cyto3")
+    download_weights()
+    model = models.CellposeModel(gpu=True, pretrained_model="cpsam")
 except Exception as e:
     print(f"Error loading model: {e}")
     exit(1)
 
+
+
+            
 def plot_flows(y):
     Y = (np.clip(normalize99(y[0][0]),0,1) - 0.5) * 2
     X = (np.clip(normalize99(y[1][0]),0,1) - 0.5) * 2
@@ -101,7 +127,7 @@ def image_resize(img, resize=400):
     
 @spaces.GPU(duration=10)
 def run_model_gpu(img):
-    masks, flows, _ = model.eval(img, channels=[0,0])
+    masks, flows, _ = model.eval(img)
     return masks, flows
 
 #@spaces.GPU(duration=10)