Added start button to reset progress based on previous state and fixed reload issue

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import os
 import numpy as np
 import cv2 as cv
-import csv
 from google.cloud import storage
 
 GCS_BUCKET_NAME = "veytel-cloud-store"
@@ -21,6 +20,7 @@ service_account_json = {
     "universe_domain": "googleapis.com"
 }
 count = 0
+fresh_start = False
 
 def set_min_dense_1(max_dense_0):
     global scaled_thresh1
@@ -77,22 +77,13 @@ def upload_csv_to_gcs(csv_content, filename):
     blob.upload_from_string(csv_content)
 
 def new__cxr(max_dense_0, max_dense_1, max_dense_2):
-    global textured_cxr, lung_noised, max_val, cxr, mask, scaled_thresh1, scaled_thresh2, scaled_thresh3, image_id, index, count
-    if count >= 0:
-        index += 1
-    if index > 3:
-        index = 1
-        image_id += 1
-
-    count += 1
-
+    global textured_cxr, lung_noised, max_val, cxr, mask, scaled_thresh1, scaled_thresh2, scaled_thresh3, image_id, index, count, label1
     csv_data = read_csv_from_gcs(user)
-    # write count, thresh1, thresh2, thresh3 to csv file
     if count < 30:
         csv_content = ""
         for row in csv_data:
             csv_content += ",".join(row) + "\n"
-        if count >0:
+        if count > 0:
             csv_content += f"{count},{max_dense_0},{max_dense_1},{max_dense_2}\n"
         filename = f"density_{user}.csv"
         upload_csv_to_gcs(csv_content, filename)
@@ -107,7 +98,7 @@ def new__cxr(max_dense_0, max_dense_1, max_dense_2):
         upload_csv_to_gcs(csv_content, filename)
         empty_image = np.zeros((256, 256), dtype=np.uint8)
         label1.update(visible=False)
-        count = -1
+        count = 0
         image_id = 1
         index = 1
         fieldnames = ['count', 'thresh_1', 'thresh_2', 'thresh_3']
@@ -115,6 +106,16 @@ def new__cxr(max_dense_0, max_dense_1, max_dense_2):
         filename = f"density_{user}.csv"
         upload_csv_to_gcs(csv_content, filename)
         return empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, 0
+
+
+    if count >= 0:
+        index += 1
+    if index > 3:
+        index = 1
+        image_id += 1
+
+    count += 1
+    # write count, thresh1, thresh2, thresh3 to csv file
     cxr_file = "cxr" + str(image_id) + "_cxr.png"
     mask_file = "cxr" + str(image_id) + "_mask.png"
     textured_cxr_file = "cxr" + str(image_id) + "_textured_" + str(index) + ".png"
@@ -138,22 +139,33 @@ def new__cxr(max_dense_0, max_dense_1, max_dense_2):
     dense_1 = np.where(((textured_cxr < scaled_thresh2) & (textured_cxr >= scaled_thresh1)), textured_cxr, 0)
     dense_2 = np.where(((textured_cxr < scaled_thresh3) & (textured_cxr >= scaled_thresh2)), textured_cxr, 0)
     dense_3 = np.where(((textured_cxr < 255) & (textured_cxr >= scaled_thresh3)), textured_cxr, 0)
-    return cxr, textured_cxr, lung_noised, lung_noised, dense_0, dense_1, dense_2, dense_3, count+1
+    return cxr, textured_cxr, lung_noised, lung_noised, dense_0, dense_1, dense_2, dense_3, count
 
 
 def create_csv():
-    global count, fieldnames, image_id, index, csv_path
+    global count, fieldnames, image_id, index, csv_path, fresh_start
     fieldnames = ['count', 'thresh_1', 'thresh_2', 'thresh_3']
     csv_data = read_csv_from_gcs(user)
     csv_content = ""
+    last_row_count = 0
     if not csv_data:
         csv_content = ",".join(fieldnames) + "\n"
     else:
-        for row in csv_data:
+        for i, row in enumerate(csv_data):
             csv_content += ",".join(row) + "\n"
+            if i != 0:
+                if isinstance(row[0], str):
+                    last_row_count = 0
+                else:
+                    last_row_count = int(row[0])
+
     #csv_content += f"{count+1},{thresh1},{thresh2},{thresh3}\n"
     filename = f"density_{user}.csv"
     upload_csv_to_gcs(csv_content, filename)
+    count = last_row_count
+    if(count>0):
+        image_id, index = get_image_id_index(count)
+    fresh_start = True
 
 
 def check_auth(username, password):
@@ -177,24 +189,69 @@ def check_auth(username, password):
     elif (user == 'swathi' and password == 'veytel'):
         create_csv()
         return True
+    elif (user == 'mike' and password == 'veytel'):
+        create_csv()
+        return True
+    elif (user == 'nischal' and password == 'veytel'):
+        create_csv()
+        return True
 
 
 def change_vis():
-    global count
+    global count, fresh_start, button1, button2
+    if fresh_start:
+        fresh_start = not fresh_start
     if count >= 30:
-        return gr.Label.update(visible=True), gr.Button.update(visible=False)
+        return gr.Label.update(visible=True), gr.Button.update(visible=False), gr.Button.update(visible=False)
     else:
-        return gr.Label.update(visible=False), gr.Button.update(visible=True)
+        return gr.Label.update(visible=False), gr.Button.update(visible=True), gr.Button.update(visible=False)
 
 
-with gr.Blocks() as demo:
-    image_id = 1
-    index = 1
-    executable_path = os.path.dirname(os.path.realpath(__file__))
-    cxr_dir = os.path.join(executable_path, "Images/cxr")
-    mask_dir = os.path.join(executable_path, "Images/mask")
-    textured_cxr_dir = os.path.join(executable_path, "Images/textured_cxr")
-    lung_noised_dir = os.path.join(executable_path, "Images/lung_noised")
+image_id = 1
+index = 1
+
+def get_image_id_index(count):
+    #global image_id, index
+    if count ==0:
+        return 1, 1
+    image_id = (count - 1) // 3 + 1
+    index = (count - 1) % 3 + 1
+    return image_id, index
+
+def set_layout():
+    global textured_cxr, lung_noised, max_val, cxr, mask, scaled_thresh1, scaled_thresh2, scaled_thresh3, image_id, index, count
+    csv_data = read_csv_from_gcs(user)
+    csv_content = ""
+    last_row_count = 0
+    if csv_data:
+        for i, row in enumerate(csv_data):
+            if i != 0:
+                if str(row[0]) == 'count':
+                    last_row_count = 0
+                else:
+                    last_row_count = int(row[0])
+
+
+    count = last_row_count
+    if (count > 0):
+        image_id, index = get_image_id_index(count)
+
+    if count >= 30:
+        empty_image = np.zeros((256, 256), dtype=np.uint8)
+        return empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, 30
+    image_id, index = get_image_id_index(count)
+    if count == 0:
+        count = 1
+        image_id = 1
+        index = 1
+    else:
+        count +=1
+        index += 1
+        if index > 3:
+            index = 1
+            image_id += 1
+
+    #new__cxr(max_dense_0, max_dense_1, max_dense_2)
     cxr_file = "cxr" + str(image_id) + "_cxr.png"
     mask_file = "cxr" + str(image_id) + "_mask.png"
     textured_cxr_file = "cxr" + str(image_id) + "_textured_" + str(index) + ".png"
@@ -203,7 +260,6 @@ with gr.Blocks() as demo:
     mask_path = os.path.join(mask_dir, mask_file)
     textured_cxr_path = os.path.join(textured_cxr_dir, textured_cxr_file)
     lung_noised_path = os.path.join(lung_noised_dir, lung_noised_file)
-    # print("*"*50, cxr_path)
     cxr = cv.imread(cxr_path, cv.IMREAD_GRAYSCALE)
     mask = cv.imread(mask_path, cv.IMREAD_GRAYSCALE)
     textured_cxr = cv.imread(textured_cxr_path, cv.IMREAD_GRAYSCALE)
@@ -219,6 +275,34 @@ with gr.Blocks() as demo:
     dense_1 = np.where(((textured_cxr < scaled_thresh2) & (textured_cxr >= scaled_thresh1)), textured_cxr, 0)
     dense_2 = np.where(((textured_cxr < scaled_thresh3) & (textured_cxr >= scaled_thresh2)), textured_cxr, 0)
     dense_3 = np.where(((textured_cxr < 255) & (textured_cxr >= scaled_thresh3)), textured_cxr, 0)
+    return cxr, textured_cxr, lung_noised, lung_noised, dense_0, dense_1, dense_2, dense_3, count
+
+
+
+def update_layout(image_id, index):
+    global cxr_dir,mask_dir, textured_cxr_dir,lung_noised_dir,  cxr, mask, textured_cxr, lung_noised, max_val, thresh1, thresh2, thresh3, scaled_thresh1, scaled_thresh2, scaled_thresh3, dense_0, dense_1, dense_2, dense_3, label_title, min_dense_0, max_dense_3, button1, button2, label1
+    executable_path = os.path.dirname(os.path.realpath(__file__))
+    cxr_dir = os.path.join(executable_path, "Images/cxr")
+    mask_dir = os.path.join(executable_path, "Images/mask")
+    textured_cxr_dir = os.path.join(executable_path, "Images/textured_cxr")
+    lung_noised_dir = os.path.join(executable_path, "Images/lung_noised")
+
+    empty_image = np.zeros((256, 256), dtype=np.uint8)
+    cxr = empty_image
+    mask = empty_image
+    textured_cxr = empty_image
+    lung_noised = empty_image
+    max_val = np.percentile(cxr, 97)  # To optimize later
+    thresh1 = 50
+    thresh2 = 100
+    thresh3 = 150
+    scaled_thresh1 = thresh1 * max_val / 255
+    scaled_thresh2 = thresh2 * max_val / 255
+    scaled_thresh3 = thresh3 * max_val / 255
+    dense_0 = np.where(textured_cxr < scaled_thresh1, textured_cxr, 0)
+    dense_1 = np.where(((textured_cxr < scaled_thresh2) & (textured_cxr >= scaled_thresh1)), textured_cxr, 0)
+    dense_2 = np.where(((textured_cxr < scaled_thresh3) & (textured_cxr >= scaled_thresh2)), textured_cxr, 0)
+    dense_3 = np.where(((textured_cxr < 255) & (textured_cxr >= scaled_thresh3)), textured_cxr, 0)
     # open csv in append mode
     # add title
     with gr.Row():
@@ -250,8 +334,9 @@ with gr.Blocks() as demo:
                     set_min_dense_1, inputs=max_dense_0, outputs=[min_dense_1, dense0, dense1])
 
         with gr.Column():
-            progress_log = gr.Slider(1, 30, value=1, step=1, label="progress")
-            button1 = gr.Button(value="Save & continue", visible=True)
+            progress_log = gr.Slider(1, 30, value=0, step=1, label="progress")
+            button1 = gr.Button(value="Save & continue", visible=fresh_start)
+            button2 = gr.Button(value="Start", visible=not fresh_start)
 
     with gr.Row():
         with gr.Column():
@@ -272,10 +357,15 @@ with gr.Blocks() as demo:
 
         with gr.Column():  # adding additional for better visualization
             im3_1 = gr.Image(lung_noised, label="Synthetic CXR")
-
     button1.click(new__cxr, inputs=[max_dense_0, max_dense_1, max_dense_2],
                   outputs=[im1, im2, im3, im3_1, dense0, dense1, dense2, dense3, progress_log])
-    button1.click(change_vis, outputs=[label1, button1])
+    button1.click(change_vis, outputs=[label1, button1, button2])
+    button2.click(set_layout,
+                  outputs=[im1, im2, im3, im3_1, dense0, dense1, dense2, dense3, progress_log])
+    button2.click(change_vis, outputs=[label1, button1, button2])
+
+with gr.Blocks() as demo:
+    update_layout(image_id, index)
 
 if __name__ == "__main__":
     demo.launch(auth=check_auth)