Update app.py

app.py

@@ -20,32 +20,32 @@ def normalize_channel(img, lower=1, upper=99):
     return img_norm.astype(np.uint8)
 
 def predict(filename, model=None, device=None, reduce_labels=True):
-	if img_name.endswith('.tif') or img_name.endswith('.tiff'):
-	    img_data = tif.imread(img_name)
-	else:
-	    img_data = io.imread(img_name)
-		# normalize image data
-	if len(img_data.shape) == 2:
-	    img_data = np.repeat(np.expand_dims(img_data, axis=-1), 3, axis=-1)
-	elif len(img_data.shape) == 3 and img_data.shape[-1] > 3:
-	    img_data = img_data[:,:, :3]
-	else:
-	    pass
-	pre_img_data = np.zeros(img_data.shape, dtype=np.uint8)
-	for i in range(3):
-	    img_channel_i = img_data[:,:,i]
-	    if len(img_channel_i[np.nonzero(img_channel_i)])>0:
-		pre_img_data[:,:,i] = normalize_channel(img_channel_i, lower=1, upper=99)
-	#dummy_input = np.zeros((512,512,3)).astype(np.uint8)
-	my_model = MultiStreamCellSegModel.from_pretrained("Lewislou/cellseg_sribd")
-	checkpoints = torch.load('model.pt')
-	my_model.__init__(ModelConfig())
-	my_model.load_checkpoints(checkpoints)
-	with torch.no_grad():
-	    output = my_model(pre_img_data)
-	overlay = visualize_instances_map(pre_img_data,star_label)
-	#cv2.imwrite('prediction.png', cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
-	return overlay
+    if img_name.endswith('.tif') or img_name.endswith('.tiff'):
+        img_data = tif.imread(img_name)
+    else:
+        img_data = io.imread(img_name)
+	# normalize image data
+    if len(img_data.shape) == 2:
+        img_data = np.repeat(np.expand_dims(img_data, axis=-1), 3, axis=-1)
+    elif len(img_data.shape) == 3 and img_data.shape[-1] > 3:
+        img_data = img_data[:,:, :3]
+    else:
+        pass
+    pre_img_data = np.zeros(img_data.shape, dtype=np.uint8)
+    for i in range(3):
+        img_channel_i = img_data[:,:,i]
+        if len(img_channel_i[np.nonzero(img_channel_i)])>0:
+	    pre_img_data[:,:,i] = normalize_channel(img_channel_i, lower=1, upper=99)
+#dummy_input = np.zeros((512,512,3)).astype(np.uint8)
+    my_model = MultiStreamCellSegModel.from_pretrained("Lewislou/cellseg_sribd")
+    checkpoints = torch.load('model.pt')
+    my_model.__init__(ModelConfig())
+    my_model.load_checkpoints(checkpoints)
+    with torch.no_grad():
+        output = my_model(pre_img_data)
+    overlay = visualize_instances_map(pre_img_data,star_label)
+#cv2.imwrite('prediction.png', cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
+    return overlay
 gr.Interface(
     predict,
     inputs=[gr.components.Image(label="Upload Input Image", type="filepath"),